AareDAQ/tests/unit/daq/test_raster_logic.py

import pytest
from types import SimpleNamespace
from unittest.mock import MagicMock

from jfjoch_client.exceptions import NotFoundException

from aare.common.coordinate import Coordinate, SmargonCoordinate
from aare.common.raster_grid import RasterGridRequest
from aare.daq.operations.raster.models import RasterContext
from aare.daq.operations.raster.service import RasterService


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 make_service() -> RasterService:
    runtime = SimpleNamespace(
        sample=None,
        sample_geometry=SimpleNamespace(),
        status=SimpleNamespace(),
    )
    context = RasterContext(
        cfg=MagicMock(),
        devs=MagicMock(),
        mlbox=MagicMock(),
        jfjoch=MagicMock(),
        aare=MagicMock(),
        runtime=runtime,
        state_controller=MagicMock(),
        services=MagicMock(),
        auto_raster_max_images=4500,
        auto_raster_min_cell_size_mm=0.005,
        auto_raster_skip_if_exceed_max_image_threshold=True,
    )
    return RasterService(context=context, logger=MagicMock())


def test_grid_image_id_from_centre_offset_single_cell():
    request = make_request(1, 1)
    image_id = RasterService._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 = RasterService._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 = RasterService._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 = RasterService._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"):
        RasterService._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"):
        RasterService._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():
    service = make_service()

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

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

    service.ctx.jfjoch.get_diffraction_image.assert_not_called()
    service.ctx.aare.upload_jpg.assert_not_called()


def test_upload_raster_diffraction_preview_ignores_not_found():
    service = make_service()

    request = make_request(2, 2)
    scan_result = MagicMock()
    scan_result.images = [MagicMock(), MagicMock(), MagicMock(), MagicMock()]
    service.ctx.jfjoch.get_diffraction_image.side_effect = NotFoundException()

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

    service.ctx.jfjoch.get_diffraction_image.assert_called_once_with(2)
    service.ctx.aare.upload_jpg.assert_not_called()


def test_upload_raster_diffraction_preview_uploads_when_present():
    service = make_service()

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

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

    service.ctx.jfjoch.get_diffraction_image.assert_called_once_with(2)
    service.ctx.aare.upload_jpg.assert_called_once_with(123, "preview", b"jpeg-bytes")