218 lines
6.3 KiB
Python
218 lines
6.3 KiB
Python
import types
|
|
|
|
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.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)
|
|
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()]
|
|
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)
|
|
service.ctx.deps.aare.upload_jpg.assert_called_once_with(123, "preview", b"jpeg-bytes") |