AareDAQ/tests/unit/daq/test_raster_logic.py

import pytest
from unittest.mock import MagicMock

import aare.daq.daq as daq_module
from aare.common.coordinate import Coordinate, SmargonCoordinate
from aare.common.raster_grid import RasterGridRequest
from aare.daq.daq import AareDAQ


def make_request(n_x: int, n_y: int, cell_x: float = 0.01, cell_y: float = 0.02) -> RasterGridRequest:
    return RasterGridRequest(
        exp_time_s=0.02,
        transmission=1.0,
        file_prefix="test",
        n_x=n_x,
        n_y=n_y,
        grid_size_mm=Coordinate(x=cell_x, y=cell_y),
        smargon_top_left=SmargonCoordinate(),
        omega_deg=0.0,
    )


def test_grid_image_id_from_centre_offset_single_cell():
    request = make_request(1, 1)
    image_id = AareDAQ._grid_image_id_from_centre_offset(
        x_mm=request.grid_size_mm.x / 2.0,
        y_mm=request.grid_size_mm.y / 2.0,
        request=request,
    )
    assert image_id == 0


def test_grid_image_id_from_centre_offset_odd_grid():
    request = make_request(5, 3, cell_x=0.01, cell_y=0.02)

    x_mm = ((request.n_x - 1) * request.grid_size_mm.x) / 2.0
    y_mm = ((request.n_y - 1) * request.grid_size_mm.y) / 2.0

    image_id = AareDAQ._grid_image_id_from_centre_offset(
        x_mm=x_mm,
        y_mm=y_mm,
        request=request,
    )

    assert image_id == 7


def test_grid_image_id_from_centre_offset_even_grid():
    request = make_request(4, 4, cell_x=0.01, cell_y=0.01)

    x_mm = ((request.n_x - 1) * request.grid_size_mm.x) / 2.0
    y_mm = ((request.n_y - 1) * request.grid_size_mm.y) / 2.0

    image_id = AareDAQ._grid_image_id_from_centre_offset(
        x_mm=x_mm,
        y_mm=y_mm,
        request=request,
    )

    assert image_id == 6


def test_grid_image_id_from_centre_offset_vertical_scan():
    request = make_request(1, 50, cell_x=0.01, cell_y=0.005)

    x_mm = request.grid_size_mm.x / 2.0
    y_mm = ((request.n_y - 1) * request.grid_size_mm.y) / 2.0

    image_id = AareDAQ._grid_image_id_from_centre_offset(
        x_mm=x_mm,
        y_mm=y_mm,
        request=request,
    )

    assert image_id == 24


def test_grid_image_id_from_centre_offset_rejects_non_positive_cell_size():
    request = make_request(3, 3)
    request.grid_size_mm = Coordinate(x=0.0, y=0.01)

    with pytest.raises(ValueError, match="grid_size_mm must be positive"):
        AareDAQ._grid_image_id_from_centre_offset(
            x_mm=0.0,
            y_mm=0.0,
            request=request,
        )


def test_grid_image_id_from_centre_offset_rejects_zero_dimensions():
    request = make_request(1, 1)
    request.n_x = 0

    with pytest.raises(ValueError, match="Raster grid dimensions must be >= 1"):
        AareDAQ._grid_image_id_from_centre_offset(
            x_mm=0.0,
            y_mm=0.0,
            request=request,
        )


def test_upload_raster_diffraction_preview_skips_out_of_range_image_id():
    daq = object.__new__(AareDAQ)
    daq._AareDAQ__jfjoch = MagicMock()
    daq._AareDAQ__aare = MagicMock()
    daq._AareDAQ__cfg = MagicMock()
    daq._AareDAQ__cfg.get_beam_mark.return_value = (0.0, 0.0)
    daq._AareDAQ__cfg.beam_size_mm = Coordinate(x=0.01, y=0.01)
    daq._AareDAQ__cfg.beam_mark_coeff = MagicMock()
    daq._AareDAQ__cfg.pixel_to_mm = MagicMock(return_value=0.001)
    daq._AareDAQ__cfg.current_sample = None

    request = make_request(2, 2)
    scan_result = MagicMock()
    scan_result.images = [MagicMock(), MagicMock(), MagicMock(), MagicMock()]

    daq._upload_raster_diffraction_preview(
        sample_id=123,
        filename="preview",
        image_id=10,
        scan_result=scan_result,
        request=request,
    )

    daq._AareDAQ__jfjoch.get_diffraction_image.assert_not_called()
    daq._AareDAQ__aare.upload_jpg.assert_not_called()


def test_upload_raster_diffraction_preview_ignores_not_found(monkeypatch):
    class FakeNotFoundException(Exception):
        pass

    monkeypatch.setattr(daq_module, "NotFoundException", FakeNotFoundException)

    daq = object.__new__(AareDAQ)
    daq._AareDAQ__jfjoch = MagicMock()
    daq._AareDAQ__aare = MagicMock()
    daq._AareDAQ__cfg = MagicMock()
    daq._AareDAQ__cfg.get_beam_mark.return_value = (0.0, 0.0)
    daq._AareDAQ__cfg.beam_size_mm = Coordinate(x=0.01, y=0.01)
    daq._AareDAQ__cfg.beam_mark_coeff = MagicMock()
    daq._AareDAQ__cfg.pixel_to_mm = MagicMock(return_value=0.001)
    daq._AareDAQ__cfg.current_sample = None

    request = make_request(2, 2)
    scan_result = MagicMock()
    scan_result.images = [MagicMock(), MagicMock(), MagicMock(), MagicMock()]
    daq._AareDAQ__jfjoch.get_diffraction_image.side_effect = FakeNotFoundException()

    daq._upload_raster_diffraction_preview(
        sample_id=123,
        filename="preview",
        image_id=2,
        scan_result=scan_result,
        request=request,
    )

    daq._AareDAQ__jfjoch.get_diffraction_image.assert_called_once_with(2)
    daq._AareDAQ__aare.upload_jpg.assert_not_called()


def test_upload_raster_diffraction_preview_uploads_when_present():
    daq = object.__new__(AareDAQ)
    daq._AareDAQ__jfjoch = MagicMock()
    daq._AareDAQ__aare = MagicMock()
    daq._AareDAQ__cfg = MagicMock()
    daq._AareDAQ__cfg.get_beam_mark.return_value = (0.0, 0.0)
    daq._AareDAQ__cfg.beam_size_mm = Coordinate(x=0.01, y=0.01)
    daq._AareDAQ__cfg.beam_mark_coeff = MagicMock()
    daq._AareDAQ__cfg.pixel_to_mm = MagicMock(return_value=0.001)
    daq._AareDAQ__cfg.current_sample = None

    request = make_request(2, 2)
    scan_result = MagicMock()
    scan_result.images = [MagicMock(), MagicMock(), MagicMock(), MagicMock()]
    daq._AareDAQ__jfjoch.get_diffraction_image.return_value = b"jpeg-bytes"

    daq._upload_raster_diffraction_preview(
        sample_id=123,
        filename="preview",
        image_id=2,
        scan_result=scan_result,
        request=request,
    )

    daq._AareDAQ__jfjoch.get_diffraction_image.assert_called_once_with(2)
    daq._AareDAQ__aare.upload_jpg.assert_called_once_with(123, "preview", b"jpeg-bytes")