From 6eed420f7cd939182b06523106d3324f5ddb438b Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 29 Feb 2024 11:48:14 +0000 Subject: [PATCH] #1192 fix document passing on logic and a test --- .../experiment_plans/flyscan_xray_centre_plan.py | 2 +- .../callbacks/ispyb_callback_base.py | 14 ++++++++++++-- .../callbacks/xray_centre/ispyb_callback.py | 6 ++++-- .../callbacks/zocalo_callback.py | 7 +++++-- .../test_flyscan_xray_centre_plan.py | 16 +++++++++++----- 5 files changed, 33 insertions(+), 12 deletions(-) diff --git a/src/hyperion/experiment_plans/flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/flyscan_xray_centre_plan.py index adf9e463d..4c300204c 100755 --- a/src/hyperion/experiment_plans/flyscan_xray_centre_plan.py +++ b/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", diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index febbae832..64aa40bdb 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/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 diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py index 998b45f6c..33da97fd3 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/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) diff --git a/src/hyperion/external_interaction/callbacks/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/zocalo_callback.py index b1a3e8fa9..0c56a6c29 100644 --- a/src/hyperion/external_interaction/callbacks/zocalo_callback.py +++ b/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() diff --git a/tests/unit_tests/experiment_plans/test_flyscan_xray_centre_plan.py b/tests/unit_tests/experiment_plans/test_flyscan_xray_centre_plan.py index 32d2e0bd8..48d34e721 100644 --- a/tests/unit_tests/experiment_plans/test_flyscan_xray_centre_plan.py +++ b/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,