diff --git a/src/mx_bluesky/common/external_interaction/callbacks/common/plan_reactive_callback.py b/src/mx_bluesky/common/external_interaction/callbacks/common/plan_reactive_callback.py index 79e0a70e1..c16e1910b 100644 --- a/src/mx_bluesky/common/external_interaction/callbacks/common/plan_reactive_callback.py +++ b/src/mx_bluesky/common/external_interaction/callbacks/common/plan_reactive_callback.py @@ -36,7 +36,7 @@ class will be activated, and on recieving the stop document for the super().__init__(emit=emit) self.emit_cb = emit # to avoid GC; base class only holds a WeakRef self.active = False - self.activity_uid = 0 + self.activity_uid = "" self.log = log def _run_activity_gated(self, name: str, func, doc, override=False): @@ -76,7 +76,7 @@ def stop(self, doc: RunStop) -> RunStop | None: do_stop = self.active if doc.get("run_start") == self.activity_uid: self.active = False - self.activity_uid = 0 + self.activity_uid = "" return ( self._run_activity_gated( "stop", self.activity_gated_stop, doc, override=True diff --git a/src/mx_bluesky/hyperion/external_interaction/callbacks/sample_handling/sample_handling_callback.py b/src/mx_bluesky/hyperion/external_interaction/callbacks/sample_handling/sample_handling_callback.py index a874ef128..dd1ac36ae 100644 --- a/src/mx_bluesky/hyperion/external_interaction/callbacks/sample_handling/sample_handling_callback.py +++ b/src/mx_bluesky/hyperion/external_interaction/callbacks/sample_handling/sample_handling_callback.py @@ -19,22 +19,27 @@ def __init__(self): super().__init__(log=ISPYB_ZOCALO_CALLBACK_LOGGER) self._sample_id: int | None = None self._descriptor: str | None = None + self._run_id: str | None = None def activity_gated_start(self, doc: RunStart): - if not self._sample_id: + if not self._sample_id and self.active: sample_id = doc.get("metadata", {}).get("sample_id") self.log.info(f"Recording sample ID at run start {sample_id}") self._sample_id = sample_id + self._run_id = self.activity_uid def activity_gated_stop(self, doc: RunStop) -> RunStop: - if doc["exit_status"] != "success": - exception_type, message = SampleException.type_and_message_from_reason( - doc.get("reason", "") - ) - self.log.info( - f"Sample handling callback intercepted exception of type {exception_type}: {message}" - ) - self._record_exception(exception_type) + if self._run_id == doc.get("run_start"): + if doc["exit_status"] != "success": + exception_type, message = SampleException.type_and_message_from_reason( + doc.get("reason", "") + ) + self.log.info( + f"Sample handling callback intercepted exception of type {exception_type}: {message}" + ) + self._record_exception(exception_type) + self._sample_id = None + self._run_id = None return doc def _record_exception(self, exception_type: str): diff --git a/tests/unit_tests/hyperion/external_interaction/callbacks/sample_handling/test_sample_handling_callback.py b/tests/unit_tests/hyperion/external_interaction/callbacks/sample_handling/test_sample_handling_callback.py index fc1769839..979cc87ef 100644 --- a/tests/unit_tests/hyperion/external_interaction/callbacks/sample_handling/test_sample_handling_callback.py +++ b/tests/unit_tests/hyperion/external_interaction/callbacks/sample_handling/test_sample_handling_callback.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock, patch +import bluesky.preprocessors as bpp import pytest from bluesky.preprocessors import run_decorator from bluesky.run_engine import RunEngine @@ -27,6 +28,69 @@ def plan_with_general_exception(exception_type: type, msg: str): raise exception_type(msg) +def plan_for_sample_id(sample_id): + def plan_with_exception(): + yield from [] + raise SampleException(f"Test exception for sample_id {sample_id}") + + yield from bpp.run_wrapper( + plan_with_exception(), + md={ + "metadata": {"sample_id": sample_id}, + "activate_callbacks": ["SampleHandlingCallback"], + }, + ) + + +def plan_with_exception_from_inner_plan(): + @run_decorator( + md={ + "metadata": {"sample_id": TEST_SAMPLE_ID}, + } + ) + def inner_plan(): + yield from [] + raise SampleException("Exception from inner plan") + + @run_decorator( + md={ + "metadata": {"sample_id": TEST_SAMPLE_ID}, + "activate_callbacks": ["SampleHandlingCallback"], + } + ) + @bpp.set_run_key_decorator("outer_plan") + def outer_plan(): + yield from inner_plan() + + yield from outer_plan() + + +def plan_with_rethrown_exception(): + @run_decorator( + md={ + "metadata": {"sample_id": TEST_SAMPLE_ID}, + } + ) + def inner_plan(): + yield from [] + raise AssertionError("Exception from inner plan") + + @run_decorator( + md={ + "metadata": {"sample_id": TEST_SAMPLE_ID}, + "activate_callbacks": ["SampleHandlingCallback"], + } + ) + @bpp.set_run_key_decorator("outer_plan") + def outer_plan(): + try: + yield from inner_plan() + except AssertionError as e: + raise SampleException("Exception from outer plan") from e + + yield from outer_plan() + + @run_decorator( md={ "metadata": {"sample_id": TEST_SAMPLE_ID}, @@ -80,3 +144,64 @@ def test_sample_handling_callback_closes_run_normally(RE: RunEngine): RE(plan_with_normal_completion()) record_exception.assert_not_called() + + +@patch( + "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback" + ".ExpeyeInteraction", +) +def test_sample_handling_callback_resets_sample_id( + mock_expeye_cls: MagicMock, RE: RunEngine +): + mock_expeye = mock_expeye_cls.return_value + callback = SampleHandlingCallback() + RE.subscribe(callback) + + with pytest.raises(SampleException): + RE(plan_for_sample_id(TEST_SAMPLE_ID)) + mock_expeye.update_sample_status.assert_called_once_with( + TEST_SAMPLE_ID, BLSampleStatus.ERROR_SAMPLE + ) + mock_expeye.reset_mock() + + with pytest.raises(SampleException): + RE(plan_for_sample_id(TEST_SAMPLE_ID + 1)) + mock_expeye.update_sample_status.assert_called_once_with( + TEST_SAMPLE_ID + 1, BLSampleStatus.ERROR_SAMPLE + ) + + +@patch( + "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback" + ".ExpeyeInteraction", +) +def test_sample_handling_callback_triggered_only_by_outermost_plan_when_exception_thrown_in_inner_plan( + mock_expeye_cls: MagicMock, RE: RunEngine +): + mock_expeye = mock_expeye_cls.return_value + callback = SampleHandlingCallback() + RE.subscribe(callback) + + with pytest.raises(SampleException): + RE(plan_with_exception_from_inner_plan()) + mock_expeye.update_sample_status.assert_called_once_with( + TEST_SAMPLE_ID, BLSampleStatus.ERROR_SAMPLE + ) + + +@patch( + "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback" + ".ExpeyeInteraction", +) +def test_sample_handling_callback_triggered_only_by_outermost_plan_when_exception_rethrown_from_outermost_plan( + mock_expeye_cls: MagicMock, RE: RunEngine +): + mock_expeye = mock_expeye_cls.return_value + callback = SampleHandlingCallback() + RE.subscribe(callback) + + with pytest.raises(SampleException): + RE(plan_with_rethrown_exception()) + mock_expeye.update_sample_status.assert_called_once_with( + TEST_SAMPLE_ID, BLSampleStatus.ERROR_SAMPLE + )