From 898b923d4ccf70efe82395f2a7345e868e871d43 Mon Sep 17 00:00:00 2001 From: Shihab Suliman Date: Wed, 8 Jan 2025 17:56:06 +0000 Subject: [PATCH] Move test file location to common --- .../common/plans/write_sample_status.py | 34 ++++ tests/unit_tests/common/test_do_fgs.py | 170 ++++++++++++++++++ .../common/test_write_sample_status.py | 65 +++++++ 3 files changed, 269 insertions(+) create mode 100644 src/mx_bluesky/common/plans/write_sample_status.py create mode 100644 tests/unit_tests/common/test_do_fgs.py create mode 100644 tests/unit_tests/common/test_write_sample_status.py diff --git a/src/mx_bluesky/common/plans/write_sample_status.py b/src/mx_bluesky/common/plans/write_sample_status.py new file mode 100644 index 000000000..cb3e780e0 --- /dev/null +++ b/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() diff --git a/tests/unit_tests/common/test_do_fgs.py b/tests/unit_tests/common/test_do_fgs.py new file mode 100644 index 000000000..3a092fbb2 --- /dev/null +++ b/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], + ) + ) diff --git a/tests/unit_tests/common/test_write_sample_status.py b/tests/unit_tests/common/test_write_sample_status.py new file mode 100644 index 000000000..33e05daea --- /dev/null +++ b/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 + )