diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py index 7444b34a..efa8112c 100644 --- a/src/aare/beamline_dispatch/default/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/default/beamline_dispatch.py @@ -2,14 +2,17 @@ import json import os from abc import ABC from importlib.resources import files -from typing import Protocol +from typing import Any -import numpy as np from aarecommon.config.beamline import MXBeamline from aarecommon.models.beam_centre import BeamCentre -from numpy.typing import ArrayLike, NDArray -from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch, Geometry +from aare.beamline_dispatch.protocols import ( + AuthDispatch, + BeamCenterModel, + BeamlineDispatch, + Geometry, +) class DefaultAuthDispatch(AuthDispatch): @@ -21,12 +24,6 @@ class DefaultAuthDispatch(AuthDispatch): return key -class BeamCenterModel(Protocol): - def predict( - self, det_z_mm: ArrayLike, det_y_mm: ArrayLike - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ... - - class DefaultGeometry(Geometry): def __init__(self, beamline: MXBeamline) -> None: super().__init__() @@ -57,3 +54,7 @@ class DefaultDispatch(BeamlineDispatch, ABC): @property def geo(self): return self._geo + + @property + def default_rotation_parameters(self) -> dict[str, Any]: + return {"exp_time_s": 0.04, "dtz": 110, "incr_omega_deg": 0.2} diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py index e3bc438d..39001d3e 100644 --- a/src/aare/beamline_dispatch/protocols.py +++ b/src/aare/beamline_dispatch/protocols.py @@ -1,7 +1,8 @@ from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Protocol -from aarecommon.math.beam_center import BeamCenterFromDetectorStage +import numpy as np +from numpy.typing import ArrayLike, NDArray class AuthDispatch(ABC): @@ -36,10 +37,16 @@ class BecMacros(ABC): def full_flux_ph_per_s() -> float: ... +class BeamCenterModel(Protocol): + def predict( + self, det_z_mm: ArrayLike, det_y_mm: ArrayLike + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ... + + class Geometry(ABC): @property @abstractmethod - def beam_centre_model(self) -> BeamCenterFromDetectorStage: ... + def beam_centre_model(self) -> BeamCenterModel: ... @property @abstractmethod @@ -57,3 +64,6 @@ class BeamlineDispatch(ABC): @property @abstractmethod def geo(self) -> Geometry: ... + @property + @abstractmethod + def default_rotation_parameters(self) -> dict[str, Any]: ... diff --git a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py index 34be4e3c..ec52c85e 100644 --- a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py @@ -119,3 +119,8 @@ class X06daDispatch(DefaultDispatch): @property def bec_macros(self) -> BecMacros: return self._bec_macros + + @property + def default_rotation_parameters(self) -> dict[str, Any]: + default = super().default_rotation_parameters + return default | {"transmission": 1.0} diff --git a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py index e3059480..160be3df 100644 --- a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py @@ -120,3 +120,8 @@ class X10saDispatch(DefaultDispatch): @property def bec_macros(self) -> BecMacros: return self._bec_macros + + @property + def default_rotation_parameters(self) -> dict[str, Any]: + default = super().default_rotation_parameters + return default | {"transmission": 0.1} diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 1f3377df..c07cebd0 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -1645,7 +1645,7 @@ class AareDAQ: ): return None, file_prefix - params = SimpleScanParameters() + params = SimpleScanParameters(**self._dispatch.default_rotation_parameters) if (exp := getattr(aaredb_params, "exposure", None)) is not None: params.exp_time_s = exp @@ -2393,7 +2393,7 @@ class AareDAQ: spreadsheet_params, _file_prefix = self.spreadsheet_params() logger.debug(f"spreadsheet_params: {spreadsheet_params}") smart_params = self._cfg.auto_params - default_params = SimpleScanParameters(exp_time_s=0.04, dtz=110, incr_omega_deg=0.2) + default_params = SimpleScanParameters(**self._dispatch.default_rotation_parameters) if prefer_smart: if smart_params: return smart_params, "smart_params" diff --git a/tests/unit/daq/conftest.py b/tests/unit/daq/conftest.py new file mode 100644 index 00000000..bcfe8aa7 --- /dev/null +++ b/tests/unit/daq/conftest.py @@ -0,0 +1,15 @@ +from unittest.mock import MagicMock + +import pytest + +from aare.daq.daq import AareDAQ + + +@pytest.fixture() +def mock_daq(): + daq = object.__new__(AareDAQ) + daq._set_state = MagicMock() + daq._devs = MagicMock() + daq._aare = MagicMock() + daq._handle_operation_error = MagicMock() + return daq diff --git a/tests/unit/daq/test_aare_daq_loop_centering.py b/tests/unit/daq/test_aare_daq_loop_centering.py index eb3f9ea2..65971675 100644 --- a/tests/unit/daq/test_aare_daq_loop_centering.py +++ b/tests/unit/daq/test_aare_daq_loop_centering.py @@ -13,27 +13,18 @@ from aarecommon.models.models import SampleShortInfo from aare.daq.daq import AareDAQ -def test_execute_loop_centering_rejects_invalid_sample(): - daq = object.__new__(AareDAQ) - daq._set_state = MagicMock() - daq._devs = MagicMock() - daq._aare = MagicMock() - daq._handle_operation_error = MagicMock() - daq._create_loop_centering_service = MagicMock() +def test_execute_loop_centering_rejects_invalid_sample(mock_daq): + mock_daq._create_loop_centering_service = MagicMock() - result = daq._execute_loop_centering(None) + result = mock_daq._execute_loop_centering(None) assert result is False - daq._create_loop_centering_service.assert_not_called() + mock_daq._create_loop_centering_service.assert_not_called() -def test_execute_loop_centering_handles_service_failure(): - daq = object.__new__(AareDAQ) - daq._set_state = MagicMock() - daq._devs = types.SimpleNamespace(lamp_light=None) - daq._aare = MagicMock() - daq._handle_operation_error = MagicMock() - daq.save_screenshot_db = MagicMock() +def test_execute_loop_centering_handles_service_failure(mock_daq): + mock_daq._devs = types.SimpleNamespace(lamp_light=None) + mock_daq.save_screenshot_db = MagicMock() sample = SampleShortInfo( db_id=1, @@ -50,21 +41,16 @@ def test_execute_loop_centering_handles_service_failure(): service.run.return_value = types.SimpleNamespace( success=False, comment="failed", error=LoopCenteringFailed("failed") ) - daq._create_loop_centering_service = MagicMock(return_value=service) + mock_daq._create_loop_centering_service = MagicMock(return_value=service) - result = daq._execute_loop_centering(sample) + result = mock_daq._execute_loop_centering(sample) assert result is False - daq._handle_operation_error.assert_called_once() + mock_daq._handle_operation_error.assert_called_once() -def test_execute_loop_centering_handles_service_success(): - daq = object.__new__(AareDAQ) - daq._set_state = MagicMock() - daq._devs = MagicMock() - daq._aare = MagicMock() - daq._handle_operation_error = MagicMock() - daq.save_screenshot_db = MagicMock() +def test_execute_loop_centering_handles_service_success(mock_daq): + mock_daq.save_screenshot_db = MagicMock() sample = SampleShortInfo( db_id=1, @@ -79,12 +65,12 @@ def test_execute_loop_centering_handles_service_success(): service = MagicMock() service.run.return_value = types.SimpleNamespace(success=True, comment=None, error=None) - daq._create_loop_centering_service = MagicMock(return_value=service) + mock_daq._create_loop_centering_service = MagicMock(return_value=service) - result = daq._execute_loop_centering(sample) + result = mock_daq._execute_loop_centering(sample) assert result is True - daq.save_screenshot_db.assert_called_once() + mock_daq.save_screenshot_db.assert_called_once() def test_record_best_effort_step_failure_marks_progress_and_logs_warning(mock_logger): diff --git a/tests/unit/daq/test_rotation_params.py b/tests/unit/daq/test_rotation_params.py new file mode 100644 index 00000000..6f317333 --- /dev/null +++ b/tests/unit/daq/test_rotation_params.py @@ -0,0 +1,45 @@ +import importlib.util +from unittest.mock import MagicMock, patch + +import pytest +from aarecommon.models.models import SampleShortInfo + +from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch +from aare.daq.daq import AareDAQ + + +@pytest.fixture() +def daq(mock_daq): + mock_daq.sync_current_sample_from_tell = MagicMock( + return_value=SampleShortInfo( + db_id=12345, + puck_name="T1", + pin=1, + dewar_name="TEST", + sample_name="test_sample", + run_number=42, + ) + ) + mock_daq.spreadsheet_params = MagicMock(return_value=(None, "/data/abcd/")) + mock_daq._cfg = MagicMock(auto_params=None) + return mock_daq + + +@pytest.mark.skipif( + importlib.util.find_spec("pxii_bec") is None, reason="run only for pxii flavour" +) +def test_pxii_default_params(daq: AareDAQ): + with patch.dict("os.environ", {"BEAMLINE": "X10SA"}): + daq._dispatch = get_beamline_dispatch() + params = daq.get_collection_params(prefer_smart=False)[0] + assert params.transmission == 0.1 + + +@pytest.mark.skipif( + importlib.util.find_spec("pxiii_bec") is None, reason="run only for pxiii flavour" +) +def test_pxiii_default_params(daq: AareDAQ): + with patch.dict("os.environ", {"BEAMLINE": "X06DA"}): + daq._dispatch = get_beamline_dispatch() + params = daq.get_collection_params(prefer_smart=False)[0] + assert params.transmission == 1