Merge pull request #959 from DiamondLightSource/950_init_callbacks_w_…

…document

Initialize callbacks using params emitted in Bluesky documents

src/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -43,7 +43,13 @@
     XrayCentreCallbackCollection,
 )
 from hyperion.parameters import external_parameters
-from hyperion.parameters.constants import SIM_BEAMLINE
+from hyperion.parameters.constants import (
+    DO_FGS,
+    GRIDSCAN_AND_MOVE,
+    GRIDSCAN_MAIN_PLAN,
+    GRIDSCAN_OUTER_PLAN,
+    SIM_BEAMLINE,
+)
 from hyperion.tracing import TRACER
 from hyperion.utils.aperturescatterguard import (
     load_default_aperture_scatterguard_positions_if_unset,
@@ -95,6 +101,7 @@ def set_aperture_for_bbox_size(
     bbox_size: list[int],
 ):
     # bbox_size is [x,y,z], for i03 we only care about x
+    assert aperture_device.aperture_positions is not None
     if bbox_size[0] < 2:
         aperture_size_positions = aperture_device.aperture_positions.MEDIUM
         selected_aperture = "MEDIUM_APERTURE"
@@ -137,13 +144,13 @@ def tidy_up_plans(fgs_composite: FlyScanXRayCentreComposite):
     yield from set_zebra_shutter_to_manual(fgs_composite.zebra)
 
 
-@bpp.set_run_key_decorator("run_gridscan")
-@bpp.run_decorator(md={"subplan_name": "run_gridscan"})
+@bpp.set_run_key_decorator(GRIDSCAN_MAIN_PLAN)
+@bpp.run_decorator(md={"subplan_name": GRIDSCAN_MAIN_PLAN})
 def run_gridscan(
     fgs_composite: FlyScanXRayCentreComposite,
     parameters: GridscanInternalParameters,
     md={
-        "plan_name": "run_gridscan",
+        "plan_name": GRIDSCAN_MAIN_PLAN,
     },
 ):
     sample_motors = fgs_composite.sample_motors
@@ -173,8 +180,8 @@ def run_gridscan(
 
     yield from wait_for_gridscan_valid(fgs_motors)
 
-    @bpp.set_run_key_decorator("do_fgs")
-    @bpp.run_decorator(md={"subplan_name": "do_fgs"})
+    @bpp.set_run_key_decorator(DO_FGS)
+    @bpp.run_decorator(md={"subplan_name": DO_FGS})
     @bpp.contingency_decorator(
         except_plan=lambda e: (yield from bps.stop(fgs_composite.eiger)),
         else_plan=lambda: (yield from bps.unstage(fgs_composite.eiger)),
@@ -204,8 +211,8 @@ def do_fgs():
     yield from bps.abs_set(fgs_motors.z_steps, 0, wait=False)
 
 
-@bpp.set_run_key_decorator("run_gridscan_and_move")
-@bpp.run_decorator(md={"subplan_name": "run_gridscan_and_move"})
+@bpp.set_run_key_decorator(GRIDSCAN_AND_MOVE)
+@bpp.run_decorator(md={"subplan_name": GRIDSCAN_AND_MOVE})
 def run_gridscan_and_move(
     fgs_composite: FlyScanXRayCentreComposite,
     parameters: GridscanInternalParameters,
@@ -268,15 +275,15 @@ def flyscan_xray_centre(
     """
     composite.eiger.set_detector_parameters(parameters.hyperion_params.detector_params)
 
-    subscriptions = XrayCentreCallbackCollection.from_params(parameters)
+    subscriptions = XrayCentreCallbackCollection.setup()
 
     @bpp.subs_decorator(  # subscribe the RE to nexus, ispyb, and zocalo callbacks
         list(subscriptions)  # must be the outermost decorator to receive the metadata
     )
-    @bpp.set_run_key_decorator("run_gridscan_move_and_tidy")
+    @bpp.set_run_key_decorator(GRIDSCAN_OUTER_PLAN)
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
-            "subplan_name": "run_gridscan_move_and_tidy",
+            "subplan_name": GRIDSCAN_OUTER_PLAN,
             "hyperion_internal_parameters": parameters.json(),
             "activate_callbacks": [
                 "XrayCentreZocaloCallback",
@@ -307,13 +314,13 @@ def run_gridscan_and_move_and_tidy(fgs_composite, params, comms):
     args = parser.parse_args()
 
     RE = RunEngine({})
-    RE.waiting_hook = ProgressBarManager()
+    RE.waiting_hook = ProgressBarManager()  # type: ignore
     from hyperion.parameters.plan_specific.gridscan_internal_params import (
         GridscanInternalParameters,
     )
 
     parameters = GridscanInternalParameters(**external_parameters.from_file())
-    subscriptions = XrayCentreCallbackCollection.from_params(parameters)
+    subscriptions = XrayCentreCallbackCollection.setup()
 
     context = setup_context(wait_for_connection=True)
     composite = create_devices(context)

src/hyperion/experiment_plans/rotation_scan_plan.py

@@ -39,6 +39,7 @@
     RotationCallbackCollection,
 )
 from hyperion.log import LOGGER
+from hyperion.parameters.constants import ROTATION_OUTER_PLAN, ROTATION_PLAN_MAIN
 from hyperion.parameters.plan_specific.rotation_scan_internal_params import (
     RotationScanParams,
 )
@@ -133,8 +134,8 @@ def set_speed(axis: EpicsMotor, image_width, exposure_time, wait=True):
     )
 
 
-@bpp.set_run_key_decorator("rotation_scan_main")
-@bpp.run_decorator(md={"subplan_name": "rotation_scan_main"})
+@bpp.set_run_key_decorator(ROTATION_PLAN_MAIN)
+@bpp.run_decorator(md={"subplan_name": ROTATION_PLAN_MAIN})
 def rotation_scan_plan(
     composite: RotationScanComposite,
     params: RotationInternalParameters,
@@ -252,13 +253,13 @@ def cleanup_plan(composite: RotationScanComposite, **kwargs):
 
 
 def rotation_scan(composite: RotationScanComposite, parameters: Any) -> MsgGenerator:
-    subscriptions = RotationCallbackCollection.from_params(parameters)
+    subscriptions = RotationCallbackCollection.setup()
 
     @bpp.subs_decorator(list(subscriptions))
     @bpp.set_run_key_decorator("rotation_scan")
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
-            "subplan_name": "rotation_scan_with_cleanup",
+            "subplan_name": ROTATION_OUTER_PLAN,
             "hyperion_internal_parameters": parameters.json(),
             "activate_callbacks": [
                 "RotationZocaloCallback",

src/hyperion/experiment_plans/stepped_grid_scan_plan.py

@@ -15,7 +15,11 @@
 from hyperion.log import LOGGER
 from hyperion.parameters import external_parameters
 from hyperion.parameters.beamline_parameters import GDABeamlineParameters
-from hyperion.parameters.constants import BEAMLINE_PARAMETER_PATHS, SIM_BEAMLINE
+from hyperion.parameters.constants import (
+    BEAMLINE_PARAMETER_PATHS,
+    GRIDSCAN_MAIN_PLAN,
+    SIM_BEAMLINE,
+)
 from hyperion.tracing import TRACER
 from hyperion.utils.context import device_composite_from_context, setup_context
 
@@ -47,7 +51,7 @@ def run_gridscan(
     composite: SteppedGridScanComposite,
     parameters: SteppedGridScanInternalParameters,
     md={
-        "plan_name": "run_gridscan",
+        "plan_name": GRIDSCAN_MAIN_PLAN,
     },
 ):
     sample_motors: Smargon = composite.smargon

src/hyperion/experiment_plans/tests/conftest.py

@@ -1,4 +1,5 @@
 from functools import partial
+from typing import Callable, Generator
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -31,8 +32,13 @@
 from hyperion.external_interaction.callbacks.xray_centre.callback_collection import (
     XrayCentreCallbackCollection,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import Store3DGridscanInIspyb
+from hyperion.external_interaction.ispyb.store_in_ispyb import (
+    IspybIds,
+    Store3DGridscanInIspyb,
+)
 from hyperion.external_interaction.system_tests.conftest import TEST_RESULT_LARGE
+from hyperion.external_interaction.zocalo.zocalo_interaction import ZocaloInteractor
+from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN
 from hyperion.parameters.external_parameters import from_file as raw_params_from_file
 from hyperion.parameters.internal_parameters import InternalParameters
 from hyperion.parameters.plan_specific.grid_scan_with_edge_detect_params import (
@@ -47,7 +53,7 @@
 
 
 def mock_set(motor: EpicsMotor, val):
-    motor.user_readback.sim_put(val)
+    motor.user_readback.sim_put(val)  # type: ignore
     return Status(done=True, success=True)
 
 
@@ -87,7 +93,7 @@ def eiger(done_status):
 
 
 @pytest.fixture
-def smargon() -> Smargon:
+def smargon() -> Generator[Smargon, None, None]:
     smargon = i03.smargon(fake_with_ophyd_sim=True)
     smargon.x.user_setpoint._use_limits = False
     smargon.y.user_setpoint._use_limits = False
@@ -294,26 +300,45 @@ def fake_fgs_composite(smargon: Smargon, test_fgs_params: InternalParameters):
         False
     )
 
-    fake_composite.fast_grid_scan.scan_invalid.sim_put(False)
-    fake_composite.fast_grid_scan.position_counter.sim_put(0)
+    fake_composite.fast_grid_scan.scan_invalid.sim_put(False)  # type: ignore
+    fake_composite.fast_grid_scan.position_counter.sim_put(0)  # type: ignore
 
     return fake_composite
 
 
+def modified_interactor_mock(assign_run_end: Callable | None = None):
+    mock = MagicMock(spec=ZocaloInteractor)
+    mock.wait_for_result.return_value = TEST_RESULT_LARGE
+    if assign_run_end:
+        mock.run_end = assign_run_end
+    return mock
+
+
+def modified_store_grid_scan_mock(*args, dcids=(0, 0), dcgid=0, **kwargs):
+    mock = MagicMock(spec=Store3DGridscanInIspyb)
+    mock.begin_deposition.return_value = IspybIds(
+        data_collection_ids=dcids, data_collection_group_id=dcgid, grid_ids=(0, 0)
+    )
+    return mock
+
+
 @pytest.fixture
 def mock_subscriptions(test_fgs_params):
-    subscriptions = XrayCentreCallbackCollection.from_params(test_fgs_params)
-    subscriptions.zocalo_handler.zocalo_interactor.wait_for_result = MagicMock()
-    subscriptions.zocalo_handler.zocalo_interactor.run_end = MagicMock()
-    subscriptions.zocalo_handler.zocalo_interactor.run_start = MagicMock()
-    subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = (
-        TEST_RESULT_LARGE
-    )
+    with patch(
+        "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor",
+        modified_interactor_mock,
+    ):
+        subscriptions = XrayCentreCallbackCollection.setup()
+        start_doc = {
+            "subplan_name": GRIDSCAN_OUTER_PLAN,
+            "hyperion_internal_parameters": test_fgs_params.json(),
+        }
+        subscriptions.ispyb_handler.activity_gated_start(start_doc)
+        subscriptions.zocalo_handler.activity_gated_start(start_doc)
     subscriptions.ispyb_handler.ispyb = MagicMock(spec=Store3DGridscanInIspyb)
-    subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: [[0, 0], 0, 0]
-    subscriptions.ispyb_handler.active = True
-    subscriptions.nexus_handler.active = True
-    subscriptions.zocalo_handler.active = True
+    subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: IspybIds(
+        data_collection_ids=(0, 0), data_collection_group_id=0, grid_ids=(0, 0)
+    )
 
     return subscriptions
 
@@ -330,7 +355,7 @@ def mock_rotation_subscriptions(test_rotation_params):
         "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationZocaloCallback",
         autospec=True,
     ):
-        subscriptions = RotationCallbackCollection.from_params(test_rotation_params)
+        subscriptions = RotationCallbackCollection.setup()
     return subscriptions