Move test file location to common

src/mx_bluesky/common/plans/write_sample_status.py

@@ -0,0 +1,34 @@
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+
+from mx_bluesky.common.utils.exceptions import SampleException
+
+
+def deposit_sample_error(exception_type, sample_id):
+    @bpp.run_decorator(
+        md={
+            "metadata": {"sample_id": sample_id},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    def _inner():
+        if exception_type == "Beamline":
+            raise AssertionError()
+        elif exception_type == "Sample":
+            raise SampleException
+        yield from bps.null()
+
+    yield from _inner()
+
+
+def deposit_loaded_sample(sample_id):
+    @bpp.run_decorator(
+        md={
+            "metadata": {"sample_id": sample_id},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    def _inner():
+        yield from bps.null()
+
+    yield from _inner()

tests/unit_tests/common/test_do_fgs.py

@@ -0,0 +1,170 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from bluesky.callbacks import CallbackBase
+from bluesky.plan_stubs import null
+from bluesky.run_engine import RunEngine
+from bluesky.simulators import RunEngineSimulator, assert_message_and_return_remaining
+from bluesky.utils import MsgGenerator
+from dodal.beamlines.i03 import eiger
+from dodal.devices.fast_grid_scan import ZebraFastGridScan
+from dodal.devices.synchrotron import Synchrotron, SynchrotronMode
+from dodal.devices.zocalo.zocalo_results import (
+    ZOCALO_STAGE_GROUP,
+)
+from event_model.documents import Event, RunStart
+from ophyd_async.core import DeviceCollector
+from ophyd_async.testing import set_mock_value
+
+from mx_bluesky.common.parameters.constants import (
+    EnvironmentConstants,
+    PlanNameConstants,
+    TriggerConstants,
+)
+from mx_bluesky.common.plans.do_fgs import kickoff_and_complete_gridscan
+
+
+@pytest.fixture
+def fgs_devices(RE):
+    with DeviceCollector(mock=True):
+        synchrotron = Synchrotron()
+        grid_scan_device = ZebraFastGridScan("zebra_fgs")
+
+    # Eiger done separately as not ophyd-async yet
+    detector = eiger(fake_with_ophyd_sim=True)
+
+    return {
+        "synchrotron": synchrotron,
+        "grid_scan_device": grid_scan_device,
+        "detector": detector,
+    }
+
+
+@patch("mx_bluesky.common.plans.do_fgs.read_hardware_for_zocalo")
+@patch("mx_bluesky.common.plans.do_fgs.check_topup_and_wait_if_necessary")
+def test_kickoff_and_complete_gridscan_correct_messages(
+    mock_check_topup,
+    mock_read_hardware,
+    sim_run_engine: RunEngineSimulator,
+    fgs_devices,
+):
+    def null_plan() -> MsgGenerator:
+        yield from null()
+
+    synchrotron = fgs_devices["synchrotron"]
+    detector = fgs_devices["detector"]
+    fgs_device = fgs_devices["grid_scan_device"]
+
+    msgs = sim_run_engine.simulate_plan(
+        kickoff_and_complete_gridscan(
+            fgs_device,
+            detector,
+            synchrotron,
+            scan_points=[],
+            scan_start_indices=[],
+            plan_during_collection=null_plan,
+        )
+    )
+
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "read"
+        and msg.obj.name == "grid_scan_device-expected_images",
+    )
+
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "read" and msg.obj.name == "eiger_cam_acquire_time",
+    )
+
+    mock_check_topup.assert_called_once()
+    mock_read_hardware.assert_called_once()
+
+    msgs = assert_message_and_return_remaining(msgs, lambda msg: msg.command == "wait")
+
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "wait" and msg.kwargs["group"] == ZOCALO_STAGE_GROUP,
+    )
+
+    msgs = assert_message_and_return_remaining(
+        msgs, lambda msg: msg.command == "kickoff"
+    )
+
+    msgs = assert_message_and_return_remaining(msgs, lambda msg: msg.command == "wait")
+
+    msgs = assert_message_and_return_remaining(msgs, lambda msg: msg.command == "null")
+
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "complete" and msg.obj.name == "grid_scan_device",
+    )
+
+    msgs = assert_message_and_return_remaining(msgs, lambda msg: msg.command == "wait")
+
+
+# This test should use the real Zocalo callbacks once https://github.com/DiamondLightSource/mx-bluesky/issues/215 is done
+def test_kickoff_and_complete_gridscan_with_run_engine_correct_documents(
+    RE: RunEngine, fgs_devices
+):
+    class TestCallback(CallbackBase):
+        def start(self, doc: RunStart):
+            self.trigger_plan = doc.get(TriggerConstants.ZOCALO)
+            self.subplan_name = doc.get("subplan_name")
+            self.scan_points = doc.get("scan_points")
+            self.scan_start_indices = doc.get("scan_start_indices")
+            self.zocalo_environment = doc.get("zocalo_environment")
+
+        def event(self, doc: Event):
+            self.event_data = list(doc.get("data").keys())
+            return doc
+
+    test_callback = TestCallback()
+
+    RE.subscribe(test_callback)
+    synchrotron = fgs_devices["synchrotron"]
+    set_mock_value(synchrotron.synchrotron_mode, SynchrotronMode.DEV)
+    detector = fgs_devices["detector"]
+    fgs_device: ZebraFastGridScan = fgs_devices["grid_scan_device"]
+
+    detector.unstage = MagicMock()
+
+    set_mock_value(fgs_device.status, 1)
+
+    with patch("mx_bluesky.common.plans.do_fgs.bps.complete"):
+        RE(
+            kickoff_and_complete_gridscan(
+                fgs_device,
+                detector,
+                synchrotron,
+                scan_points=[],
+                scan_start_indices=[],
+            )
+        )
+
+    assert test_callback.trigger_plan == PlanNameConstants.DO_FGS
+    assert test_callback.subplan_name == PlanNameConstants.DO_FGS
+    assert test_callback.scan_points == []
+    assert test_callback.scan_start_indices == []
+    assert test_callback.zocalo_environment == EnvironmentConstants.ZOCALO_ENV
+    assert len(test_callback.event_data) == 1
+    assert test_callback.event_data[0] == "eiger_odin_file_writer_id"
+
+
+@patch("mx_bluesky.common.plans.do_fgs.check_topup_and_wait_if_necessary")
+def test_error_if_kickoff_and_complete_gridscan_parameters_wrong_lengths(
+    mock_check_topup, sim_run_engine: RunEngineSimulator, fgs_devices
+):
+    synchrotron = fgs_devices["synchrotron"]
+    detector = fgs_devices["detector"]
+    fgs_device = fgs_devices["grid_scan_device"]
+    with pytest.raises(AssertionError):
+        sim_run_engine.simulate_plan(
+            kickoff_and_complete_gridscan(
+                fgs_device,
+                detector,
+                synchrotron,
+                scan_points=[],
+                scan_start_indices=[0],
+            )
+        )

tests/unit_tests/common/test_write_sample_status.py

@@ -0,0 +1,65 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from bluesky.run_engine import RunEngine
+
+from mx_bluesky.common.external_interaction.ispyb.exp_eye_store import BLSampleStatus
+from mx_bluesky.common.plans.write_sample_status import (
+    deposit_loaded_sample,
+    deposit_sample_error,
+)
+from mx_bluesky.common.utils.exceptions import SampleException
+from mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback import (
+    SampleHandlingCallback,
+)
+
+TEST_SAMPLE_ID = 123456
+
+
+@pytest.mark.parametrize(
+    "exception_type, expected_sample_status, expected_raised_exception",
+    [
+        ["Beamline", BLSampleStatus.ERROR_BEAMLINE, AssertionError],
+        ["Sample", BLSampleStatus.ERROR_SAMPLE, SampleException],
+    ],
+)
+def test_depositing_sample_error_with_sample_or_beamline_exception(
+    RE: RunEngine,
+    exception_type: str,
+    expected_sample_status: BLSampleStatus,
+    expected_raised_exception: type,
+):
+    sample_handling_callback = SampleHandlingCallback()
+    RE.subscribe(sample_handling_callback)
+
+    mock_expeye = MagicMock()
+    with (
+        patch(
+            "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback"
+            ".ExpeyeInteraction",
+            return_value=mock_expeye,
+        ),
+        pytest.raises(expected_raised_exception),
+    ):
+        RE(deposit_sample_error(exception_type, TEST_SAMPLE_ID))
+        mock_expeye.update_sample_status.assert_called_once_with(
+            TEST_SAMPLE_ID, expected_sample_status
+        )
+
+
+def test_depositing_sample_loaded(
+    RE: RunEngine,
+):
+    sample_handling_callback = SampleHandlingCallback()
+    RE.subscribe(sample_handling_callback)
+
+    mock_expeye = MagicMock()
+    with patch(
+        "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback"
+        ".ExpeyeInteraction",
+        return_value=mock_expeye,
+    ):
+        RE(deposit_loaded_sample(TEST_SAMPLE_ID))
+        mock_expeye.update_sample_status.assert_called_once_with(
+            TEST_SAMPLE_ID, BLSampleStatus.LOADED
+        )