import types from unittest.mock import MagicMock import pytest from aarecommon.math.coordinate import Coordinate, SmargonCoordinate from aarecommon.models.raster_grid import RasterGridRequest from jfjoch_client.exceptions import NotFoundException from aare.daq.operations.common.runtime import DAQRuntimeState from aare.daq.operations.common.services import OperationServices from aare.daq.operations.raster.models import RasterContext, RasterDependencies, RasterSettings 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_raster_context(*, jfjoch, aare, sample=None): deps = RasterDependencies( cfg=types.SimpleNamespace( simulated_detector=False, abr_meas_pos=types.SimpleNamespace(at_mm=Coordinate(x=0.0, y=0.0, z=0.0)), get_beam_mark=lambda zoom: (0.0, 0.0), ), devs=types.SimpleNamespace(aerotech_omega=0.0, zoom=100.0), mlbox=types.SimpleNamespace(), jfjoch=jfjoch, aare=aare, ) runtime = DAQRuntimeState( sample_provider=types.SimpleNamespace(sample=sample), sample_geometry_provider=types.SimpleNamespace(sample_geometry=types.SimpleNamespace()), status_provider=types.SimpleNamespace(status=None), ) services = OperationServices( screenshots=types.SimpleNamespace(save_to_db=lambda *args, **kwargs: None) ) settings = RasterSettings( auto_raster_max_images=4500, auto_raster_min_cell_size_mm=0.005, auto_raster_skip_if_exceed_max_image_threshold=True, ) return RasterContext(deps=deps, runtime=runtime, services=services, settings=settings) def make_service() -> RasterService: jfjoch = MagicMock() aare = MagicMock() context = _make_raster_context(jfjoch=jfjoch, aare=aare) 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.deps.jfjoch.get_diffraction_image.assert_not_called() service.ctx.deps.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.deps.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.deps.jfjoch.get_diffraction_image.assert_called_once_with( 2, show_spots=True, show_res_est=True, show_beam_center=True ) service.ctx.deps.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()] scan_result.images[2].spots = 479 scan_result.images[2].res = 2.23 service.ctx.deps.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.deps.jfjoch.get_diffraction_image.assert_called_once_with( 2, show_spots=True, show_res_est=True, show_beam_center=True ) service.ctx.deps.aare.upload_jpg.assert_called_once_with( 123, "preview", b"jpeg-bytes", message="Raster diffraction (479 spots, 2.23 Å)" )