DiamondLightSource/hyperion#1192 fix most rotation tests

src/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -59,6 +59,7 @@
     GRIDSCAN_MAIN_PLAN,
     GRIDSCAN_OUTER_PLAN,
     SIM_BEAMLINE,
+    TRIGGER_ZOCALO_ON,
 )
 from hyperion.tracing import TRACER
 from hyperion.utils.aperturescatterguard import (
@@ -340,7 +341,7 @@ def flyscan_xray_centre(
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
             "subplan_name": GRIDSCAN_OUTER_PLAN,
-            "trigger_zocalo_on": DO_FGS,
+            TRIGGER_ZOCALO_ON: DO_FGS,
             "hyperion_internal_parameters": parameters.json(),
             "activate_callbacks": [
                 "GridscanISPyBCallback",

src/hyperion/experiment_plans/panda_flyscan_xray_centre_plan.py

@@ -51,6 +51,7 @@
     GRIDSCAN_MAIN_PLAN,
     GRIDSCAN_OUTER_PLAN,
     SIM_BEAMLINE,
+    TRIGGER_ZOCALO_ON,
 )
 from hyperion.tracing import TRACER
 from hyperion.utils.context import device_composite_from_context, setup_context
@@ -271,7 +272,7 @@ def panda_flyscan_xray_centre(
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
             "subplan_name": GRIDSCAN_OUTER_PLAN,
-            "trigger_zocalo_on": DO_FGS,
+            TRIGGER_ZOCALO_ON: DO_FGS,
             "hyperion_internal_parameters": parameters.json(),
             "activate_callbacks": [
                 "GridscanISPyBCallback",

src/hyperion/experiment_plans/rotation_scan_plan.py

@@ -39,7 +39,11 @@
     setup_zebra_for_rotation,
 )
 from hyperion.log import LOGGER
-from hyperion.parameters.constants import ROTATION_OUTER_PLAN, ROTATION_PLAN_MAIN
+from hyperion.parameters.constants import (
+    ROTATION_OUTER_PLAN,
+    ROTATION_PLAN_MAIN,
+    TRIGGER_ZOCALO_ON,
+)
 from hyperion.parameters.plan_specific.rotation_scan_internal_params import (
     RotationScanParams,
 )
@@ -266,7 +270,7 @@ def rotation_scan(composite: RotationScanComposite, parameters: Any) -> MsgGener
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
             "subplan_name": ROTATION_OUTER_PLAN,
-            "trigger_zocalo_on": ROTATION_PLAN_MAIN,
+            TRIGGER_ZOCALO_ON: ROTATION_PLAN_MAIN,
             "hyperion_internal_parameters": parameters.json(),
             "activate_callbacks": [
                 "RotationISPyBCallback",

src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py

@@ -91,13 +91,15 @@ def activity_gated_start(self, doc: RunStart):
         ISPYB_LOGGER.info("ISPYB handler received start document.")
         if doc.get("subplan_name") == ROTATION_PLAN_MAIN:
             self.uid_to_finalize_on = doc.get("uid")
+        return self._tag_doc(doc)
 
     def activity_gated_event(self, doc: Event):
         doc = super().activity_gated_event(doc)
         set_dcgid_tag(self.ispyb_ids.data_collection_group_id)
-        return doc
+        return self._tag_doc(doc)
 
     def activity_gated_stop(self, doc: RunStop):
         if doc.get("run_start") == self.uid_to_finalize_on:
             self.uid_to_finalize_on = None
-            super().activity_gated_stop(doc)
+            return super().activity_gated_stop(doc)
+        return self._tag_doc(doc)

src/hyperion/external_interaction/callbacks/zocalo_callback.py

@@ -9,6 +9,7 @@
 
 from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
 from hyperion.log import ISPYB_LOGGER
+from hyperion.parameters.constants import TRIGGER_ZOCALO_ON
 
 if TYPE_CHECKING:
     from event_model.documents import RunStart, RunStop
@@ -39,7 +40,7 @@ def __init__(
 
     def start(self, doc: RunStart):
         ISPYB_LOGGER.info("Zocalo handler received start document.")
-        if triggering_plan := doc.get("trigger_zocalo_on"):
+        if triggering_plan := doc.get(TRIGGER_ZOCALO_ON):
             self.triggering_plan = triggering_plan
         if self.triggering_plan and doc.get("subplan_name") == self.triggering_plan:
             assert isinstance(zocalo_environment := doc.get("zocalo_environment"), str)

src/hyperion/parameters/constants.py

@@ -28,6 +28,11 @@
 ROTATION_PLAN_MAIN = "rotation_scan_main"
 ########################################################################################
 
+# Plan metadata keys ###################################################################
+# Gridscan
+TRIGGER_ZOCALO_ON = "trigger_zocalo_on"
+########################################################################################
+
 
 class Actions(Enum):
     START = "start"

tests/unit_tests/experiment_plans/conftest.py

@@ -11,9 +11,6 @@
 from ophyd.sim import make_fake_device
 from ophyd_async.core.async_status import AsyncStatus
 
-from hyperion.external_interaction.callbacks.rotation.callback_collection import (
-    RotationCallbackCollection,
-)
 from hyperion.external_interaction.callbacks.xray_centre.callback_collection import (
     XrayCentreCallbackCollection,
 )
@@ -31,6 +28,7 @@
     GRIDSCAN_OUTER_PLAN,
     ISPYB_HARDWARE_READ_PLAN,
     ISPYB_TRANSMISSION_FLUX_READ_PLAN,
+    TRIGGER_ZOCALO_ON,
 )
 from hyperion.parameters.plan_specific.gridscan_internal_params import (
     GridscanInternalParameters,
@@ -153,29 +151,13 @@ def mock_subscriptions(test_fgs_params):
         start_doc = {
             "subplan_name": GRIDSCAN_OUTER_PLAN,
             "hyperion_internal_parameters": test_fgs_params.json(),
-            "trigger_zocalo_on": DO_FGS,
+            TRIGGER_ZOCALO_ON: DO_FGS,
         }
         subscriptions.ispyb_handler.activity_gated_start(start_doc)
 
     return subscriptions
 
 
-@pytest.fixture
-def mock_rotation_subscriptions(test_rotation_params):
-    with patch(
-        "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationNexusFileCallback",
-        autospec=True,
-    ), patch(
-        "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationISPyBCallback",
-        autospec=True,
-    ), patch(
-        "hyperion.external_interaction.callbacks.rotation.callback_collection.ZocaloCallback",
-        autospec=True,
-    ):
-        subscriptions = RotationCallbackCollection()
-    return subscriptions
-
-
 def fake_read(obj, initial_positions, _):
     initial_positions[obj] = 0
     yield Msg("null", obj)

tests/unit_tests/experiment_plans/test_flyscan_xray_centre_plan.py

@@ -50,10 +50,7 @@
     IspybIds,
 )
 from hyperion.parameters import external_parameters
-from hyperion.parameters.constants import (
-    DO_FGS,
-    GRIDSCAN_OUTER_PLAN,
-)
+from hyperion.parameters.constants import DO_FGS, GRIDSCAN_OUTER_PLAN, TRIGGER_ZOCALO_ON
 from hyperion.parameters.plan_specific.gridscan_internal_params import (
     GridscanInternalParameters,
 )
@@ -747,7 +744,7 @@ def test_kickoff_and_complete_gridscan_triggers_zocalo(
     assert isinstance(zocalo_cb := ispyb_cb.emit_cb, ZocaloCallback)
     zocalo_env = "dev_env"
 
-    zocalo_cb.start({"trigger_zocalo_on": DO_FGS})  # type: ignore
+    zocalo_cb.start({TRIGGER_ZOCALO_ON: DO_FGS})  # type: ignore
     assert zocalo_cb.triggering_plan == DO_FGS
 
     mock_zocalo_trigger_class.return_value = (mock_zocalo_trigger := MagicMock())

tests/unit_tests/experiment_plans/test_rotation_scan_plan.py

@@ -19,9 +19,6 @@
     rotation_scan,
     rotation_scan_plan,
 )
-from hyperion.external_interaction.callbacks.rotation.callback_collection import (
-    RotationCallbackCollection,
-)
 from hyperion.parameters.plan_specific.rotation_scan_internal_params import (
     RotationInternalParameters,
 )
@@ -37,47 +34,37 @@
 
 
 def do_rotation_main_plan_for_tests(
-    run_eng_w_subs: tuple[RunEngine, RotationCallbackCollection],
+    run_eng: RunEngine,
     expt_params: RotationInternalParameters,
     devices: RotationScanComposite,
     plan=rotation_scan_plan,
 ):
-    run_engine, _ = run_eng_w_subs
     with patch(
         "bluesky.preprocessors.__read_and_stash_a_motor",
         fake_read,
     ):
-        run_engine(
+        run_eng(
             plan(
                 devices,
                 expt_params,
             ),
         )
 
 
-@pytest.fixture
-def RE_with_subs(
-    RE: RunEngine, mock_rotation_subscriptions: RotationCallbackCollection
-):
-    for cb in mock_rotation_subscriptions:
-        RE.subscribe(cb)
-    return RE, mock_rotation_subscriptions
-
-
 @pytest.fixture
 def run_full_rotation_plan(
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_rotation_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
 ):
     do_rotation_main_plan_for_tests(
-        RE_with_subs, test_rotation_params, fake_create_rotation_devices, rotation_scan
+        RE, test_rotation_params, fake_create_rotation_devices, rotation_scan
     )
     return fake_create_rotation_devices
 
 
 def setup_and_run_rotation_plan_for_tests(
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
 ):
@@ -94,13 +81,13 @@ def side_set_w_return(obj, *args):
 
     with patch("bluesky.plan_stubs.wait", autospec=True):
         do_rotation_main_plan_for_tests(
-            RE_with_subs,
+            RE,
             test_params,
             fake_create_rotation_devices,
         )
 
     return {
-        "RE_with_subs": RE_with_subs,
+        "RE_with_subs": RE,
         "test_rotation_params": test_params,
         "smargon": fake_create_rotation_devices.smargon,
         "zebra": fake_create_rotation_devices.zebra,
@@ -109,25 +96,25 @@ def side_set_w_return(obj, *args):
 
 @pytest.fixture
 def setup_and_run_rotation_plan_for_tests_standard(
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_rotation_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
 ):
     return setup_and_run_rotation_plan_for_tests(
-        RE_with_subs,
+        RE,
         test_rotation_params,
         fake_create_rotation_devices,
     )
 
 
 @pytest.fixture
 def setup_and_run_rotation_plan_for_tests_nomove(
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_rotation_params_nomove: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
 ):
     return setup_and_run_rotation_plan_for_tests(
-        RE_with_subs,
+        RE,
         test_rotation_params_nomove,
         fake_create_rotation_devices,
     )
@@ -169,11 +156,10 @@ def test_move_to_end(smargon: Smargon, RE):
 @patch("hyperion.experiment_plans.rotation_scan_plan.rotation_scan_plan", autospec=True)
 def test_rotation_scan(
     plan: MagicMock,
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_rotation_params,
     fake_create_rotation_devices: RotationScanComposite,
 ):
-    RE, _ = RE_with_subs
 
     composite = fake_create_rotation_devices
     RE(rotation_scan(composite, test_rotation_params))
@@ -183,10 +169,7 @@ def test_rotation_scan(
 
 
 def test_rotation_plan_runs(setup_and_run_rotation_plan_for_tests_standard) -> None:
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection] = (
-        setup_and_run_rotation_plan_for_tests_standard["RE_with_subs"]
-    )
-    RE, _ = RE_with_subs
+    RE: RunEngine = setup_and_run_rotation_plan_for_tests_standard["RE_with_subs"]
     assert RE._exit_status == "success"
 
 
@@ -270,11 +253,10 @@ def test_rotation_plan_smargon_doesnt_move_xyz_if_not_given_in_params(
 def test_cleanup_happens(
     bps_wait: MagicMock,
     cleanup_plan: MagicMock,
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_rotation_params,
     fake_create_rotation_devices: RotationScanComposite,
 ):
-    RE, _ = RE_with_subs
 
     class MyTestException(Exception):
         pass
@@ -310,14 +292,14 @@ class MyTestException(Exception):
 )
 def test_acceleration_offset_calculated_correctly(
     mock_move_to_start: MagicMock,
-    RE_with_subs: tuple[RunEngine, RotationCallbackCollection],
+    RE: RunEngine,
     test_rotation_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
 ):
     composite = fake_create_rotation_devices
     composite.smargon.omega.acceleration.sim_put(0.2)  # type: ignore
     setup_and_run_rotation_plan_for_tests(
-        RE_with_subs,
+        RE,
         test_rotation_params,
         fake_create_rotation_devices,
     )