#1192 fix document passing on logic and a test

src/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -340,7 +340,7 @@ def flyscan_xray_centre(
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
             "subplan_name": GRIDSCAN_OUTER_PLAN,
-            "trigger_zocalo_on": GRIDSCAN_MAIN_PLAN,
+            "trigger_zocalo_on": DO_FGS,
             "hyperion_internal_parameters": parameters.json(),
             "activate_callbacks": [
                 "GridscanISPyBCallback",

src/hyperion/external_interaction/callbacks/ispyb_callback_base.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, TypeVar
 
 from hyperion.external_interaction.callbacks.plan_reactive_callback import (
     PlanReactiveCallback,
@@ -24,6 +24,7 @@
     RotationInternalParameters,
 )
 
+D = TypeVar("D")
 if TYPE_CHECKING:
     from event_model.documents import Event, EventDescriptor, RunStart, RunStop
 
@@ -62,9 +63,11 @@ def activity_gated_start(self, doc: RunStart):
         ISPYB_LOGGER.debug("ISPyB Callback Start Triggered")
         if self.uid_to_finalize_on is None:
             self.uid_to_finalize_on = doc.get("uid")
+        return self._tag_doc(doc)
 
     def activity_gated_descriptor(self, doc: EventDescriptor):
         self.descriptors[doc["uid"]] = doc
+        return self._tag_doc(doc)
 
     def activity_gated_event(self, doc: Event) -> Event:
         """Subclasses should extend this to add a call to set_dcig_tag from
@@ -112,7 +115,7 @@ def activity_gated_event(self, doc: Event) -> Event:
             ISPYB_LOGGER.info("Updating ispyb entry.")
             self.ispyb_ids = self.ispyb.update_deposition()
             ISPYB_LOGGER.info(f"Recieved ISPYB IDs: {self.ispyb_ids}")
-        return doc
+        return self._tag_doc(doc)
 
     def activity_gated_stop(self, doc: RunStop):
         """Subclasses must check that they are recieving a stop document for the correct
@@ -132,6 +135,7 @@ def activity_gated_stop(self, doc: RunStop):
             ISPYB_LOGGER.warning(
                 f"Failed to finalise ISPyB deposition on stop document: {doc} with exception: {e}"
             )
+        return self._tag_doc(doc)
 
     def _append_to_comment(self, id: int, comment: str):
         assert isinstance(self.ispyb, StoreInIspyb)
@@ -145,3 +149,9 @@ def _append_to_comment(self, id: int, comment: str):
     def append_to_comment(self, comment: str):
         for id in self.ispyb_ids.data_collection_ids:
             self._append_to_comment(id, comment)
+
+    def _tag_doc(self, doc: D) -> D:
+        assert isinstance(doc, dict)
+        if self.ispyb_ids:
+            doc["ispyb_dcids"] = self.ispyb_ids.data_collection_ids
+        return doc

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

@@ -79,6 +79,7 @@ def activity_gated_start(self, doc: RunStart):
                 else Store2DGridscanInIspyb(self.ispyb_config, self.params)
             )
             self.ispyb_ids = self.ispyb.begin_deposition()
+        return self._tag_doc(doc)
 
     def activity_gated_event(self, doc: Event):
         doc = super().activity_gated_event(doc)
@@ -119,7 +120,7 @@ def activity_gated_event(self, doc: Event):
             )
 
         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._start_of_fgs_uid:
@@ -131,4 +132,5 @@ def activity_gated_stop(self, doc: RunStop):
             )
             if self.ispyb_ids == IspybIds():
                 raise ISPyBDepositionNotMade("ispyb was not initialised at run start")
-            super().activity_gated_stop(doc)
+            return super().activity_gated_stop(doc)
+        return self._tag_doc(doc)

src/hyperion/external_interaction/callbacks/zocalo_callback.py

@@ -46,7 +46,10 @@ def start(self, doc: RunStart):
             ISPYB_LOGGER.info(f"Zocalo environment set to {zocalo_environment}.")
             self.zocalo_interactor = ZocaloTrigger(zocalo_environment)
             self.run_uid = doc.get("uid")
-            if isinstance(ispyb_ids := doc.get("ispyb_dcids"), tuple):
+            if (
+                isinstance(ispyb_ids := doc.get("ispyb_dcids"), tuple)
+                and len(ispyb_ids) > 0
+            ):
                 self.ispyb_ids = ispyb_ids
                 for id in self.ispyb_ids:
                     self.zocalo_interactor.run_start(id)
@@ -60,7 +63,7 @@ def stop(self, doc: RunStop):
             ISPYB_LOGGER.info(
                 f"Zocalo handler received stop document, for run {doc.get('run_start')}."
             )
-            if self.ispyb_ids:
+            if self.ispyb_ids and self.zocalo_interactor:
                 for id in self.ispyb_ids:
                     self.zocalo_interactor.run_end(id)
             self._reset_state()

tests/unit_tests/experiment_plans/test_flyscan_xray_centre_plan.py

@@ -51,6 +51,7 @@
 )
 from hyperion.parameters import external_parameters
 from hyperion.parameters.constants import (
+    DO_FGS,
     GRIDSCAN_OUTER_PLAN,
 )
 from hyperion.parameters.plan_specific.gridscan_internal_params import (
@@ -741,13 +742,18 @@ def test_kickoff_and_complete_gridscan_triggers_zocalo(
     RE: RunEngine,
     fake_fgs_composite: FlyScanXRayCentreComposite,
 ):
-    mock_ispyb_handler = MagicMock()
-    mock_ispyb_handler.ispyb_ids.data_collection_ids = (100, 200)
+    cbs = XrayCentreCallbackCollection()
+    ispyb_cb = cbs.ispyb_handler
+    ispyb_cb.active = True
+    ispyb_cb.ispyb_ids.data_collection_ids = (1, 2)
+    assert isinstance(zocalo_cb := ispyb_cb.emit_cb, ZocaloCallback)
     zocalo_env = "dev_env"
-    zocalo_callback = ZocaloCallback()
-    zocalo_callback.active = True
+
+    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())
-    RE.subscribe(zocalo_callback)
+    RE.subscribe(ispyb_cb)
     RE(
         kickoff_and_complete_gridscan(
             fake_fgs_composite.fast_grid_scan,