From ab13176e58a61f293009eaac5d49bd6c6164b602 Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 3 Nov 2023 14:48:16 +0000 Subject: [PATCH 01/15] # 950 remove params from callback init and replace plan names with constants --- .../flyscan_xray_centre_plan.py | 33 +++++++++++-------- .../experiment_plans/rotation_scan_plan.py | 7 ++-- .../stepped_grid_scan_plan.py | 8 +++-- .../experiment_plans/tests/conftest.py | 2 +- .../abstract_plan_callback_collection.py | 10 ++---- .../callbacks/ispyb_callback_base.py | 30 ++++++++++++----- .../callbacks/rotation/ispyb_callback.py | 14 ++++++-- .../callbacks/rotation/nexus_callback.py | 2 +- .../rotation/tests/test_rotation_callbacks.py | 8 ++--- .../xray_centre/callback_collection.py | 10 ++---- .../callbacks/xray_centre/ispyb_callback.py | 26 ++++++++++----- .../callbacks/xray_centre/nexus_callback.py | 4 +-- .../callbacks/xray_centre/tests/conftest.py | 14 ++++---- .../xray_centre/tests/test_nexus_handler.py | 10 ++++-- .../callbacks/xray_centre/zocalo_callback.py | 20 +++++++---- .../system_tests/test_write_rotation_nexus.py | 2 +- src/hyperion/parameters/constants.py | 11 +++++++ 17 files changed, 135 insertions(+), 76 deletions(-) diff --git a/src/hyperion/experiment_plans/flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/flyscan_xray_centre_plan.py index b9ee89337..6aecd98d7 100755 --- a/src/hyperion/experiment_plans/flyscan_xray_centre_plan.py +++ b/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, @@ -263,15 +270,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(), } ) @@ -297,13 +304,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) diff --git a/src/hyperion/experiment_plans/rotation_scan_plan.py b/src/hyperion/experiment_plans/rotation_scan_plan.py index 94d54911e..8907d03ce 100644 --- a/src/hyperion/experiment_plans/rotation_scan_plan.py +++ b/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, @@ -258,7 +259,7 @@ def rotation_scan(composite: RotationScanComposite, parameters: Any) -> MsgGener @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(), } ) diff --git a/src/hyperion/experiment_plans/stepped_grid_scan_plan.py b/src/hyperion/experiment_plans/stepped_grid_scan_plan.py index 298a962e8..21da0c744 100644 --- a/src/hyperion/experiment_plans/stepped_grid_scan_plan.py +++ b/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 diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index ab2a335e3..0ef519088 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/src/hyperion/experiment_plans/tests/conftest.py @@ -285,7 +285,7 @@ def fake_fgs_composite(smargon: Smargon, test_fgs_params: InternalParameters): @pytest.fixture def mock_subscriptions(test_fgs_params): - subscriptions = XrayCentreCallbackCollection.from_params(test_fgs_params) + subscriptions = XrayCentreCallbackCollection.setup() 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() diff --git a/src/hyperion/external_interaction/callbacks/abstract_plan_callback_collection.py b/src/hyperion/external_interaction/callbacks/abstract_plan_callback_collection.py index e346feb34..17cda3e07 100644 --- a/src/hyperion/external_interaction/callbacks/abstract_plan_callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/abstract_plan_callback_collection.py @@ -2,10 +2,6 @@ from abc import ABC, abstractmethod from dataclasses import fields -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from hyperion.parameters.internal_parameters import InternalParameters class AbstractPlanCallbackCollection(ABC): @@ -17,15 +13,15 @@ class AbstractPlanCallbackCollection(ABC): @classmethod @abstractmethod - def from_params(cls, params: InternalParameters): + def setup(cls): ... def __iter__(self): - for field in fields(self): + for field in fields(self): # type: ignore # subclasses must be dataclass yield getattr(self, field.name) class NullPlanCallbackCollection(AbstractPlanCallbackCollection): @classmethod - def from_params(cls, params: InternalParameters): + def setup(cls): pass diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index 753239531..228af7c50 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from typing import Dict, Optional +from typing import TYPE_CHECKING, Dict, Optional from bluesky.callbacks import CallbackBase @@ -12,16 +12,28 @@ ISPYB_TRANSMISSION_FLUX_READ_PLAN, SIM_ISPYB_CONFIG, ) -from hyperion.parameters.internal_parameters import InternalParameters + +if TYPE_CHECKING: + from hyperion.external_interaction.ispyb.store_in_ispyb import ( + Store2DGridscanInIspyb, + Store3DGridscanInIspyb, + StoreRotationInIspyb, + ) + from hyperion.parameters.plan_specific.gridscan_internal_params import ( + GridscanInternalParameters, + ) + from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, + ) class BaseISPyBCallback(CallbackBase): - def __init__(self, parameters: InternalParameters): - """Subclasses should run super().__init__() with parameters, then set + def __init__(self) -> None: + """Subclasses should run super().__init__(), then in their start() method set self.ispyb to the type of ispyb relevant to the experiment and define the type for self.ispyb_ids.""" - self.ispyb: StoreInIspyb - self.params = parameters + self.params: GridscanInternalParameters | RotationInternalParameters + self.ispyb: StoreRotationInIspyb | Store3DGridscanInIspyb | Store2DGridscanInIspyb self.descriptors: Dict[str, dict] = {} self.ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) if self.ispyb_config == SIM_ISPYB_CONFIG: @@ -88,8 +100,10 @@ def stop(self, doc: dict): self.ispyb, StoreInIspyb ), "ISPyB handler recieved stop document, but deposition object doesn't exist!" LOGGER.debug("ISPyB handler received stop document.") - exit_status = doc.get("exit_status") - reason = doc.get("reason") + exit_status = ( + doc.get("exit_status") or "Exit status not available in stop document!" + ) + reason = doc.get("reason") or "Unknown failure reason!" set_dcgid_tag(None) try: self.ispyb.end_deposition(exit_status, reason) diff --git a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py index e0d0a14a2..af99df885 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py @@ -5,6 +5,7 @@ ) from hyperion.external_interaction.ispyb.store_in_ispyb import StoreRotationInIspyb from hyperion.log import LOGGER, set_dcgid_tag +from hyperion.parameters.constants import ROTATION_OUTER_PLAN, ROTATION_PLAN_MAIN from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, ) @@ -27,9 +28,9 @@ class RotationISPyBCallback(BaseISPyBCallback): Usually used as part of a RotationCallbackCollection. """ - def __init__(self, parameters: RotationInternalParameters): + def __init__(self) -> None: self.params: RotationInternalParameters - super().__init__(parameters) + super().__init__() self.ispyb: StoreRotationInIspyb = StoreRotationInIspyb( self.ispyb_config, self.params ) @@ -40,8 +41,15 @@ def append_to_comment(self, comment: str): self._append_to_comment(self.ispyb_ids[0], comment) def start(self, doc: dict): + if doc.get("subplan_name") == ROTATION_OUTER_PLAN: + self.run_uid = doc.get("uid") + LOGGER.info( + "ISPyB callback recieved start document with experiment parameters." + ) + json_params = doc.get("hyperion_internal_parameters") + self.parameters = RotationInternalParameters.from_json(json_params) LOGGER.info("ISPYB handler received start document.") - if doc.get("subplan_name") == "rotation_scan_main": + if doc.get("subplan_name") == ROTATION_PLAN_MAIN: self.uid_to_finalize_on = doc.get("uid") def event(self, doc: dict): diff --git a/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py b/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py index 3135da361..0ad84c043 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py @@ -30,7 +30,7 @@ def __init__(self): self.writer: NexusWriter | None = None def start(self, doc: dict): - if doc.get("subplan_name") == "rotation_scan_with_cleanup": + if doc.get("subplan_name") == ROTATION_OUTER_PLAN: self.run_uid = doc.get("uid") LOGGER.info( "Nexus writer recieved start document with experiment parameters." diff --git a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index f74d5460d..3cf733f87 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -51,7 +51,7 @@ def fake_rotation_scan( @bpp.set_run_key_decorator("rotation_scan_with_cleanup_and_subs") @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(), } ) @@ -59,10 +59,10 @@ def plan(): if after_open_assert: after_open_assert(subscriptions) - @bpp.set_run_key_decorator("rotation_scan_main") + @bpp.set_run_key_decorator(ROTATION_PLAN_MAIN) @bpp.run_decorator( md={ - "subplan_name": "rotation_scan_main", + "subplan_name": ROTATION_PLAN_MAIN, } ) def fake_main_plan(): @@ -101,7 +101,7 @@ def test_nexus_handler_gets_documents_in_mock_plan( call_content_outer = cb.nexus_handler.start.call_args_list[0].args[0] assert call_content_outer["hyperion_internal_parameters"] == params.json() call_content_inner = cb.nexus_handler.start.call_args_list[1].args[0] - assert call_content_inner["subplan_name"] == "rotation_scan_main" + assert call_content_inner["subplan_name"] == ROTATION_PLAN_MAIN @patch( diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py b/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py index 65b1eb4b1..d254c624e 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/callback_collection.py @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING from hyperion.external_interaction.callbacks.abstract_plan_callback_collection import ( AbstractPlanCallbackCollection, @@ -16,9 +15,6 @@ XrayCentreZocaloCallback, ) -if TYPE_CHECKING: - from hyperion.parameters.internal_parameters import InternalParameters - @dataclass(frozen=True, order=True) class XrayCentreCallbackCollection(AbstractPlanCallbackCollection): @@ -31,10 +27,10 @@ class XrayCentreCallbackCollection(AbstractPlanCallbackCollection): zocalo_handler: XrayCentreZocaloCallback @classmethod - def from_params(cls, parameters: InternalParameters): + def setup(cls): nexus_handler = GridscanNexusFileCallback() - ispyb_handler = GridscanISPyBCallback(parameters) - zocalo_handler = XrayCentreZocaloCallback(parameters, ispyb_handler) + ispyb_handler = GridscanISPyBCallback() + zocalo_handler = XrayCentreZocaloCallback(ispyb_handler) callback_collection = cls( nexus_handler=nexus_handler, ispyb_handler=ispyb_handler, 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 ccebb3f8e..d22f0f49d 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -9,7 +9,7 @@ Store3DGridscanInIspyb, StoreGridscanInIspyb, ) -from hyperion.log import set_dcgid_tag +from hyperion.log import LOGGER, set_dcgid_tag from hyperion.parameters.plan_specific.gridscan_internal_params import ( GridscanInternalParameters, ) @@ -32,15 +32,25 @@ class GridscanISPyBCallback(BaseISPyBCallback): Usually used as part of an FGSCallbackCollection. """ - def __init__(self, parameters: GridscanInternalParameters): - super().__init__(parameters) - self.ispyb: StoreGridscanInIspyb = ( - Store3DGridscanInIspyb(self.ispyb_config, self.params) - if self.params.experiment_params.is_3d_grid_scan - else Store2DGridscanInIspyb(self.ispyb_config, self.params) - ) + def __init__(self) -> None: + super().__init__() + self.ispyb: StoreGridscanInIspyb self.ispyb_ids: tuple = (None, None, None) + def start(self, doc: dict): + if doc.get("subplan_name") == "run_gridscan_move_and_tidy": + LOGGER.info( + "Nexus writer recieved start document with experiment parameters." + ) + json_params = doc.get("hyperion_internal_parameters") + self.params = GridscanInternalParameters.from_json(json_params) + self.ispyb = ( + Store3DGridscanInIspyb(self.ispyb_config, self.params) + if self.params.experiment_params.is_3d_grid_scan + else Store2DGridscanInIspyb(self.ispyb_config, self.params) + ) + self.run_start_uid = doc.get("uid") + def append_to_comment(self, comment: str): for id in self.ispyb_ids[0]: self._append_to_comment(id, comment) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py index d68ee31e9..70ad45a82 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py @@ -4,7 +4,7 @@ from hyperion.external_interaction.nexus.write_nexus import NexusWriter from hyperion.log import LOGGER -from hyperion.parameters.constants import ISPYB_HARDWARE_READ_PLAN +from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN, ISPYB_HARDWARE_READ_PLAN from hyperion.parameters.plan_specific.gridscan_internal_params import ( GridscanInternalParameters, ) @@ -36,7 +36,7 @@ def __init__(self) -> None: self.nexus_writer_2: NexusWriter | None = None def start(self, doc: dict): - if doc.get("subplan_name") == "run_gridscan_move_and_tidy": + if doc.get("subplan_name") == GRIDSCAN_OUTER_PLAN: LOGGER.info( "Nexus writer recieved start document with experiment parameters." ) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py index e6c8811cc..1ff359548 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py @@ -3,6 +3,8 @@ import pytest from hyperion.parameters.constants import ( + GRIDSCAN_AND_MOVE, + GRIDSCAN_MAIN_PLAN, ISPYB_HARDWARE_READ_PLAN, ISPYB_TRANSMISSION_FLUX_READ_PLAN, ) @@ -66,7 +68,7 @@ class TestData: "versions": {"ophyd": "1.6.4.post76+g0895f9f", "bluesky": "1.8.3"}, "scan_id": 1, "plan_type": "generator", - "plan_name": "run_gridscan_and_move", + "plan_name": GRIDSCAN_AND_MOVE, } test_run_gridscan_start_document: dict = { "uid": "d8bee3ee-f614-4e7a-a516-25d6b9e87ef3", @@ -74,8 +76,8 @@ class TestData: "versions": {"ophyd": "1.6.4.post76+g0895f9f", "bluesky": "1.8.3"}, "scan_id": 1, "plan_type": "generator", - "plan_name": "run_gridscan_and_move", - "subplan_name": "run_gridscan", + "plan_name": GRIDSCAN_AND_MOVE, + "subplan_name": GRIDSCAN_MAIN_PLAN, } test_do_fgs_start_document: dict = { "uid": "d8bee3ee-f614-4e7a-a516-25d6b9e87ef3", @@ -83,7 +85,7 @@ class TestData: "versions": {"ophyd": "1.6.4.post76+g0895f9f", "bluesky": "1.8.3"}, "scan_id": 1, "plan_type": "generator", - "plan_name": "run_gridscan_and_move", + "plan_name": GRIDSCAN_AND_MOVE, "subplan_name": "do_fgs", } test_descriptor_document_pre_data_collection: dict = { @@ -137,7 +139,7 @@ class TestData: "exit_status": "success", "reason": "", "num_events": {"fake_ispyb_params": 1, "primary": 1}, - "subplan_name": "run_gridscan", + "subplan_name": GRIDSCAN_MAIN_PLAN, } test_do_fgs_gridscan_stop_document: dict = { "run_start": "d8bee3ee-f614-4e7a-a516-25d6b9e87ef3", @@ -163,5 +165,5 @@ class TestData: "exit_status": "fail", "reason": "could not connect to devices", "num_events": {"fake_ispyb_params": 1, "primary": 1}, - "subplan_name": "run_gridscan", + "subplan_name": GRIDSCAN_MAIN_PLAN, } diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py index 376ee401d..215a2bb3c 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_nexus_handler.py @@ -5,7 +5,11 @@ from hyperion.external_interaction.callbacks.xray_centre.nexus_callback import ( GridscanNexusFileCallback, ) -from hyperion.parameters.constants import ISPYB_HARDWARE_READ_PLAN +from hyperion.parameters.constants import ( + GRIDSCAN_AND_MOVE, + GRIDSCAN_MAIN_PLAN, + ISPYB_HARDWARE_READ_PLAN, +) from hyperion.parameters.external_parameters import from_file as default_raw_params from hyperion.parameters.plan_specific.gridscan_internal_params import ( GridscanInternalParameters, @@ -17,7 +21,7 @@ "versions": {"ophyd": "1.6.4.post76+g0895f9f", "bluesky": "1.8.3"}, "scan_id": 1, "plan_type": "generator", - "plan_name": "run_gridscan_and_move", + "plan_name": GRIDSCAN_AND_MOVE, } @@ -94,7 +98,7 @@ def test_writers_do_create_one_file_each_on_start_doc_for_run_gridscan( ) nexus_handler.start( { - "subplan_name": "run_gridscan", + "subplan_name": GRIDSCAN_MAIN_PLAN, } ) with patch( diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py index 2883bc110..4f30fc7e6 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py @@ -16,6 +16,7 @@ ZocaloInteractor, ) from hyperion.log import LOGGER +from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN from hyperion.parameters.plan_specific.gridscan_internal_params import ( GridscanInternalParameters, ) @@ -43,21 +44,26 @@ class XrayCentreZocaloCallback(CallbackBase): def __init__( self, - parameters: GridscanInternalParameters, ispyb_handler: GridscanISPyBCallback, ): - self.grid_position_to_motor_position: Callable[ - [ndarray], ndarray - ] = parameters.experiment_params.grid_position_to_motor_position self.processing_start_time = 0.0 self.processing_time = 0.0 self.do_fgs_uid: Optional[str] = None self.ispyb: GridscanISPyBCallback = ispyb_handler - self.zocalo_interactor = ZocaloInteractor( - parameters.hyperion_params.zocalo_environment - ) def start(self, doc: dict): + if doc.get("subplan_name") == GRIDSCAN_OUTER_PLAN: + LOGGER.info( + "Zocalo callback recieved start document with experiment parameters." + ) + json_params = doc.get("hyperion_internal_parameters") + params = GridscanInternalParameters.from_json(json_params) + self.zocalo_interactor = ZocaloInteractor( + params.hyperion_params.zocalo_environment + ) + self.grid_position_to_motor_position: Callable[ + [ndarray], ndarray + ] = params.experiment_params.grid_position_to_motor_position LOGGER.info("Zocalo handler received start document.") if doc.get("subplan_name") == "do_fgs": self.do_fgs_uid = doc.get("uid") diff --git a/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py b/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py index ef483f6db..5fb92b9ba 100644 --- a/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py +++ b/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py @@ -51,7 +51,7 @@ def fake_rotation_scan( @bpp.set_run_key_decorator("rotation_scan_with_cleanup_and_subs") @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(), } ) diff --git a/src/hyperion/parameters/constants.py b/src/hyperion/parameters/constants.py index e54553257..919ae54a8 100644 --- a/src/hyperion/parameters/constants.py +++ b/src/hyperion/parameters/constants.py @@ -18,6 +18,17 @@ PARAMETER_SCHEMA_DIRECTORY = "src/hyperion/parameters/schemas/" OAV_REFRESH_DELAY = 0.3 +# Plan section names ################################################################### +# Gridscan +GRIDSCAN_OUTER_PLAN = "run_gridscan_move_and_tidy" +GRIDSCAN_AND_MOVE = "run_gridscan_and_move" +GRIDSCAN_MAIN_PLAN = "run_gridscan" +DO_FGS = "do_fgs" +# Rotation scan +ROTATION_OUTER_PLAN = "rotation_scan_with_cleanup" +ROTATION_PLAN_MAIN = "rotation_scan_main" +######################################################################################## + class Actions(Enum): START = "start" From a27e9a08d5b9df344e81210e6199d2c60e74c212 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 7 Nov 2023 08:34:16 +0000 Subject: [PATCH 02/15] complete instantiation of callbacks on start doc --- .../experiment_plans/rotation_scan_plan.py | 2 +- .../callbacks/ispyb_callback_base.py | 23 ++++++++++++------- .../callbacks/rotation/ispyb_callback.py | 14 ++++------- .../callbacks/rotation/nexus_callback.py | 3 ++- .../callbacks/rotation/zocalo_callback.py | 12 ++++++++-- .../callbacks/xray_centre/ispyb_callback.py | 1 + .../callbacks/xray_centre/zocalo_callback.py | 12 ++++------ .../ispyb/ispyb_dataclass.py | 2 +- 8 files changed, 40 insertions(+), 29 deletions(-) diff --git a/src/hyperion/experiment_plans/rotation_scan_plan.py b/src/hyperion/experiment_plans/rotation_scan_plan.py index 8907d03ce..902c94ca3 100644 --- a/src/hyperion/experiment_plans/rotation_scan_plan.py +++ b/src/hyperion/experiment_plans/rotation_scan_plan.py @@ -253,7 +253,7 @@ 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") diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index 228af7c50..d62b5f6ba 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -12,6 +12,12 @@ ISPYB_TRANSMISSION_FLUX_READ_PLAN, SIM_ISPYB_CONFIG, ) +from hyperion.parameters.plan_specific.gridscan_internal_params import ( + GridscanInternalParameters, +) +from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, +) if TYPE_CHECKING: from hyperion.external_interaction.ispyb.store_in_ispyb import ( @@ -19,12 +25,6 @@ Store3DGridscanInIspyb, StoreRotationInIspyb, ) - from hyperion.parameters.plan_specific.gridscan_internal_params import ( - GridscanInternalParameters, - ) - from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( - RotationInternalParameters, - ) class BaseISPyBCallback(CallbackBase): @@ -32,7 +32,9 @@ def __init__(self) -> None: """Subclasses should run super().__init__(), then in their start() method set self.ispyb to the type of ispyb relevant to the experiment and define the type for self.ispyb_ids.""" - self.params: GridscanInternalParameters | RotationInternalParameters + self.params: GridscanInternalParameters | RotationInternalParameters | None = ( + None + ) self.ispyb: StoreRotationInIspyb | Store3DGridscanInIspyb | Store2DGridscanInIspyb self.descriptors: Dict[str, dict] = {} self.ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) @@ -42,6 +44,9 @@ def __init__(self) -> None: " set the ISPYB_CONFIG_PATH environment variable." ) self.uid_to_finalize_on: Optional[str] = None + self.ispyb_ids: tuple[int, int] | tuple[int, int, int] | tuple[ + None, None, None + ] = (None, None, None) def _append_to_comment(self, id: int, comment: str): assert isinstance(self.ispyb, StoreInIspyb) @@ -66,7 +71,9 @@ def event(self, doc: dict): self.ispyb, StoreInIspyb ), "ISPyB deposition can't be initialised!" event_descriptor = self.descriptors[doc["descriptor"]] - + assert isinstance(self.params, GridscanInternalParameters) or isinstance( + self.params, RotationInternalParameters + ), "ISPyB handler params set with wrong type" if event_descriptor.get("name") == ISPYB_HARDWARE_READ_PLAN: self.params.hyperion_params.ispyb_params.undulator_gap = doc["data"][ "undulator_gap" diff --git a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py index af99df885..0c15a03ef 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py @@ -28,14 +28,6 @@ class RotationISPyBCallback(BaseISPyBCallback): Usually used as part of a RotationCallbackCollection. """ - def __init__(self) -> None: - self.params: RotationInternalParameters - super().__init__() - self.ispyb: StoreRotationInIspyb = StoreRotationInIspyb( - self.ispyb_config, self.params - ) - self.ispyb_ids: tuple[int, int] | tuple[None, None] = (None, None) - def append_to_comment(self, comment: str): assert self.ispyb_ids[0] is not None self._append_to_comment(self.ispyb_ids[0], comment) @@ -47,7 +39,11 @@ def start(self, doc: dict): "ISPyB callback recieved start document with experiment parameters." ) json_params = doc.get("hyperion_internal_parameters") - self.parameters = RotationInternalParameters.from_json(json_params) + self.params = RotationInternalParameters.from_json(json_params) + self.ispyb: StoreRotationInIspyb = StoreRotationInIspyb( + self.ispyb_config, self.params + ) + self.ispyb_ids: tuple[int, int] | tuple[None, None] = (None, None) LOGGER.info("ISPYB handler received start document.") if doc.get("subplan_name") == ROTATION_PLAN_MAIN: self.uid_to_finalize_on = doc.get("uid") diff --git a/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py b/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py index 0ad84c043..54428ce27 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/nexus_callback.py @@ -4,6 +4,7 @@ from hyperion.external_interaction.nexus.write_nexus import NexusWriter from hyperion.log import LOGGER +from hyperion.parameters.constants import ROTATION_OUTER_PLAN from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, ) @@ -24,7 +25,7 @@ class RotationNexusFileCallback(CallbackBase): Usually used as part of a RotationCallbackCollection. """ - def __init__(self): + def __init__(self) -> None: self.run_uid: str | None = None self.parameters: RotationInternalParameters | None = None self.writer: NexusWriter | None = None diff --git a/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py index 3791b62b1..bd7ca92c6 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py @@ -8,6 +8,7 @@ from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade from hyperion.external_interaction.zocalo.zocalo_interaction import ZocaloInteractor from hyperion.log import LOGGER +from hyperion.parameters.constants import ROTATION_OUTER_PLAN class RotationZocaloCallback(CallbackBase): @@ -17,15 +18,22 @@ class RotationZocaloCallback(CallbackBase): def __init__( self, - zocalo_environment: str, ispyb_handler: RotationISPyBCallback, ): self.ispyb: RotationISPyBCallback = ispyb_handler - self.zocalo_interactor = ZocaloInteractor(zocalo_environment) self.run_uid = None def start(self, doc: dict): LOGGER.info("Zocalo handler received start document.") + if doc.get("subplan_name") == ROTATION_OUTER_PLAN: + LOGGER.info( + "Zocalo callback recieved start document with experiment parameters." + ) + assert ( + self.ispyb.params is not None + ), "ISPyB handler attached to Zocalo handler did not recieve parameters" + zocalo_environment = self.ispyb.params.hyperion_params.zocalo_environment + self.zocalo_interactor = ZocaloInteractor(zocalo_environment) if self.run_uid is None: self.run_uid = doc.get("uid") 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 d22f0f49d..8568529cb 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -34,6 +34,7 @@ class GridscanISPyBCallback(BaseISPyBCallback): def __init__(self) -> None: super().__init__() + self.params: GridscanInternalParameters self.ispyb: StoreGridscanInIspyb self.ispyb_ids: tuple = (None, None, None) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py index 4f30fc7e6..9d46e451f 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py @@ -17,9 +17,6 @@ ) from hyperion.log import LOGGER from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN -from hyperion.parameters.plan_specific.gridscan_internal_params import ( - GridscanInternalParameters, -) class XrayCentreZocaloCallback(CallbackBase): @@ -56,14 +53,15 @@ def start(self, doc: dict): LOGGER.info( "Zocalo callback recieved start document with experiment parameters." ) - json_params = doc.get("hyperion_internal_parameters") - params = GridscanInternalParameters.from_json(json_params) + assert ( + self.ispyb.params is not None + ), "ISPyB handler attached to Zocalo handler did not recieve parameters" self.zocalo_interactor = ZocaloInteractor( - params.hyperion_params.zocalo_environment + self.ispyb.params.hyperion_params.zocalo_environment ) self.grid_position_to_motor_position: Callable[ [ndarray], ndarray - ] = params.experiment_params.grid_position_to_motor_position + ] = self.ispyb.params.experiment_params.grid_position_to_motor_position LOGGER.info("Zocalo handler received start document.") if doc.get("subplan_name") == "do_fgs": self.do_fgs_uid = doc.get("uid") diff --git a/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py b/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py index cf69c22b5..06eca874d 100644 --- a/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py @@ -66,7 +66,7 @@ def _parse_position( comment: str resolution: float - sample_id: Optional[int] = None + sample_id: Optional[str] = None sample_barcode: Optional[str] = None # Optional from GDA as populated by Ophyd From 939e1b9402abde5df5748215382f7471a8dca4fb Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 7 Nov 2023 08:34:22 +0000 Subject: [PATCH 03/15] fix tests --- .../experiment_plans/tests/conftest.py | 51 +++++++--- .../tests/test_flyscan_xray_centre_plan.py | 78 +++++++++++----- .../tests/test_rotation_scan_plan.py | 34 +++---- .../callbacks/rotation/callback_collection.py | 14 +-- .../rotation/tests/test_rotation_callbacks.py | 93 ++++++++++++------- .../callbacks/xray_centre/tests/conftest.py | 22 ++++- .../xray_centre/tests/test_ispyb_handler.py | 11 ++- .../test_xraycentre_callback_collection.py | 40 ++++---- .../xray_centre/tests/test_zocalo_handler.py | 17 ++-- .../system_tests/test_write_rotation_nexus.py | 35 +++---- .../system_tests/test_zocalo_system.py | 4 +- .../unit_tests/test_store_in_ispyb.py | 3 + src/hyperion/system_tests/test_fgs_plan.py | 8 +- 13 files changed, 259 insertions(+), 151 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index 0ef519088..bfbb1e60c 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/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,6 +32,8 @@ ) from hyperion.external_interaction.ispyb.store_in_ispyb import 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 ( @@ -45,7 +48,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) @@ -82,7 +85,7 @@ def eiger(): @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 @@ -277,21 +280,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, **kwargs): + mock = MagicMock(spec=Store3DGridscanInIspyb) + mock.begin_deposition.return_value = ((0, 0), 0, 0) + return mock + + @pytest.fixture def mock_subscriptions(test_fgs_params): - subscriptions = XrayCentreCallbackCollection.setup() - 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() + subscriptions.ispyb_handler.start( + { + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + ) + subscriptions.zocalo_handler.start( + { + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + ) subscriptions.ispyb_handler.ispyb = MagicMock(spec=Store3DGridscanInIspyb) subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: [[0, 0], 0, 0] @@ -310,7 +337,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 diff --git a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py index 24e62de7c..2526bc498 100644 --- a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py +++ b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py @@ -1,7 +1,7 @@ import types -from unittest.mock import MagicMock, call, patch +from unittest.mock import ANY, MagicMock, call, patch -import bluesky.plan_stubs as bps +import bluesky.preprocessors as bpp import numpy as np import pytest from bluesky.run_engine import RunEngine @@ -27,6 +27,10 @@ run_gridscan_and_move, wait_for_gridscan_valid, ) +from hyperion.experiment_plans.tests.conftest import ( + modified_interactor_mock, + modified_store_grid_scan_mock, +) from hyperion.external_interaction.callbacks.logging_callback import ( VerbosePlanExecutionLoggingCallback, ) @@ -36,15 +40,18 @@ from hyperion.external_interaction.callbacks.xray_centre.ispyb_callback import ( GridscanISPyBCallback, ) +from hyperion.external_interaction.callbacks.xray_centre.tests.conftest import TestData from hyperion.external_interaction.ispyb.store_in_ispyb import Store3DGridscanInIspyb from hyperion.external_interaction.system_tests.conftest import ( TEST_RESULT_LARGE, TEST_RESULT_MEDIUM, TEST_RESULT_SMALL, ) +from hyperion.external_interaction.zocalo.zocalo_interaction import ZocaloInteractor from hyperion.log import set_up_logging_handlers from hyperion.parameters import external_parameters from hyperion.parameters.constants import ( + GRIDSCAN_OUTER_PLAN, ISPYB_HARDWARE_READ_PLAN, ISPYB_TRANSMISSION_FLUX_READ_PLAN, ) @@ -83,32 +90,37 @@ def test_read_hardware_for_ispyb_updates_from_ophyd_devices( ): undulator_test_value = 1.234 - fake_fgs_composite.undulator.gap.user_readback.sim_put(undulator_test_value) + fake_fgs_composite.undulator.gap.user_readback.sim_put(undulator_test_value) # type: ignore synchrotron_test_value = "test" - fake_fgs_composite.synchrotron.machine_status.synchrotron_mode.sim_put( + fake_fgs_composite.synchrotron.machine_status.synchrotron_mode.sim_put( # type: ignore synchrotron_test_value ) transmission_test_value = 0.01 - fake_fgs_composite.attenuator.actual_transmission.sim_put(transmission_test_value) + fake_fgs_composite.attenuator.actual_transmission.sim_put(transmission_test_value) # type: ignore xgap_test_value = 0.1234 ygap_test_value = 0.2345 - fake_fgs_composite.s4_slit_gaps.xgap.user_readback.sim_put(xgap_test_value) - fake_fgs_composite.s4_slit_gaps.ygap.user_readback.sim_put(ygap_test_value) + fake_fgs_composite.s4_slit_gaps.xgap.user_readback.sim_put(xgap_test_value) # type: ignore + fake_fgs_composite.s4_slit_gaps.ygap.user_readback.sim_put(ygap_test_value) # type: ignore flux_test_value = 10.0 - fake_fgs_composite.flux.flux_reading.sim_put(flux_test_value) + fake_fgs_composite.flux.flux_reading.sim_put(flux_test_value) # type: ignore - test_ispyb_callback = GridscanISPyBCallback(test_fgs_params) + test_ispyb_callback = GridscanISPyBCallback() test_ispyb_callback.ispyb = MagicMock(spec=Store3DGridscanInIspyb) RE.subscribe(test_ispyb_callback) + @bpp.set_run_key_decorator(GRIDSCAN_OUTER_PLAN) + @bpp.run_decorator( # attach experiment metadata to the start document + md={ + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + ) def standalone_read_hardware_for_ispyb(und, syn, slits, attn, fl): - yield from bps.open_run() yield from read_hardware_for_ispyb_pre_collection(und, syn, slits) yield from read_hardware_for_ispyb_during_collection(attn, fl) - yield from bps.close_run() RE( standalone_read_hardware_for_ispyb( @@ -151,6 +163,12 @@ def test_results_adjusted_and_passed_to_move_xyz( set_up_logging_handlers(logging_level="INFO", dev_mode=True) RE.subscribe(VerbosePlanExecutionLoggingCallback()) + mock_subscriptions.ispyb_handler.start( + { + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + ) mock_subscriptions.ispyb_handler.descriptor( {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} ) @@ -208,7 +226,7 @@ def test_results_adjusted_and_passed_to_move_xyz( mock_subscriptions, ) ) - + assert fake_fgs_composite.aperture_scatterguard.aperture_positions is not None ap_call_large = call( *(fake_fgs_composite.aperture_scatterguard.aperture_positions.LARGE) ) @@ -274,6 +292,10 @@ def test_results_passed_to_move_motors( @patch("hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True) @patch("hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True) @patch("bluesky.plan_stubs.rd") +@patch( + "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor", + modified_interactor_mock, +) def test_individual_plans_triggered_once_and_only_once_in_composite_run( rd: MagicMock, move_xyz: MagicMock, @@ -284,6 +306,9 @@ def test_individual_plans_triggered_once_and_only_once_in_composite_run( test_fgs_params: GridscanInternalParameters, RE: RunEngine, ): + td = TestData() + mock_subscriptions.ispyb_handler.start(td.test_start_document) + mock_subscriptions.zocalo_handler.start(td.test_start_document) mock_subscriptions.ispyb_handler.descriptor( {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} ) @@ -346,6 +371,7 @@ def test_logging_within_plan( test_fgs_params: GridscanInternalParameters, RE: RunEngine, ): + mock_subscriptions.ispyb_handler.start({}) mock_subscriptions.ispyb_handler.descriptor( {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} ) @@ -397,8 +423,8 @@ def test_GIVEN_scan_already_valid_THEN_wait_for_GRIDSCAN_returns_immediately( ): test_fgs: FastGridScan = make_fake_device(FastGridScan)("prefix", name="fake_fgs") - test_fgs.scan_invalid.sim_put(False) - test_fgs.position_counter.sim_put(0) + test_fgs.scan_invalid.sim_put(False) # type: ignore + test_fgs.position_counter.sim_put(0) # type: ignore RE(wait_for_gridscan_valid(test_fgs)) @@ -411,8 +437,8 @@ def test_GIVEN_scan_not_valid_THEN_wait_for_GRIDSCAN_raises_and_sleeps_called( ): test_fgs: FastGridScan = make_fake_device(FastGridScan)("prefix", name="fake_fgs") - test_fgs.scan_invalid.sim_put(True) - test_fgs.position_counter.sim_put(0) + test_fgs.scan_invalid.sim_put(True) # type: ignore + test_fgs.position_counter.sim_put(0) # type: ignore with pytest.raises(WarningException): RE(wait_for_gridscan_valid(test_fgs)) @@ -448,22 +474,26 @@ def test_when_grid_scan_ran_then_eiger_disarmed_before_zocalo_end( mock_parent = MagicMock() fake_fgs_composite.eiger.disarm_detector = mock_parent.disarm - fake_fgs_composite.eiger.filewriters_finished = Status() - fake_fgs_composite.eiger.filewriters_finished.set_finished() + fake_fgs_composite.eiger.filewriters_finished = Status(done=True, success=True) fake_fgs_composite.eiger.odin.check_odin_state = MagicMock(return_value=True) - fake_fgs_composite.eiger.odin.file_writer.num_captured.sim_put(1200) + fake_fgs_composite.eiger.odin.file_writer.num_captured.sim_put(1200) # type: ignore fake_fgs_composite.eiger.stage = MagicMock( return_value=Status(None, None, 0, True, True) ) - fake_fgs_composite.xbpm_feedback.pos_stable.sim_put(1) + fake_fgs_composite.xbpm_feedback.pos_stable.sim_put(1) # type: ignore - mock_subscriptions.zocalo_handler.zocalo_interactor.run_end = mock_parent.run_end with patch( - "hyperion.experiment_plans.flyscan_xray_centre_plan.XrayCentreCallbackCollection.from_params", - lambda _: mock_subscriptions, + "hyperion.experiment_plans.flyscan_xray_centre_plan.XrayCentreCallbackCollection.setup", + lambda: mock_subscriptions, ), patch( "hyperion.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter.create_nexus_file", autospec=True, + ), patch( + "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor", + lambda _: modified_interactor_mock(mock_parent.run_end), + ), patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + modified_store_grid_scan_mock, ): RE(flyscan_xray_centre(fake_fgs_composite, test_fgs_params)) @@ -511,7 +541,7 @@ class CompleteException(Exception): fake_fgs_composite.eiger.disable_roi_mode = MagicMock() # Without the complete finishing we will not get all the images - fake_fgs_composite.eiger.ALL_FRAMES_TIMEOUT = 0.1 + fake_fgs_composite.eiger.ALL_FRAMES_TIMEOUT = 0.1 # type: ignore # Want to get the underlying completion error, not the one raised from unstage with pytest.raises(CompleteException): diff --git a/src/hyperion/experiment_plans/tests/test_rotation_scan_plan.py b/src/hyperion/experiment_plans/tests/test_rotation_scan_plan.py index 97ac52dcb..bf8624f29 100644 --- a/src/hyperion/experiment_plans/tests/test_rotation_scan_plan.py +++ b/src/hyperion/experiment_plans/tests/test_rotation_scan_plan.py @@ -81,8 +81,8 @@ def do_rotation_main_plan_for_tests( fake_read, ), patch( - "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", - lambda _: callbacks, + "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.setup", + lambda: callbacks, ), patch("dodal.beamlines.i03.undulator", lambda: sim_und), patch("dodal.beamlines.i03.synchrotron", lambda: sim_synch), @@ -116,8 +116,8 @@ def run_full_rotation_plan( fake_read, ), patch( - "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", - lambda _: mock_rotation_subscriptions, + "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.setup", + lambda: mock_rotation_subscriptions, ), ): RE(rotation_scan(fake_create_rotation_devices, test_rotation_params)) @@ -339,8 +339,8 @@ def test_rotation_scan( return_value=detector_motion, ), patch( - "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", - lambda _: mock_rotation_subscriptions, + "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.setup", + lambda: mock_rotation_subscriptions, ), ): composite = RotationScanComposite( @@ -361,12 +361,14 @@ def test_rotation_scan( eiger.unstage.assert_called() -def test_rotation_plan_runs(setup_and_run_rotation_plan_for_tests_standard): +def test_rotation_plan_runs(setup_and_run_rotation_plan_for_tests_standard) -> None: RE: RunEngine = setup_and_run_rotation_plan_for_tests_standard["RE"] assert RE._exit_status == "success" -def test_rotation_plan_zebra_settings(setup_and_run_rotation_plan_for_tests_standard): +def test_rotation_plan_zebra_settings( + setup_and_run_rotation_plan_for_tests_standard, +) -> None: zebra: Zebra = setup_and_run_rotation_plan_for_tests_standard["zebra"] params: RotationInternalParameters = setup_and_run_rotation_plan_for_tests_standard[ "test_rotation_params" @@ -381,12 +383,12 @@ def test_rotation_plan_zebra_settings(setup_and_run_rotation_plan_for_tests_stan def test_full_rotation_plan_smargon_settings( run_full_rotation_plan, test_rotation_params, -): +) -> None: smargon: Smargon = run_full_rotation_plan.smargon params: RotationInternalParameters = test_rotation_params expt_params = params.experiment_params - omega_set: MagicMock = smargon.omega.set + omega_set: MagicMock = smargon.omega.set # type: ignore rotation_speed = ( expt_params.image_width / params.hyperion_params.detector_params.exposure_time ) @@ -405,7 +407,7 @@ def test_full_rotation_plan_smargon_settings( def test_rotation_plan_smargon_doesnt_move_xyz_if_not_given_in_params( setup_and_run_rotation_plan_for_tests_nomove, -): +) -> None: smargon: Smargon = setup_and_run_rotation_plan_for_tests_nomove["smargon"] params: RotationInternalParameters = setup_and_run_rotation_plan_for_tests_nomove[ "test_rotation_params" @@ -477,8 +479,8 @@ class MyTestException(Exception): cleanup_plan.assert_not_called() # check that failure is handled in composite plan with patch( - "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", - lambda _: mock_rotation_subscriptions, + "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.setup", + lambda: mock_rotation_subscriptions, ): with pytest.raises(MyTestException) as exc: RE( @@ -525,7 +527,7 @@ def test_ispyb_deposition_in_plan( test_rotation_params.hyperion_params.ispyb_params.current_energy_ev = ( convert_angstrom_to_eV(test_wl) ) - callbacks = RotationCallbackCollection.from_params(test_rotation_params) + callbacks = RotationCallbackCollection.setup() callbacks.ispyb_handler.ispyb.ISPYB_CONFIG_PATH = DEV_ISPYB_DATABASE_CFG composite = RotationScanComposite( @@ -547,8 +549,8 @@ def test_ispyb_deposition_in_plan( fake_read, ), patch( - "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.from_params", - lambda _: callbacks, + "hyperion.experiment_plans.rotation_scan_plan.RotationCallbackCollection.setup", + lambda: callbacks, ), ): RE( diff --git a/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py b/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py index d23be80de..dd6c80a6a 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/rotation/callback_collection.py @@ -1,7 +1,6 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING from hyperion.external_interaction.callbacks.abstract_plan_callback_collection import ( AbstractPlanCallbackCollection, @@ -16,11 +15,6 @@ RotationZocaloCallback, ) -if TYPE_CHECKING: - from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( - RotationInternalParameters, - ) - @dataclass(frozen=True, order=True) class RotationCallbackCollection(AbstractPlanCallbackCollection): @@ -32,12 +26,10 @@ class RotationCallbackCollection(AbstractPlanCallbackCollection): zocalo_handler: RotationZocaloCallback @classmethod - def from_params(cls, parameters: RotationInternalParameters): + def setup(cls): nexus_handler = RotationNexusFileCallback() - ispyb_handler = RotationISPyBCallback(parameters) - zocalo_handler = RotationZocaloCallback( - parameters.hyperion_params.zocalo_environment, ispyb_handler - ) + ispyb_handler = RotationISPyBCallback() + zocalo_handler = RotationZocaloCallback(ispyb_handler) callback_collection = cls( nexus_handler=nexus_handler, ispyb_handler=ispyb_handler, diff --git a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index 3cf733f87..81720dc87 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -16,8 +16,13 @@ from hyperion.external_interaction.callbacks.rotation.callback_collection import ( RotationCallbackCollection, ) +from hyperion.external_interaction.callbacks.xray_centre.tests.conftest import TestData from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade -from hyperion.external_interaction.ispyb.store_in_ispyb import StoreInIspyb +from hyperion.external_interaction.ispyb.store_in_ispyb import ( + StoreInIspyb, + StoreRotationInIspyb, +) +from hyperion.parameters.constants import ROTATION_OUTER_PLAN, ROTATION_PLAN_MAIN from hyperion.parameters.external_parameters import from_file from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, @@ -33,13 +38,21 @@ def params(): ) +@pytest.fixture +def test_start_doc(params: RotationInternalParameters): + return { + "subplan_name": ROTATION_OUTER_PLAN, + "hyperion_internal_parameters": params.json(), + } + + @pytest.fixture def RE(): return RunEngine({}) def fake_rotation_scan( - parameters: RotationInternalParameters, + params: RotationInternalParameters, subscriptions: RotationCallbackCollection, after_open_assert: Callable | None = None, after_main_assert: Callable | None = None, @@ -52,7 +65,7 @@ def fake_rotation_scan( @bpp.run_decorator( # attach experiment metadata to the start document md={ "subplan_name": ROTATION_OUTER_PLAN, - "hyperion_internal_parameters": parameters.json(), + "hyperion_internal_parameters": params.json(), } ) def plan(): @@ -67,6 +80,7 @@ def plan(): ) def fake_main_plan(): yield from read_hardware_for_ispyb_during_collection(attenuator, flux) + subscriptions.ispyb_handler.ispyb_ids = (0, 0) if after_main_assert: after_main_assert(subscriptions) yield from bps.sleep(0) @@ -86,11 +100,12 @@ def test_nexus_handler_gets_documents_in_mock_plan( with patch( "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationZocaloCallback", autospec=True, + ), patch( + "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationISPyBCallback", + autospec=True, ): - cb = RotationCallbackCollection.from_params(params) + cb = RotationCallbackCollection.setup() cb.nexus_handler.start = MagicMock(autospec=True) - cb.ispyb_handler.start = MagicMock(autospec=True) - cb.ispyb_handler.stop = MagicMock(autospec=True) RE(fake_rotation_scan(params, cb)) @@ -113,18 +128,22 @@ def test_nexus_handler_gets_documents_in_mock_plan( autospec=True, ) def test_nexus_handler_only_writes_once( - ispyb, nexus_writer, RE: RunEngine, params: RotationInternalParameters + ispyb, + nexus_writer, + RE: RunEngine, + params: RotationInternalParameters, + test_start_doc, ): with patch( "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationZocaloCallback", autospec=True, ): - cb = RotationCallbackCollection.from_params(params) - cb.ispyb_handler.start = MagicMock(autospec=True) - cb.ispyb_handler.stop = MagicMock(autospec=True) + cb = RotationCallbackCollection.setup() + cb.ispyb_handler.start(test_start_doc) RE(fake_rotation_scan(params, cb)) nexus_writer.assert_called_once() + assert cb.nexus_handler.writer is not None cb.nexus_handler.writer.create_nexus_file.assert_called_once() @@ -133,9 +152,7 @@ def test_nexus_handler_only_writes_once( autospec=True, ) def test_nexus_handler_triggers_write_file_when_told( - ispyb, - RE: RunEngine, - params: RotationInternalParameters, + ispyb, RE: RunEngine, params: RotationInternalParameters ): if os.path.isfile("/tmp/file_name_0.nxs"): os.remove("/tmp/file_name_0.nxs") @@ -146,10 +163,12 @@ def test_nexus_handler_triggers_write_file_when_told( "hyperion.external_interaction.callbacks.rotation.callback_collection.RotationZocaloCallback", autospec=True, ): - cb = RotationCallbackCollection.from_params(params) + cb = RotationCallbackCollection.setup() cb.ispyb_handler.start = MagicMock(autospec=True) cb.ispyb_handler.stop = MagicMock(autospec=True) + cb.ispyb_handler.ispyb = ispyb + cb.ispyb_handler.params = params RE(fake_rotation_scan(params, cb)) @@ -173,11 +192,13 @@ def test_zocalo_start_and_end_triggered_once( RE: RunEngine, params: RotationInternalParameters, ): - cb = RotationCallbackCollection.from_params(params) + cb = RotationCallbackCollection.setup() cb.nexus_handler.start = MagicMock(autospec=True) cb.ispyb_handler.start = MagicMock(autospec=True) cb.ispyb_handler.stop = MagicMock(autospec=True) + cb.ispyb_handler.ispyb = MagicMock(spec=StoreRotationInIspyb) + cb.ispyb_handler.params = params cb.ispyb_handler.ispyb_ids = (0, 0) RE(fake_rotation_scan(params, cb)) @@ -192,32 +213,40 @@ def test_zocalo_start_and_end_triggered_once( autospec=True, ) def test_zocalo_start_and_end_not_triggered_if_ispyb_ids_not_present( - zocalo, - RE: RunEngine, - params: RotationInternalParameters, + zocalo, RE: RunEngine, params: RotationInternalParameters, test_start_doc ): - cb = RotationCallbackCollection.from_params(params) + cb = RotationCallbackCollection.setup() cb.nexus_handler.start = MagicMock(autospec=True) cb.ispyb_handler.start = MagicMock(autospec=True) cb.ispyb_handler.stop = MagicMock(autospec=True) cb.ispyb_handler.event = MagicMock(autospec=True) + cb.ispyb_handler.ispyb = MagicMock(spec=StoreRotationInIspyb) + cb.ispyb_handler.params = params cb.ispyb_handler.ispyb = MagicMock(autospec=True) with pytest.raises(ISPyBDepositionNotMade): RE(fake_rotation_scan(params, cb)) +@patch( + "hyperion.external_interaction.callbacks.rotation.nexus_callback.NexusWriter", + autospec=True, +) @patch( "hyperion.external_interaction.callbacks.rotation.zocalo_callback.ZocaloInteractor", autospec=True, ) def test_zocalo_starts_on_opening_and_ispyb_on_main_so_ispyb_triggered_before_zocalo( zocalo, + nexus_writer, RE: RunEngine, params: RotationInternalParameters, + test_start_doc, ): - cb = RotationCallbackCollection.from_params(params) - cb.nexus_handler.start = MagicMock(autospec=True) + cb = RotationCallbackCollection.setup() + cb.nexus_handler.start(test_start_doc) + cb.ispyb_handler.start(test_start_doc) + cb.zocalo_handler.start(test_start_doc) cb.ispyb_handler.ispyb = MagicMock(spec=StoreInIspyb) cb.ispyb_handler.ispyb_ids = (0, 0) cb.zocalo_handler.zocalo_interactor.run_start = MagicMock() @@ -230,7 +259,11 @@ def after_main_assert(callbacks: RotationCallbackCollection): callbacks.ispyb_handler.ispyb.begin_deposition.assert_called_once() cb.zocalo_handler.zocalo_interactor.run_end.assert_not_called() - RE(fake_rotation_scan(params, cb, after_open_assert, after_main_assert)) + with patch( + "hyperion.external_interaction.callbacks.rotation.ispyb_callback.StoreRotationInIspyb", + autospec=True, + ): + RE(fake_rotation_scan(params, cb, after_open_assert, after_main_assert)) cb.zocalo_handler.zocalo_interactor.run_end.assert_called_once() @@ -240,19 +273,13 @@ def after_main_assert(callbacks: RotationCallbackCollection): autospec=True, ) def test_ispyb_handler_grabs_uid_from_main_plan_and_not_first_start_doc( - zocalo, - RE: RunEngine, - params: RotationInternalParameters, + zocalo, RE: RunEngine, params: RotationInternalParameters, test_start_doc ): - cb = RotationCallbackCollection.from_params(params) + cb = RotationCallbackCollection.setup() cb.nexus_handler.start = MagicMock(autospec=True) cb.ispyb_handler.start = MagicMock( autospec=True, side_effect=cb.ispyb_handler.start ) - cb.ispyb_handler.ispyb = MagicMock(spec=StoreInIspyb) - cb.ispyb_handler.ispyb_ids = (0, 0) - cb.zocalo_handler.zocalo_interactor.run_start = MagicMock() - cb.zocalo_handler.zocalo_interactor.run_end = MagicMock() def after_open_assert(callbacks: RotationCallbackCollection): callbacks.ispyb_handler.start.assert_called_once() @@ -262,4 +289,8 @@ def after_main_assert(callbacks: RotationCallbackCollection): assert callbacks.ispyb_handler.start.call_count == 2 assert callbacks.ispyb_handler.uid_to_finalize_on is not None - RE(fake_rotation_scan(params, cb, after_open_assert, after_main_assert)) + with patch( + "hyperion.external_interaction.callbacks.rotation.ispyb_callback.StoreRotationInIspyb", + autospec=True, + ): + RE(fake_rotation_scan(params, cb, after_open_assert, after_main_assert)) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py index 1ff359548..9ebf1eaa7 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/conftest.py @@ -2,12 +2,20 @@ import pytest +from hyperion.external_interaction.callbacks.xray_centre.ispyb_callback import ( + GridscanISPyBCallback, +) from hyperion.parameters.constants import ( GRIDSCAN_AND_MOVE, GRIDSCAN_MAIN_PLAN, + GRIDSCAN_OUTER_PLAN, ISPYB_HARDWARE_READ_PLAN, ISPYB_TRANSMISSION_FLUX_READ_PLAN, ) +from hyperion.parameters.external_parameters import from_file as default_raw_params +from hyperion.parameters.plan_specific.gridscan_internal_params import ( + GridscanInternalParameters, +) @pytest.fixture @@ -58,6 +66,16 @@ def mock_ispyb_end_deposition(): yield p +@pytest.fixture +def ispyb_handler(): + return GridscanISPyBCallback() + + +def dummy_params(): + dummy_params = GridscanInternalParameters(**default_raw_params()) + return dummy_params + + class TestData: DUMMY_TIME_STRING: str = "1970-01-01 00:00:00" GOOD_ISPYB_RUN_STATUS: str = "DataCollection Successful" @@ -68,7 +86,9 @@ class TestData: "versions": {"ophyd": "1.6.4.post76+g0895f9f", "bluesky": "1.8.3"}, "scan_id": 1, "plan_type": "generator", - "plan_name": GRIDSCAN_AND_MOVE, + "plan_name": GRIDSCAN_OUTER_PLAN, + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": dummy_params().json(), } test_run_gridscan_start_document: dict = { "uid": "d8bee3ee-f614-4e7a-a516-25d6b9e87ef3", diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py index 8d442d319..fa348fc2e 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py @@ -29,11 +29,12 @@ def test_fgs_failing_results_in_bad_run_status_in_ispyb( mock_ispyb_get_time: MagicMock, mock_ispyb_store_grid_scan: MagicMock, dummy_params, + ispyb_handler, ): mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING mock_ispyb_update_time_and_status.return_value = None - ispyb_handler = GridscanISPyBCallback(dummy_params) + ispyb_handler.start(td.test_start_document) ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) ispyb_handler.event(td.test_event_document_pre_data_collection) @@ -61,11 +62,12 @@ def test_fgs_raising_no_exception_results_in_good_run_status_in_ispyb( mock_ispyb_get_time: MagicMock, mock_ispyb_store_grid_scan: MagicMock, dummy_params, + ispyb_handler, ): mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING mock_ispyb_update_time_and_status.return_value = None - ispyb_handler = GridscanISPyBCallback(dummy_params) + ispyb_handler = GridscanISPyBCallback() ispyb_handler.start(td.test_start_document) ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) ispyb_handler.event(td.test_event_document_pre_data_collection) @@ -107,7 +109,7 @@ def test_given_ispyb_callback_started_writing_to_ispyb_when_messages_logged_then mock_emit, mock_ispyb_store_grid_scan: MagicMock, dummy_params ): mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] - ispyb_handler = GridscanISPyBCallback(dummy_params) + ispyb_handler = GridscanISPyBCallback() ispyb_handler.start(td.test_start_document) ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) ispyb_handler.event(td.test_event_document_pre_data_collection) @@ -126,11 +128,12 @@ def test_given_ispyb_callback_finished_writing_to_ispyb_when_messages_logged_the mock_ispyb_update_time_and_status: MagicMock, mock_ispyb_get_time: MagicMock, dummy_params, + ispyb_handler, ): mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING mock_ispyb_update_time_and_status.return_value = None - ispyb_handler = GridscanISPyBCallback(dummy_params) + ispyb_handler = GridscanISPyBCallback() ispyb_handler.start(td.test_start_document) ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) ispyb_handler.event(td.test_event_document_pre_data_collection) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py index 7ef2d6eb6..ab69fc133 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py @@ -1,8 +1,8 @@ +from typing import Generator from unittest.mock import MagicMock import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp -import numpy as np import pytest from bluesky.run_engine import RunEngine from dodal.devices.eiger import DetectorParams, EigerDetector @@ -23,7 +23,7 @@ def test_callback_collection_init(): test_parameters = GridscanInternalParameters(**default_raw_params()) - callbacks = XrayCentreCallbackCollection.from_params(test_parameters) + callbacks = XrayCentreCallbackCollection.setup() assert ( callbacks.ispyb_handler.params.experiment_params == test_parameters.experiment_params @@ -46,7 +46,7 @@ def test_callback_collection_init(): @pytest.fixture() -def eiger(): +def eiger() -> Generator[EigerDetector, None, None]: detector_params: DetectorParams = DetectorParams( current_energy_ev=100, exposure_time=0.1, @@ -55,7 +55,8 @@ def eiger(): detector_distance=100.0, omega_start=0.0, omega_increment=0.1, - num_images=50, + num_images_per_trigger=50, + num_triggers=1, use_roi_mode=False, run_number=0, det_dist_to_beam_converter_path="src/hyperion/unit_tests/test_lookup_table.txt", @@ -68,7 +69,7 @@ def eiger(): eiger.cam.manual_trigger.put("Yes") # S03 currently does not have StaleParameters_RBV - eiger.wait_for_stale_parameters = lambda: None + # eiger.wait_for_stale_parameters = lambda: None eiger.odin.check_odin_initialised = lambda: (True, "") yield eiger @@ -82,7 +83,6 @@ def test_communicator_in_composite_run( nexus_writer: MagicMock, ispyb_begin_deposition: MagicMock, ispyb_end_deposition: MagicMock, - eiger: EigerDetector, ): nexus_writer.side_effect = [MagicMock(), MagicMock()] RE = RunEngine({}) @@ -91,34 +91,33 @@ def test_communicator_in_composite_run( params.hyperion_params.beamline = SIM_BEAMLINE ispyb_begin_deposition.return_value = ([1, 2], None, 4) - callbacks = XrayCentreCallbackCollection.from_params(params) - callbacks.zocalo_handler._wait_for_result = MagicMock() - callbacks.zocalo_handler._run_end = MagicMock() - callbacks.zocalo_handler._run_start = MagicMock() - callbacks.zocalo_handler.xray_centre_motor_position = np.array([1, 2, 3]) + callbacks = XrayCentreCallbackCollection.setup() + callbacks.zocalo_handler.zocalo_interactor.wait_for_result = MagicMock() + callbacks.zocalo_handler.zocalo_interactor.run_start = MagicMock() + callbacks.zocalo_handler.zocalo_interactor.run_end = MagicMock() + # callbacks.zocalo_handler.xray_centre_motor_position = np.array([1, 2, 3]) flyscan_xray_centre_composite = MagicMock(spec=FlyScanXRayCentreComposite) # this is where it's currently getting stuck: # flyscan_xray_centre_composite.fast_grid_scan.is_invalid = lambda: False # but this is not a solution # Would be better to use flyscan_xray_centre instead but eiger doesn't work well in S03 - RE(run_gridscan_and_move(flyscan_xray_centre_composite, eiger, params, callbacks)) + RE(run_gridscan_and_move(flyscan_xray_centre_composite, params, callbacks)) # nexus writing - callbacks.nexus_handler.nexus_writer_1.assert_called_once() - callbacks.nexus_handler.nexus_writer_2.assert_called_once() + callbacks.nexus_handler.nexus_writer_1.assert_called_once() # type: ignore + callbacks.nexus_handler.nexus_writer_2.assert_called_once() # type: ignore # ispyb ispyb_begin_deposition.assert_called_once() ispyb_end_deposition.assert_called_once() # zocalo - callbacks.zocalo_handler._run_start.assert_called() - callbacks.zocalo_handler._run_end.assert_called() - callbacks.zocalo_handler._wait_for_result.assert_called_once() + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once() + callbacks.zocalo_handler.zocalo_interactor.run_start.assert_called_once() + callbacks.zocalo_handler.zocalo_interactor.run_end.assert_called_once() def test_callback_collection_list(): - test_parameters = GridscanInternalParameters(**default_raw_params()) - callbacks = XrayCentreCallbackCollection.from_params(test_parameters) + callbacks = XrayCentreCallbackCollection.setup() callback_list = list(callbacks) assert len(callback_list) == 3 assert callbacks.ispyb_handler in callback_list @@ -127,8 +126,7 @@ def test_callback_collection_list(): def test_subscribe_in_plan(): - test_parameters = GridscanInternalParameters(**default_raw_params()) - callbacks = XrayCentreCallbackCollection.from_params(test_parameters) + callbacks = XrayCentreCallbackCollection.setup() document_event_mock = MagicMock() callbacks.ispyb_handler.start = document_event_mock callbacks.ispyb_handler.stop = document_event_mock diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py index 259838e04..4d71d72e3 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py @@ -31,6 +31,8 @@ def dummy_params(): def mock_zocalo_functions(callbacks: XrayCentreCallbackCollection): + callbacks.ispyb_handler.start(td.test_start_document) + callbacks.zocalo_handler.start(td.test_start_document) callbacks.zocalo_handler.zocalo_interactor.wait_for_result = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_end = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_start = MagicMock() @@ -50,10 +52,10 @@ def test_execution_of_run_gridscan_triggers_zocalo_calls( mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING mock_ispyb_update_time_and_status.return_value = None - callbacks = XrayCentreCallbackCollection.from_params(dummy_params) + callbacks = XrayCentreCallbackCollection.setup() mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) + callbacks.zocalo_handler.start(td.test_run_gridscan_start_document) callbacks.ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) callbacks.ispyb_handler.descriptor( @@ -87,8 +89,7 @@ def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_cal store_3d_grid_scan, dummy_params: GridscanInternalParameters, ): - callbacks = XrayCentreCallbackCollection.from_params(dummy_params) - callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) + callbacks = XrayCentreCallbackCollection.setup() callbacks.ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) @@ -134,7 +135,7 @@ def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_ store_3d_grid_scan, dummy_params, ): - callbacks = XrayCentreCallbackCollection.from_params(dummy_params) + callbacks = XrayCentreCallbackCollection.setup() mock_zocalo_functions(callbacks) callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) callbacks.zocalo_handler.zocalo_interactor.wait_for_result.side_effect = ( @@ -158,7 +159,7 @@ def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_excepti store_3d_grid_scan, dummy_params, ): - callbacks = XrayCentreCallbackCollection.from_params(dummy_params) + callbacks = XrayCentreCallbackCollection.setup() mock_zocalo_functions(callbacks) with pytest.raises(ISPyBDepositionNotMade): @@ -173,7 +174,7 @@ def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_b store_3d_grid_scan, dummy_params: GridscanInternalParameters, ): - callbacks = XrayCentreCallbackCollection.from_params(dummy_params) + callbacks = XrayCentreCallbackCollection.setup() mock_zocalo_functions(callbacks) callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) expected_centre_grid_coords = np.array([4, 6, 2]) @@ -214,6 +215,6 @@ def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_b ) ) np.testing.assert_array_equal(found_centre, expected_centre_motor_coords) - + assert isinstance(found_bbox, np.ndarray) expected_bbox_size = np.array([8, 8, 7]) - np.array([2, 2, 2]) np.testing.assert_array_equal(found_bbox, expected_bbox_size) diff --git a/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py b/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py index 5fb92b9ba..1becc793f 100644 --- a/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py +++ b/src/hyperion/external_interaction/system_tests/test_write_rotation_nexus.py @@ -12,6 +12,7 @@ from hyperion.external_interaction.callbacks.rotation.callback_collection import ( RotationCallbackCollection, ) +from hyperion.parameters.constants import ROTATION_OUTER_PLAN from hyperion.parameters.external_parameters import from_file from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( RotationInternalParameters, @@ -80,7 +81,7 @@ def test_rotation_scan_nexus_output_compared_to_existing_file( RE = RunEngine({}) - cb = RotationCallbackCollection.from_params(test_params) + cb = RotationCallbackCollection.setup() cb.ispyb_handler.start = MagicMock() cb.ispyb_handler.stop = MagicMock() cb.ispyb_handler.event = MagicMock() @@ -97,18 +98,18 @@ def test_rotation_scan_nexus_output_compared_to_existing_file( h5py.File(str(TEST_DIRECTORY / TEST_EXAMPLE_NEXUS_FILE), "r") as example_nexus, h5py.File(nexus_filename, "r") as hyperion_nexus, ): - assert hyperion_nexus["/entry/start_time"][()] == b"test_timeZ" - assert hyperion_nexus["/entry/end_time_estimated"][()] == b"test_timeZ" + assert hyperion_nexus["/entry/start_time"][()] == b"test_timeZ" # type: ignore + assert hyperion_nexus["/entry/end_time_estimated"][()] == b"test_timeZ" # type: ignore # we used to write the positions wrong... hyperion_omega: np.ndarray = np.array( - hyperion_nexus["/entry/data/omega"][:] + hyperion_nexus["/entry/data/omega"][:] # type: ignore ) * (3599 / 3600) - example_omega: np.ndarray = example_nexus["/entry/data/omega"][:] + example_omega: np.ndarray = example_nexus["/entry/data/omega"][:] # type: ignore assert np.allclose(hyperion_omega, example_omega) - hyperion_data_shape = hyperion_nexus["/entry/data/data"].shape - example_data_shape = example_nexus["/entry/data/data"].shape + hyperion_data_shape = hyperion_nexus["/entry/data/data"].shape # type: ignore + example_data_shape = example_nexus["/entry/data/data"].shape # type: ignore assert hyperion_data_shape == example_data_shape @@ -117,12 +118,12 @@ def test_rotation_scan_nexus_output_compared_to_existing_file( transmission = "attenuator/attenuator_transmission" wavelength = "beam/incident_wavelength" assert np.isclose( - hyperion_instrument[transmission][()], - example_instrument[transmission][()], + hyperion_instrument[transmission][()], # type: ignore + example_instrument[transmission][()], # type: ignore ) assert np.isclose( - hyperion_instrument[wavelength][()], - example_instrument[wavelength][()], + hyperion_instrument[wavelength][()], # type: ignore + example_instrument[wavelength][()], # type: ignore ) hyperion_sam_x = hyperion_nexus["/entry/sample/sample_x/sam_x"] @@ -141,16 +142,16 @@ def test_rotation_scan_nexus_output_compared_to_existing_file( example_sam_omega = example_nexus["/entry/sample/sample_omega/omega"] assert np.isclose( - hyperion_sam_x[()], - example_sam_x[()], + hyperion_sam_x[()], # type: ignore + example_sam_x[()], # type: ignore ) assert np.isclose( - hyperion_sam_y[()], - example_sam_y[()], + hyperion_sam_y[()], # type: ignore + example_sam_y[()], # type: ignore ) assert np.isclose( - hyperion_sam_z[()], - example_sam_z[()], + hyperion_sam_z[()], # type: ignore + example_sam_z[()], # type: ignore ) assert hyperion_sam_x.attrs.get("depends_on") == example_sam_x.attrs.get( diff --git a/src/hyperion/external_interaction/system_tests/test_zocalo_system.py b/src/hyperion/external_interaction/system_tests/test_zocalo_system.py index 80d09f201..6dc5120ed 100644 --- a/src/hyperion/external_interaction/system_tests/test_zocalo_system.py +++ b/src/hyperion/external_interaction/system_tests/test_zocalo_system.py @@ -20,7 +20,7 @@ @pytest.mark.s03 def test_when_running_start_stop_then_get_expected_returned_results(zocalo_env): params = GridscanInternalParameters(**default_raw_params()) - zc: XrayCentreZocaloCallback = XrayCentreCallbackCollection.from_params( + zc: XrayCentreZocaloCallback = XrayCentreCallbackCollection.setup( params ).zocalo_handler dcids = [1, 2] @@ -37,7 +37,7 @@ def test_when_running_start_stop_then_get_expected_returned_results(zocalo_env): def run_zocalo_with_dev_ispyb(dummy_params: GridscanInternalParameters, dummy_ispyb_3d): def inner(sample_name="", fallback=np.array([0, 0, 0])): dummy_params.hyperion_params.detector_params.prefix = sample_name - zc: XrayCentreZocaloCallback = XrayCentreCallbackCollection.from_params( + zc: XrayCentreZocaloCallback = XrayCentreCallbackCollection.setup( dummy_params ).zocalo_handler zc.ispyb.ispyb.ISPYB_CONFIG_PATH = dummy_ispyb_3d.ISPYB_CONFIG_PATH diff --git a/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py b/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py index cb9924c7a..386ec4075 100644 --- a/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py @@ -345,6 +345,8 @@ def test_store_3d_grid_scan( assert dummy_ispyb_3d.y_step_size == dummy_params.experiment_params.z_step_size assert dummy_ispyb_3d.y_steps == dummy_params.experiment_params.z_steps + assert dummy_ispyb_3d.upper_left is not None + assert dummy_ispyb_3d.upper_left[0] == x assert dummy_ispyb_3d.upper_left[1] == z @@ -507,6 +509,7 @@ def test_ispyb_deposition_rounds_to_int( mock_ispyb_conn.return_value.__enter__.return_value.mx_acquisition ) mock_upsert_data_collection = mock_mx_aquisition.upsert_data_collection + assert dummy_ispyb.full_params is not None dummy_ispyb.full_params.hyperion_params.ispyb_params.upper_left = np.array( [0.01, 100, 50] ) diff --git a/src/hyperion/system_tests/test_fgs_plan.py b/src/hyperion/system_tests/test_fgs_plan.py index e9050cc28..7e978a1fc 100644 --- a/src/hyperion/system_tests/test_fgs_plan.py +++ b/src/hyperion/system_tests/test_fgs_plan.py @@ -146,13 +146,13 @@ def test_full_plan_tidies_at_end( params: GridscanInternalParameters, RE: RunEngine, ): - callbacks = XrayCentreCallbackCollection.from_params(params) + callbacks = XrayCentreCallbackCollection.setup(params) callbacks.nexus_handler.nexus_writer_1 = MagicMock() callbacks.nexus_handler.nexus_writer_2 = MagicMock() callbacks.ispyb_handler.ispyb_ids = MagicMock() callbacks.ispyb_handler.ispyb.datacollection_ids = MagicMock() with patch( - "hyperion.experiment_plans.flyscan_xray_centre_plan.XrayCentreCallbackCollection.from_params", + "hyperion.experiment_plans.flyscan_xray_centre_plan.XrayCentreCallbackCollection.setup", return_value=callbacks, ): RE(flyscan_xray_centre(fgs_composite, params)) @@ -201,7 +201,7 @@ def test_GIVEN_scan_invalid_WHEN_plan_run_THEN_ispyb_entry_made_but_no_zocalo_en # Currently s03 calls anything with z_steps > 1 invalid params.experiment_params.z_steps = 100 - callbacks = XrayCentreCallbackCollection.from_params(params) + callbacks = XrayCentreCallbackCollection.setup(params) callbacks.ispyb_handler.ispyb.ISPYB_CONFIG_PATH = ISPYB_CONFIG mock_start_zocalo = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_start = mock_start_zocalo @@ -241,7 +241,7 @@ def test_WHEN_plan_run_THEN_move_to_centre_returned_from_zocalo_expected_centre( fgs_composite.eiger.stage = MagicMock() fgs_composite.eiger.unstage = MagicMock() - callbacks = XrayCentreCallbackCollection.from_params(params) + callbacks = XrayCentreCallbackCollection.setup(params) callbacks.ispyb_handler.ispyb.ISPYB_CONFIG_PATH = ISPYB_CONFIG RE(flyscan_xray_centre(fgs_composite, params)) From c631ae6d912cadc629d447cc5ddca282c833d7cd Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 7 Nov 2023 10:11:01 +0000 Subject: [PATCH 04/15] fix more tests --- .../experiment_plans/tests/conftest.py | 4 +- .../tests/test_flyscan_xray_centre_plan.py | 924 +++++++++--------- .../callbacks/ispyb_callback_base.py | 2 +- .../rotation/tests/test_rotation_callbacks.py | 19 +- .../callbacks/xray_centre/ispyb_callback.py | 1 + .../xray_centre/tests/test_ispyb_handler.py | 219 ++--- .../test_xraycentre_callback_collection.py | 102 -- .../xray_centre/tests/test_zocalo_handler.py | 347 +++---- .../ispyb/store_in_ispyb.py | 49 +- 9 files changed, 817 insertions(+), 850 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index bfbb1e60c..994acfec5 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/src/hyperion/experiment_plans/tests/conftest.py @@ -294,9 +294,9 @@ def modified_interactor_mock(assign_run_end: Callable | None = None): return mock -def modified_store_grid_scan_mock(*args, **kwargs): +def modified_store_grid_scan_mock(*args, dcids=(0, 0), dcgid=0, **kwargs): mock = MagicMock(spec=Store3DGridscanInIspyb) - mock.begin_deposition.return_value = ((0, 0), 0, 0) + mock.begin_deposition.return_value = (dcids, dcgid, 0) return mock diff --git a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py index 2526bc498..ceebe284a 100644 --- a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py +++ b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py @@ -1,5 +1,5 @@ import types -from unittest.mock import ANY, MagicMock, call, patch +from unittest.mock import MagicMock, call, patch import bluesky.preprocessors as bpp import numpy as np @@ -47,7 +47,6 @@ TEST_RESULT_MEDIUM, TEST_RESULT_SMALL, ) -from hyperion.external_interaction.zocalo.zocalo_interaction import ZocaloInteractor from hyperion.log import set_up_logging_handlers from hyperion.parameters import external_parameters from hyperion.parameters.constants import ( @@ -60,492 +59,529 @@ ) -def test_given_full_parameters_dict_when_detector_name_used_and_converted_then_detector_constants_correct( - test_fgs_params: GridscanInternalParameters, -): - assert ( - test_fgs_params.hyperion_params.detector_params.detector_size_constants.det_type_string - == EIGER_TYPE_EIGER2_X_16M - ) - raw_params_dict = external_parameters.from_file() - raw_params_dict["hyperion_params"]["detector_params"][ - "detector_size_constants" - ] = EIGER_TYPE_EIGER2_X_4M - params: GridscanInternalParameters = GridscanInternalParameters(**raw_params_dict) - det_dimension = ( - params.hyperion_params.detector_params.detector_size_constants.det_dimension - ) - assert det_dimension == EIGER2_X_4M_DIMENSION +@patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + modified_store_grid_scan_mock, +) +class TestFlyscanXrayCentrePlan: + def test_given_full_parameters_dict_when_detector_name_used_and_converted_then_detector_constants_correct( + self, + test_fgs_params: GridscanInternalParameters, + ): + assert ( + test_fgs_params.hyperion_params.detector_params.detector_size_constants.det_type_string + == EIGER_TYPE_EIGER2_X_16M + ) + raw_params_dict = external_parameters.from_file() + raw_params_dict["hyperion_params"]["detector_params"][ + "detector_size_constants" + ] = EIGER_TYPE_EIGER2_X_4M + params: GridscanInternalParameters = GridscanInternalParameters( + **raw_params_dict + ) + det_dimension = ( + params.hyperion_params.detector_params.detector_size_constants.det_dimension + ) + assert det_dimension == EIGER2_X_4M_DIMENSION + def test_when_run_gridscan_called_then_generator_returned( + self, + ): + plan = run_gridscan(MagicMock(), MagicMock()) + assert isinstance(plan, types.GeneratorType) + + def test_read_hardware_for_ispyb_updates_from_ophyd_devices( + self, + fake_fgs_composite: FlyScanXRayCentreComposite, + test_fgs_params: GridscanInternalParameters, + RE: RunEngine, + ): + undulator_test_value = 1.234 -def test_when_run_gridscan_called_then_generator_returned(): - plan = run_gridscan(MagicMock(), MagicMock()) - assert isinstance(plan, types.GeneratorType) + fake_fgs_composite.undulator.gap.user_readback.sim_put(undulator_test_value) # type: ignore + synchrotron_test_value = "test" + fake_fgs_composite.synchrotron.machine_status.synchrotron_mode.sim_put( # type: ignore + synchrotron_test_value + ) -def test_read_hardware_for_ispyb_updates_from_ophyd_devices( - fake_fgs_composite: FlyScanXRayCentreComposite, - test_fgs_params: GridscanInternalParameters, - RE: RunEngine, -): - undulator_test_value = 1.234 + transmission_test_value = 0.01 + fake_fgs_composite.attenuator.actual_transmission.sim_put(transmission_test_value) # type: ignore + + xgap_test_value = 0.1234 + ygap_test_value = 0.2345 + fake_fgs_composite.s4_slit_gaps.xgap.user_readback.sim_put(xgap_test_value) # type: ignore + fake_fgs_composite.s4_slit_gaps.ygap.user_readback.sim_put(ygap_test_value) # type: ignore + flux_test_value = 10.0 + fake_fgs_composite.flux.flux_reading.sim_put(flux_test_value) # type: ignore + + test_ispyb_callback = GridscanISPyBCallback() + test_ispyb_callback.ispyb = MagicMock(spec=Store3DGridscanInIspyb) + RE.subscribe(test_ispyb_callback) + + @bpp.set_run_key_decorator(GRIDSCAN_OUTER_PLAN) + @bpp.run_decorator( # attach experiment metadata to the start document + md={ + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + ) + def standalone_read_hardware_for_ispyb(und, syn, slits, attn, fl): + yield from read_hardware_for_ispyb_pre_collection(und, syn, slits) + yield from read_hardware_for_ispyb_during_collection(attn, fl) + + RE( + standalone_read_hardware_for_ispyb( + fake_fgs_composite.undulator, + fake_fgs_composite.synchrotron, + fake_fgs_composite.s4_slit_gaps, + fake_fgs_composite.attenuator, + fake_fgs_composite.flux, + ) + ) + params = test_ispyb_callback.params - fake_fgs_composite.undulator.gap.user_readback.sim_put(undulator_test_value) # type: ignore + assert params.hyperion_params.ispyb_params.undulator_gap == undulator_test_value + assert ( + params.hyperion_params.ispyb_params.synchrotron_mode + == synchrotron_test_value + ) + assert params.hyperion_params.ispyb_params.slit_gap_size_x == xgap_test_value + assert params.hyperion_params.ispyb_params.slit_gap_size_y == ygap_test_value + assert ( + params.hyperion_params.ispyb_params.transmission_fraction + == transmission_test_value + ) + assert params.hyperion_params.ispyb_params.flux == flux_test_value + + @patch( + "dodal.devices.aperturescatterguard.ApertureScatterguard._safe_move_within_datacollection_range" + ) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True + ) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True + ) + def test_results_adjusted_and_passed_to_move_xyz( + self, + move_x_y_z: MagicMock, + run_gridscan: MagicMock, + move_aperture: MagicMock, + fake_fgs_composite: FlyScanXRayCentreComposite, + mock_subscriptions: XrayCentreCallbackCollection, + test_fgs_params: GridscanInternalParameters, + RE: RunEngine, + ): + set_up_logging_handlers(logging_level="INFO", dev_mode=True) + RE.subscribe(VerbosePlanExecutionLoggingCallback()) + + mock_subscriptions.ispyb_handler.start( + { + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + ) + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "123abc", + "data": { + "undulator_gap": 0, + "synchrotron_machine_status_synchrotron_mode": 0, + "s4_slit_gaps_xgap": 0, + "s4_slit_gaps_ygap": 0, + }, + } + ) + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "abc123", + "data": { + "attenuator_actual_transmission": 0, + "flux_flux_reading": 10, + }, + } + ) - synchrotron_test_value = "test" - fake_fgs_composite.synchrotron.machine_status.synchrotron_mode.sim_put( # type: ignore - synchrotron_test_value - ) + mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( + TEST_RESULT_LARGE + ) + RE( + run_gridscan_and_move( + fake_fgs_composite, + test_fgs_params, + mock_subscriptions, + ) + ) + mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( + TEST_RESULT_MEDIUM + ) + RE( + run_gridscan_and_move( + fake_fgs_composite, + test_fgs_params, + mock_subscriptions, + ) + ) + mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( + TEST_RESULT_SMALL + ) + RE( + run_gridscan_and_move( + fake_fgs_composite, + test_fgs_params, + mock_subscriptions, + ) + ) + assert fake_fgs_composite.aperture_scatterguard.aperture_positions is not None + ap_call_large = call( + *(fake_fgs_composite.aperture_scatterguard.aperture_positions.LARGE) + ) + ap_call_medium = call( + *(fake_fgs_composite.aperture_scatterguard.aperture_positions.MEDIUM) + ) - transmission_test_value = 0.01 - fake_fgs_composite.attenuator.actual_transmission.sim_put(transmission_test_value) # type: ignore - - xgap_test_value = 0.1234 - ygap_test_value = 0.2345 - fake_fgs_composite.s4_slit_gaps.xgap.user_readback.sim_put(xgap_test_value) # type: ignore - fake_fgs_composite.s4_slit_gaps.ygap.user_readback.sim_put(ygap_test_value) # type: ignore - flux_test_value = 10.0 - fake_fgs_composite.flux.flux_reading.sim_put(flux_test_value) # type: ignore - - test_ispyb_callback = GridscanISPyBCallback() - test_ispyb_callback.ispyb = MagicMock(spec=Store3DGridscanInIspyb) - RE.subscribe(test_ispyb_callback) - - @bpp.set_run_key_decorator(GRIDSCAN_OUTER_PLAN) - @bpp.run_decorator( # attach experiment metadata to the start document - md={ - "subplan_name": GRIDSCAN_OUTER_PLAN, - "hyperion_internal_parameters": test_fgs_params.json(), - } - ) - def standalone_read_hardware_for_ispyb(und, syn, slits, attn, fl): - yield from read_hardware_for_ispyb_pre_collection(und, syn, slits) - yield from read_hardware_for_ispyb_during_collection(attn, fl) - - RE( - standalone_read_hardware_for_ispyb( - fake_fgs_composite.undulator, - fake_fgs_composite.synchrotron, - fake_fgs_composite.s4_slit_gaps, - fake_fgs_composite.attenuator, - fake_fgs_composite.flux, + move_aperture.assert_has_calls( + [ap_call_large, ap_call_large, ap_call_medium], any_order=True ) - ) - params = test_ispyb_callback.params - assert params.hyperion_params.ispyb_params.undulator_gap == undulator_test_value - assert ( - params.hyperion_params.ispyb_params.synchrotron_mode == synchrotron_test_value - ) - assert params.hyperion_params.ispyb_params.slit_gap_size_x == xgap_test_value - assert params.hyperion_params.ispyb_params.slit_gap_size_y == ygap_test_value - assert ( - params.hyperion_params.ispyb_params.transmission_fraction - == transmission_test_value - ) - assert params.hyperion_params.ispyb_params.flux == flux_test_value + mv_call_large = call( + fake_fgs_composite.sample_motors, 0.05, pytest.approx(0.15), 0.25, wait=True + ) + mv_call_medium = call( + fake_fgs_composite.sample_motors, 0.05, pytest.approx(0.15), 0.25, wait=True + ) + move_x_y_z.assert_has_calls( + [mv_call_large, mv_call_large, mv_call_medium], any_order=True + ) + @patch("bluesky.plan_stubs.abs_set", autospec=True) + def test_results_passed_to_move_motors( + self, + bps_abs_set: MagicMock, + test_fgs_params: GridscanInternalParameters, + fake_fgs_composite: FlyScanXRayCentreComposite, + RE: RunEngine, + ): + from hyperion.device_setup_plans.manipulate_sample import move_x_y_z + + set_up_logging_handlers(logging_level="INFO", dev_mode=True) + RE.subscribe(VerbosePlanExecutionLoggingCallback()) + motor_position = ( + test_fgs_params.experiment_params.grid_position_to_motor_position( + np.array([1, 2, 3]) + ) + ) + RE(move_x_y_z(fake_fgs_composite.sample_motors, *motor_position)) + bps_abs_set.assert_has_calls( + [ + call( + fake_fgs_composite.sample_motors.x, + motor_position[0], + group="move_x_y_z", + ), + call( + fake_fgs_composite.sample_motors.y, + motor_position[1], + group="move_x_y_z", + ), + call( + fake_fgs_composite.sample_motors.z, + motor_position[2], + group="move_x_y_z", + ), + ], + any_order=True, + ) -@patch( - "dodal.devices.aperturescatterguard.ApertureScatterguard._safe_move_within_datacollection_range" -) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True) -def test_results_adjusted_and_passed_to_move_xyz( - move_x_y_z: MagicMock, - run_gridscan: MagicMock, - move_aperture: MagicMock, - fake_fgs_composite: FlyScanXRayCentreComposite, - mock_subscriptions: XrayCentreCallbackCollection, - test_fgs_params: GridscanInternalParameters, - RE: RunEngine, -): - set_up_logging_handlers(logging_level="INFO", dev_mode=True) - RE.subscribe(VerbosePlanExecutionLoggingCallback()) - - mock_subscriptions.ispyb_handler.start( - { - "subplan_name": GRIDSCAN_OUTER_PLAN, - "hyperion_internal_parameters": test_fgs_params.json(), - } - ) - mock_subscriptions.ispyb_handler.descriptor( - {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} - ) - mock_subscriptions.ispyb_handler.event( - { - "descriptor": "123abc", - "data": { - "undulator_gap": 0, - "synchrotron_machine_status_synchrotron_mode": 0, - "s4_slit_gaps_xgap": 0, - "s4_slit_gaps_ygap": 0, - }, - } + @patch( + "dodal.devices.aperturescatterguard.ApertureScatterguard._safe_move_within_datacollection_range", ) - mock_subscriptions.ispyb_handler.descriptor( - {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True ) - mock_subscriptions.ispyb_handler.event( - { - "descriptor": "abc123", - "data": { - "attenuator_actual_transmission": 0, - "flux_flux_reading": 10, - }, - } + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True ) + @patch("bluesky.plan_stubs.rd") + @patch( + "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor", + modified_interactor_mock, + ) + def test_individual_plans_triggered_once_and_only_once_in_composite_run( + self, + rd: MagicMock, + move_xyz: MagicMock, + run_gridscan: MagicMock, + move_aperture: MagicMock, + mock_subscriptions: XrayCentreCallbackCollection, + fake_fgs_composite: FlyScanXRayCentreComposite, + test_fgs_params: GridscanInternalParameters, + RE: RunEngine, + ): + td = TestData() + mock_subscriptions.ispyb_handler.start(td.test_start_document) + mock_subscriptions.zocalo_handler.start(td.test_start_document) + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} + ) - mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( - TEST_RESULT_LARGE - ) - RE( - run_gridscan_and_move( - fake_fgs_composite, - test_fgs_params, - mock_subscriptions, + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "123abc", + "data": { + "undulator_gap": 0, + "synchrotron_machine_status_synchrotron_mode": 0, + "s4_slit_gaps_xgap": 0, + "s4_slit_gaps_ygap": 0, + }, + } ) - ) - mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( - TEST_RESULT_MEDIUM - ) - RE( - run_gridscan_and_move( - fake_fgs_composite, - test_fgs_params, - mock_subscriptions, + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} ) - ) - mock_subscriptions.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( - TEST_RESULT_SMALL - ) - RE( - run_gridscan_and_move( - fake_fgs_composite, - test_fgs_params, - mock_subscriptions, + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "abc123", + "data": { + "attenuator_actual_transmission": 0, + "flux_flux_reading": 10, + }, + } ) - ) - assert fake_fgs_composite.aperture_scatterguard.aperture_positions is not None - ap_call_large = call( - *(fake_fgs_composite.aperture_scatterguard.aperture_positions.LARGE) - ) - ap_call_medium = call( - *(fake_fgs_composite.aperture_scatterguard.aperture_positions.MEDIUM) - ) - move_aperture.assert_has_calls( - [ap_call_large, ap_call_large, ap_call_medium], any_order=True - ) + set_up_logging_handlers(logging_level="INFO", dev_mode=True) + RE.subscribe(VerbosePlanExecutionLoggingCallback()) - mv_call_large = call( - fake_fgs_composite.sample_motors, 0.05, pytest.approx(0.15), 0.25, wait=True - ) - mv_call_medium = call( - fake_fgs_composite.sample_motors, 0.05, pytest.approx(0.15), 0.25, wait=True - ) - move_x_y_z.assert_has_calls( - [mv_call_large, mv_call_large, mv_call_medium], any_order=True + RE( + run_gridscan_and_move( + fake_fgs_composite, + test_fgs_params, + mock_subscriptions, + ) + ) + + run_gridscan.assert_called_once_with(fake_fgs_composite, test_fgs_params) + array_arg = move_xyz.call_args.args[1:4] + np.testing.assert_allclose(array_arg, np.array([0.05, 0.15, 0.25])) + move_xyz.assert_called_once() + + @patch( + "dodal.devices.aperturescatterguard.ApertureScatterguard._safe_move_within_datacollection_range", + autospec=True, ) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True + ) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True + ) + @patch("bluesky.plan_stubs.rd") + def test_logging_within_plan( + self, + rd: MagicMock, + move_xyz: MagicMock, + run_gridscan: MagicMock, + move_aperture: MagicMock, + mock_subscriptions: XrayCentreCallbackCollection, + fake_fgs_composite: FlyScanXRayCentreComposite, + test_fgs_params: GridscanInternalParameters, + RE: RunEngine, + ): + mock_subscriptions.ispyb_handler.start({}) + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "123abc", + "data": { + "undulator_gap": 0, + "synchrotron_machine_status_synchrotron_mode": 0, + "s4_slit_gaps_xgap": 0, + "s4_slit_gaps_ygap": 0, + }, + } + ) + mock_subscriptions.ispyb_handler.descriptor( + {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} + ) + mock_subscriptions.ispyb_handler.event( + { + "descriptor": "abc123", + "data": { + "attenuator_actual_transmission": 0, + "flux_flux_reading": 10, + }, + } + ) -@patch("bluesky.plan_stubs.abs_set", autospec=True) -def test_results_passed_to_move_motors( - bps_abs_set: MagicMock, - test_fgs_params: GridscanInternalParameters, - fake_fgs_composite: FlyScanXRayCentreComposite, - RE: RunEngine, -): - from hyperion.device_setup_plans.manipulate_sample import move_x_y_z + set_up_logging_handlers(logging_level="INFO", dev_mode=True) + RE.subscribe(VerbosePlanExecutionLoggingCallback()) - set_up_logging_handlers(logging_level="INFO", dev_mode=True) - RE.subscribe(VerbosePlanExecutionLoggingCallback()) - motor_position = test_fgs_params.experiment_params.grid_position_to_motor_position( - np.array([1, 2, 3]) - ) - RE(move_x_y_z(fake_fgs_composite.sample_motors, *motor_position)) - bps_abs_set.assert_has_calls( - [ - call( - fake_fgs_composite.sample_motors.x, - motor_position[0], - group="move_x_y_z", - ), - call( - fake_fgs_composite.sample_motors.y, - motor_position[1], - group="move_x_y_z", - ), - call( - fake_fgs_composite.sample_motors.z, - motor_position[2], - group="move_x_y_z", - ), - ], - any_order=True, - ) + RE( + run_gridscan_and_move( + fake_fgs_composite, + test_fgs_params, + mock_subscriptions, + ) + ) + run_gridscan.assert_called_once_with(fake_fgs_composite, test_fgs_params) + array_arg = move_xyz.call_args.args[1:4] + np.testing.assert_array_almost_equal(array_arg, np.array([0.05, 0.15, 0.25])) + move_xyz.assert_called_once() -@patch( - "dodal.devices.aperturescatterguard.ApertureScatterguard._safe_move_within_datacollection_range", -) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True) -@patch("bluesky.plan_stubs.rd") -@patch( - "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor", - modified_interactor_mock, -) -def test_individual_plans_triggered_once_and_only_once_in_composite_run( - rd: MagicMock, - move_xyz: MagicMock, - run_gridscan: MagicMock, - move_aperture: MagicMock, - mock_subscriptions: XrayCentreCallbackCollection, - fake_fgs_composite: FlyScanXRayCentreComposite, - test_fgs_params: GridscanInternalParameters, - RE: RunEngine, -): - td = TestData() - mock_subscriptions.ispyb_handler.start(td.test_start_document) - mock_subscriptions.zocalo_handler.start(td.test_start_document) - mock_subscriptions.ispyb_handler.descriptor( - {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.sleep", autospec=True ) + def test_GIVEN_scan_already_valid_THEN_wait_for_GRIDSCAN_returns_immediately( + self, patch_sleep: MagicMock, RE: RunEngine + ): + test_fgs: FastGridScan = make_fake_device(FastGridScan)( + "prefix", name="fake_fgs" + ) - mock_subscriptions.ispyb_handler.event( - { - "descriptor": "123abc", - "data": { - "undulator_gap": 0, - "synchrotron_machine_status_synchrotron_mode": 0, - "s4_slit_gaps_xgap": 0, - "s4_slit_gaps_ygap": 0, - }, - } - ) - mock_subscriptions.ispyb_handler.descriptor( - {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} - ) - mock_subscriptions.ispyb_handler.event( - { - "descriptor": "abc123", - "data": { - "attenuator_actual_transmission": 0, - "flux_flux_reading": 10, - }, - } - ) + test_fgs.scan_invalid.sim_put(False) # type: ignore + test_fgs.position_counter.sim_put(0) # type: ignore - set_up_logging_handlers(logging_level="INFO", dev_mode=True) - RE.subscribe(VerbosePlanExecutionLoggingCallback()) + RE(wait_for_gridscan_valid(test_fgs)) - RE( - run_gridscan_and_move( - fake_fgs_composite, - test_fgs_params, - mock_subscriptions, - ) + patch_sleep.assert_not_called() + + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.sleep", autospec=True ) + def test_GIVEN_scan_not_valid_THEN_wait_for_GRIDSCAN_raises_and_sleeps_called( + self, patch_sleep: MagicMock, RE: RunEngine + ): + test_fgs: FastGridScan = make_fake_device(FastGridScan)( + "prefix", name="fake_fgs" + ) - run_gridscan.assert_called_once_with(fake_fgs_composite, test_fgs_params) - array_arg = move_xyz.call_args.args[1:4] - np.testing.assert_allclose(array_arg, np.array([0.05, 0.15, 0.25])) - move_xyz.assert_called_once() + test_fgs.scan_invalid.sim_put(True) # type: ignore + test_fgs.position_counter.sim_put(0) # type: ignore + with pytest.raises(WarningException): + RE(wait_for_gridscan_valid(test_fgs)) + patch_sleep.assert_called() -@patch( - "dodal.devices.aperturescatterguard.ApertureScatterguard._safe_move_within_datacollection_range", - autospec=True, -) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.run_gridscan", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.move_x_y_z", autospec=True) -@patch("bluesky.plan_stubs.rd") -def test_logging_within_plan( - rd: MagicMock, - move_xyz: MagicMock, - run_gridscan: MagicMock, - move_aperture: MagicMock, - mock_subscriptions: XrayCentreCallbackCollection, - fake_fgs_composite: FlyScanXRayCentreComposite, - test_fgs_params: GridscanInternalParameters, - RE: RunEngine, -): - mock_subscriptions.ispyb_handler.start({}) - mock_subscriptions.ispyb_handler.descriptor( - {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.abs_set", autospec=True ) - - mock_subscriptions.ispyb_handler.event( - { - "descriptor": "123abc", - "data": { - "undulator_gap": 0, - "synchrotron_machine_status_synchrotron_mode": 0, - "s4_slit_gaps_xgap": 0, - "s4_slit_gaps_ygap": 0, - }, - } + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.kickoff", autospec=True ) - mock_subscriptions.ispyb_handler.descriptor( - {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.complete", autospec=True ) - mock_subscriptions.ispyb_handler.event( - { - "descriptor": "abc123", - "data": { - "attenuator_actual_transmission": 0, - "flux_flux_reading": 10, - }, - } + @patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.mv", autospec=True) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.wait_for_gridscan_valid", + autospec=True, ) - - set_up_logging_handlers(logging_level="INFO", dev_mode=True) - RE.subscribe(VerbosePlanExecutionLoggingCallback()) - - RE( - run_gridscan_and_move( - fake_fgs_composite, - test_fgs_params, - mock_subscriptions, + @patch( + "hyperion.external_interaction.nexus.write_nexus.NexusWriter", + autospec=True, + spec_set=True, + ) + def test_when_grid_scan_ran_then_eiger_disarmed_before_zocalo_end( + self, + nexuswriter, + wait_for_valid, + mock_mv, + mock_complete, + mock_kickoff, + mock_abs_set, + fake_fgs_composite: FlyScanXRayCentreComposite, + test_fgs_params: GridscanInternalParameters, + mock_subscriptions: XrayCentreCallbackCollection, + RE: RunEngine, + ): + # Put both mocks in a parent to easily capture order + mock_parent = MagicMock() + fake_fgs_composite.eiger.disarm_detector = mock_parent.disarm + + fake_fgs_composite.eiger.filewriters_finished = Status(done=True, success=True) # type: ignore + fake_fgs_composite.eiger.odin.check_odin_state = MagicMock(return_value=True) + fake_fgs_composite.eiger.odin.file_writer.num_captured.sim_put(1200) # type: ignore + fake_fgs_composite.eiger.stage = MagicMock( + return_value=Status(None, None, 0, True, True) ) - ) - - run_gridscan.assert_called_once_with(fake_fgs_composite, test_fgs_params) - array_arg = move_xyz.call_args.args[1:4] - np.testing.assert_array_almost_equal(array_arg, np.array([0.05, 0.15, 0.25])) - move_xyz.assert_called_once() - - -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.sleep", autospec=True) -def test_GIVEN_scan_already_valid_THEN_wait_for_GRIDSCAN_returns_immediately( - patch_sleep: MagicMock, RE: RunEngine -): - test_fgs: FastGridScan = make_fake_device(FastGridScan)("prefix", name="fake_fgs") - - test_fgs.scan_invalid.sim_put(False) # type: ignore - test_fgs.position_counter.sim_put(0) # type: ignore - - RE(wait_for_gridscan_valid(test_fgs)) - - patch_sleep.assert_not_called() - - -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.sleep", autospec=True) -def test_GIVEN_scan_not_valid_THEN_wait_for_GRIDSCAN_raises_and_sleeps_called( - patch_sleep: MagicMock, RE: RunEngine -): - test_fgs: FastGridScan = make_fake_device(FastGridScan)("prefix", name="fake_fgs") - - test_fgs.scan_invalid.sim_put(True) # type: ignore - test_fgs.position_counter.sim_put(0) # type: ignore - with pytest.raises(WarningException): - RE(wait_for_gridscan_valid(test_fgs)) + fake_fgs_composite.xbpm_feedback.pos_stable.sim_put(1) # type: ignore + + with patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.XrayCentreCallbackCollection.setup", + lambda: mock_subscriptions, + ), patch( + "hyperion.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter.create_nexus_file", + autospec=True, + ), patch( + "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor", + lambda _: modified_interactor_mock(mock_parent.run_end), + ): + RE(flyscan_xray_centre(fake_fgs_composite, test_fgs_params)) + + mock_parent.assert_has_calls([call.disarm(), call.run_end(0), call.run_end(0)]) + + @patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.wait", autospec=True) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.complete", autospec=True + ) + def test_fgs_arms_eiger_without_grid_detect( + self, + mock_complete, + mock_wait, + fake_fgs_composite: FlyScanXRayCentreComposite, + test_fgs_params: GridscanInternalParameters, + RE: RunEngine, + ): + fake_fgs_composite.eiger.stage = MagicMock() + fake_fgs_composite.eiger.unstage = MagicMock() - patch_sleep.assert_called() + RE(run_gridscan(fake_fgs_composite, test_fgs_params)) + fake_fgs_composite.eiger.stage.assert_called_once() + fake_fgs_composite.eiger.unstage.assert_called_once() + + @patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.wait", autospec=True) + @patch( + "hyperion.experiment_plans.flyscan_xray_centre_plan.bps.complete", autospec=True + ) + def test_when_grid_scan_fails_then_detector_disarmed_and_correct_exception_returned( + self, + mock_complete, + mock_wait, + fake_fgs_composite: FlyScanXRayCentreComposite, + test_fgs_params: GridscanInternalParameters, + RE: RunEngine, + ): + class CompleteException(Exception): + pass + mock_complete.side_effect = CompleteException() -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.abs_set", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.kickoff", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.complete", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.mv", autospec=True) -@patch( - "hyperion.experiment_plans.flyscan_xray_centre_plan.wait_for_gridscan_valid", - autospec=True, -) -@patch( - "hyperion.external_interaction.nexus.write_nexus.NexusWriter", - autospec=True, - spec_set=True, -) -def test_when_grid_scan_ran_then_eiger_disarmed_before_zocalo_end( - nexuswriter, - wait_for_valid, - mock_mv, - mock_complete, - mock_kickoff, - mock_abs_set, - fake_fgs_composite: FlyScanXRayCentreComposite, - test_fgs_params: GridscanInternalParameters, - mock_subscriptions: XrayCentreCallbackCollection, - RE: RunEngine, -): - # Put both mocks in a parent to easily capture order - mock_parent = MagicMock() - fake_fgs_composite.eiger.disarm_detector = mock_parent.disarm - - fake_fgs_composite.eiger.filewriters_finished = Status(done=True, success=True) - fake_fgs_composite.eiger.odin.check_odin_state = MagicMock(return_value=True) - fake_fgs_composite.eiger.odin.file_writer.num_captured.sim_put(1200) # type: ignore - fake_fgs_composite.eiger.stage = MagicMock( - return_value=Status(None, None, 0, True, True) - ) - fake_fgs_composite.xbpm_feedback.pos_stable.sim_put(1) # type: ignore - - with patch( - "hyperion.experiment_plans.flyscan_xray_centre_plan.XrayCentreCallbackCollection.setup", - lambda: mock_subscriptions, - ), patch( - "hyperion.external_interaction.callbacks.xray_centre.nexus_callback.NexusWriter.create_nexus_file", - autospec=True, - ), patch( - "hyperion.external_interaction.callbacks.xray_centre.zocalo_callback.ZocaloInteractor", - lambda _: modified_interactor_mock(mock_parent.run_end), - ), patch( - "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", - modified_store_grid_scan_mock, - ): - RE(flyscan_xray_centre(fake_fgs_composite, test_fgs_params)) - - mock_parent.assert_has_calls([call.disarm(), call.run_end(0), call.run_end(0)]) - - -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.wait", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.complete", autospec=True) -def test_fgs_arms_eiger_without_grid_detect( - mock_complete, - mock_wait, - fake_fgs_composite: FlyScanXRayCentreComposite, - test_fgs_params: GridscanInternalParameters, - RE: RunEngine, -): - fake_fgs_composite.eiger.stage = MagicMock() - fake_fgs_composite.eiger.unstage = MagicMock() - - RE(run_gridscan(fake_fgs_composite, test_fgs_params)) - fake_fgs_composite.eiger.stage.assert_called_once() - fake_fgs_composite.eiger.unstage.assert_called_once() - - -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.wait", autospec=True) -@patch("hyperion.experiment_plans.flyscan_xray_centre_plan.bps.complete", autospec=True) -def test_when_grid_scan_fails_then_detector_disarmed_and_correct_exception_returned( - mock_complete, - mock_wait, - fake_fgs_composite: FlyScanXRayCentreComposite, - test_fgs_params: GridscanInternalParameters, - RE: RunEngine, -): - class CompleteException(Exception): - pass - - mock_complete.side_effect = CompleteException() - - fake_fgs_composite.eiger.stage = MagicMock( - return_value=Status(None, None, 0, True, True) - ) + fake_fgs_composite.eiger.stage = MagicMock( + return_value=Status(None, None, 0, True, True) + ) - fake_fgs_composite.eiger.odin.check_odin_state = MagicMock() + fake_fgs_composite.eiger.odin.check_odin_state = MagicMock() - fake_fgs_composite.eiger.disarm_detector = MagicMock() - fake_fgs_composite.eiger.disable_roi_mode = MagicMock() + fake_fgs_composite.eiger.disarm_detector = MagicMock() + fake_fgs_composite.eiger.disable_roi_mode = MagicMock() - # Without the complete finishing we will not get all the images - fake_fgs_composite.eiger.ALL_FRAMES_TIMEOUT = 0.1 # type: ignore + # Without the complete finishing we will not get all the images + fake_fgs_composite.eiger.ALL_FRAMES_TIMEOUT = 0.1 # type: ignore - # Want to get the underlying completion error, not the one raised from unstage - with pytest.raises(CompleteException): - RE(run_gridscan(fake_fgs_composite, test_fgs_params)) + # Want to get the underlying completion error, not the one raised from unstage + with pytest.raises(CompleteException): + RE(run_gridscan(fake_fgs_composite, test_fgs_params)) - fake_fgs_composite.eiger.disable_roi_mode.assert_called() - fake_fgs_composite.eiger.disarm_detector.assert_called() + fake_fgs_composite.eiger.disable_roi_mode.assert_called() + fake_fgs_composite.eiger.disarm_detector.assert_called() diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index d62b5f6ba..cde6b1ee3 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -110,7 +110,7 @@ def stop(self, doc: dict): exit_status = ( doc.get("exit_status") or "Exit status not available in stop document!" ) - reason = doc.get("reason") or "Unknown failure reason!" + reason = doc.get("reason") or "" set_dcgid_tag(None) try: self.ispyb.end_deposition(exit_status, reason) diff --git a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index 81720dc87..c9e0eb750 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -54,8 +54,8 @@ def RE(): def fake_rotation_scan( params: RotationInternalParameters, subscriptions: RotationCallbackCollection, - after_open_assert: Callable | None = None, - after_main_assert: Callable | None = None, + after_open_do: Callable | None = None, + after_main_do: Callable | None = None, ): attenuator = make_fake_device(Attenuator)(name="attenuator") flux = make_fake_device(Flux)(name="flux") @@ -69,8 +69,8 @@ def fake_rotation_scan( } ) def plan(): - if after_open_assert: - after_open_assert(subscriptions) + if after_open_do: + after_open_do(subscriptions) @bpp.set_run_key_decorator(ROTATION_PLAN_MAIN) @bpp.run_decorator( @@ -80,9 +80,8 @@ def plan(): ) def fake_main_plan(): yield from read_hardware_for_ispyb_during_collection(attenuator, flux) - subscriptions.ispyb_handler.ispyb_ids = (0, 0) - if after_main_assert: - after_main_assert(subscriptions) + if after_main_do: + after_main_do(subscriptions) yield from bps.sleep(0) yield from fake_main_plan() @@ -199,9 +198,11 @@ def test_zocalo_start_and_end_triggered_once( cb.ispyb_handler.stop = MagicMock(autospec=True) cb.ispyb_handler.ispyb = MagicMock(spec=StoreRotationInIspyb) cb.ispyb_handler.params = params - cb.ispyb_handler.ispyb_ids = (0, 0) - RE(fake_rotation_scan(params, cb)) + def set_ispyb_ids(cbs): + cbs.ispyb_handler.ispyb_ids = (0, 0) + + RE(fake_rotation_scan(params, cb, after_open_do=set_ispyb_ids)) zocalo.assert_called_once() cb.zocalo_handler.zocalo_interactor.run_start.assert_called_once() 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 8568529cb..3c3a8526b 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -51,6 +51,7 @@ def start(self, doc: dict): else Store2DGridscanInIspyb(self.ispyb_config, self.params) ) self.run_start_uid = doc.get("uid") + self.uid_to_finalize_on = doc.get("uid") def append_to_comment(self, comment: str): for id in self.ispyb_ids[0]: diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py index fa348fc2e..d42f11ee6 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py @@ -8,6 +8,7 @@ GridscanISPyBCallback, ) from hyperion.external_interaction.callbacks.xray_centre.tests.conftest import TestData +from hyperion.external_interaction.ispyb.store_in_ispyb import Store3DGridscanInIspyb from hyperion.log import LOGGER, set_up_logging_handlers from hyperion.parameters.external_parameters import from_file as default_raw_params from hyperion.parameters.plan_specific.gridscan_internal_params import ( @@ -19,77 +20,6 @@ td = TestData() -@pytest.fixture -def dummy_params(): - return GridscanInternalParameters(**default_raw_params()) - - -def test_fgs_failing_results_in_bad_run_status_in_ispyb( - mock_ispyb_update_time_and_status: MagicMock, - mock_ispyb_get_time: MagicMock, - mock_ispyb_store_grid_scan: MagicMock, - dummy_params, - ispyb_handler, -): - mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] - mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING - mock_ispyb_update_time_and_status.return_value = None - - ispyb_handler.start(td.test_start_document) - ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) - ispyb_handler.event(td.test_event_document_pre_data_collection) - ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) - ispyb_handler.event(td.test_event_document_during_data_collection) - ispyb_handler.stop(td.test_run_gridscan_failed_stop_document) - - mock_ispyb_update_time_and_status.assert_has_calls( - [ - call( - td.DUMMY_TIME_STRING, - td.BAD_ISPYB_RUN_STATUS, - "could not connect to devices", - id, - DCG_ID, - ) - for id in DC_IDS - ] - ) - assert mock_ispyb_update_time_and_status.call_count == len(DC_IDS) - - -def test_fgs_raising_no_exception_results_in_good_run_status_in_ispyb( - mock_ispyb_update_time_and_status: MagicMock, - mock_ispyb_get_time: MagicMock, - mock_ispyb_store_grid_scan: MagicMock, - dummy_params, - ispyb_handler, -): - mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] - mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING - mock_ispyb_update_time_and_status.return_value = None - ispyb_handler = GridscanISPyBCallback() - ispyb_handler.start(td.test_start_document) - ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) - ispyb_handler.event(td.test_event_document_pre_data_collection) - ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) - ispyb_handler.event(td.test_event_document_during_data_collection) - ispyb_handler.stop(td.test_do_fgs_gridscan_stop_document) - - mock_ispyb_update_time_and_status.assert_has_calls( - [ - call( - td.DUMMY_TIME_STRING, - td.GOOD_ISPYB_RUN_STATUS, - "", - id, - DCG_ID, - ) - for id in DC_IDS - ] - ) - assert mock_ispyb_update_time_and_status.call_count == len(DC_IDS) - - @pytest.fixture def mock_emit(): with patch("hyperion.log.setup_dodal_logging"): @@ -105,44 +35,109 @@ def mock_emit(): dodal_logger.removeHandler(test_handler) -def test_given_ispyb_callback_started_writing_to_ispyb_when_messages_logged_then_they_contain_dcgid( - mock_emit, mock_ispyb_store_grid_scan: MagicMock, dummy_params -): - mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] - ispyb_handler = GridscanISPyBCallback() - ispyb_handler.start(td.test_start_document) - ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) - ispyb_handler.event(td.test_event_document_pre_data_collection) - ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) - ispyb_handler.event(td.test_event_document_during_data_collection) - - for logger in [LOGGER, dodal_logger]: - logger.info("test") - latest_record = mock_emit.call_args.args[-1] - assert latest_record.dc_group_id == DCG_ID - - -def test_given_ispyb_callback_finished_writing_to_ispyb_when_messages_logged_then_they_do_not_contain_dcgid( - mock_emit, - mock_ispyb_store_grid_scan: MagicMock, - mock_ispyb_update_time_and_status: MagicMock, - mock_ispyb_get_time: MagicMock, - dummy_params, - ispyb_handler, -): - mock_ispyb_store_grid_scan.return_value = [DC_IDS, None, DCG_ID] - mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING - mock_ispyb_update_time_and_status.return_value = None - ispyb_handler = GridscanISPyBCallback() - ispyb_handler.start(td.test_start_document) - ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) - ispyb_handler.event(td.test_event_document_pre_data_collection) - ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) - ispyb_handler.event(td.test_event_document_during_data_collection) - ispyb_handler.stop(td.test_run_gridscan_failed_stop_document) - - for logger in [LOGGER, dodal_logger]: - logger.info("test") - - latest_record = mock_emit.call_args.args[-1] - assert not hasattr(latest_record, "dc_group_id") +@pytest.fixture +def dummy_params(): + return GridscanInternalParameters(**default_raw_params()) + + +def mock_store_in_ispyb(dummy_params, *args, **kwargs) -> Store3DGridscanInIspyb: + mock = Store3DGridscanInIspyb(None, dummy_params) + mock.store_grid_scan = MagicMock(return_value=[DC_IDS, None, DCG_ID]) + mock.get_current_time_string = MagicMock(return_value=td.DUMMY_TIME_STRING) + mock.update_scan_with_end_time_and_status = MagicMock(return_value=None) + return mock + + +@patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + mock_store_in_ispyb, +) +class TestXrayCentreIspybHandler: + def test_fgs_failing_results_in_bad_run_status_in_ispyb( + self, + ): + ispyb_handler = GridscanISPyBCallback() + ispyb_handler.start(td.test_start_document) + ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) + ispyb_handler.event(td.test_event_document_pre_data_collection) + ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) + ispyb_handler.event(td.test_event_document_during_data_collection) + ispyb_handler.stop(td.test_run_gridscan_failed_stop_document) + + ispyb_handler.ispyb.update_scan_with_end_time_and_status.assert_has_calls( + [ + call( + td.DUMMY_TIME_STRING, + td.BAD_ISPYB_RUN_STATUS, + "could not connect to devices", + id, + DCG_ID, + ) + for id in DC_IDS + ] + ) + assert ( + ispyb_handler.ispyb.update_scan_with_end_time_and_status.call_count + == len(DC_IDS) + ) + + def test_fgs_raising_no_exception_results_in_good_run_status_in_ispyb( + self, + ): + ispyb_handler = GridscanISPyBCallback() + ispyb_handler.start(td.test_start_document) + ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) + ispyb_handler.event(td.test_event_document_pre_data_collection) + ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) + ispyb_handler.event(td.test_event_document_during_data_collection) + ispyb_handler.stop(td.test_do_fgs_gridscan_stop_document) + + ispyb_handler.ispyb.update_scan_with_end_time_and_status.assert_has_calls( + [ + call( + td.DUMMY_TIME_STRING, + td.GOOD_ISPYB_RUN_STATUS, + "", + id, + DCG_ID, + ) + for id in DC_IDS + ] + ) + assert ( + ispyb_handler.ispyb.update_scan_with_end_time_and_status.call_count + == len(DC_IDS) + ) + + def test_given_ispyb_callback_started_writing_to_ispyb_when_messages_logged_then_they_contain_dcgid( + self, mock_emit + ): + ispyb_handler = GridscanISPyBCallback() + ispyb_handler.start(td.test_start_document) + ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) + ispyb_handler.event(td.test_event_document_pre_data_collection) + ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) + ispyb_handler.event(td.test_event_document_during_data_collection) + + for logger in [LOGGER, dodal_logger]: + logger.info("test") + latest_record = mock_emit.call_args.args[-1] + assert latest_record.dc_group_id == DCG_ID + + def test_given_ispyb_callback_finished_writing_to_ispyb_when_messages_logged_then_they_do_not_contain_dcgid( + self, + mock_emit, + ): + ispyb_handler = GridscanISPyBCallback() + ispyb_handler.start(td.test_start_document) + ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) + ispyb_handler.event(td.test_event_document_pre_data_collection) + ispyb_handler.descriptor(td.test_descriptor_document_during_data_collection) + ispyb_handler.event(td.test_event_document_during_data_collection) + ispyb_handler.stop(td.test_run_gridscan_failed_stop_document) + + for logger in [LOGGER, dodal_logger]: + logger.info("test") + + latest_record = mock_emit.call_args.args[-1] + assert not hasattr(latest_record, "dc_group_id") diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py index ab69fc133..ace8475ef 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_xraycentre_callback_collection.py @@ -1,121 +1,19 @@ -from typing import Generator from unittest.mock import MagicMock import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp -import pytest from bluesky.run_engine import RunEngine -from dodal.devices.eiger import DetectorParams, EigerDetector -from hyperion.experiment_plans.flyscan_xray_centre_plan import ( - FlyScanXRayCentreComposite, - run_gridscan_and_move, -) from hyperion.external_interaction.callbacks.xray_centre.callback_collection import ( XrayCentreCallbackCollection, ) -from hyperion.parameters.constants import SIM_BEAMLINE -from hyperion.parameters.external_parameters import from_file as default_raw_params -from hyperion.parameters.plan_specific.gridscan_internal_params import ( - GridscanInternalParameters, -) def test_callback_collection_init(): - test_parameters = GridscanInternalParameters(**default_raw_params()) callbacks = XrayCentreCallbackCollection.setup() - assert ( - callbacks.ispyb_handler.params.experiment_params - == test_parameters.experiment_params - ) - assert ( - callbacks.ispyb_handler.params.hyperion_params.detector_params - == test_parameters.hyperion_params.detector_params - ) - assert ( - callbacks.ispyb_handler.params.hyperion_params.ispyb_params - == test_parameters.hyperion_params.ispyb_params - ) - assert ( - callbacks.ispyb_handler.params.hyperion_params - == test_parameters.hyperion_params - ) - assert callbacks.ispyb_handler.params == test_parameters - assert callbacks.zocalo_handler.ispyb == callbacks.ispyb_handler assert len(list(callbacks)) == 3 -@pytest.fixture() -def eiger() -> Generator[EigerDetector, None, None]: - detector_params: DetectorParams = DetectorParams( - current_energy_ev=100, - exposure_time=0.1, - directory="/tmp", - prefix="file_name", - detector_distance=100.0, - omega_start=0.0, - omega_increment=0.1, - num_images_per_trigger=50, - num_triggers=1, - use_roi_mode=False, - run_number=0, - det_dist_to_beam_converter_path="src/hyperion/unit_tests/test_lookup_table.txt", - ) - eiger = EigerDetector( - detector_params=detector_params, name="eiger", prefix="BL03S-EA-EIGER-01:" - ) - - # Otherwise odin moves too fast to be tested - eiger.cam.manual_trigger.put("Yes") - - # S03 currently does not have StaleParameters_RBV - # eiger.wait_for_stale_parameters = lambda: None - eiger.odin.check_odin_initialised = lambda: (True, "") - - yield eiger - - -@pytest.mark.skip( - reason="Needs better S03 or some other workaround for eiger/odin timeout." -) -@pytest.mark.s03 -def test_communicator_in_composite_run( - nexus_writer: MagicMock, - ispyb_begin_deposition: MagicMock, - ispyb_end_deposition: MagicMock, -): - nexus_writer.side_effect = [MagicMock(), MagicMock()] - RE = RunEngine({}) - - params = GridscanInternalParameters(**default_raw_params()) - params.hyperion_params.beamline = SIM_BEAMLINE - ispyb_begin_deposition.return_value = ([1, 2], None, 4) - - callbacks = XrayCentreCallbackCollection.setup() - callbacks.zocalo_handler.zocalo_interactor.wait_for_result = MagicMock() - callbacks.zocalo_handler.zocalo_interactor.run_start = MagicMock() - callbacks.zocalo_handler.zocalo_interactor.run_end = MagicMock() - # callbacks.zocalo_handler.xray_centre_motor_position = np.array([1, 2, 3]) - - flyscan_xray_centre_composite = MagicMock(spec=FlyScanXRayCentreComposite) - # this is where it's currently getting stuck: - # flyscan_xray_centre_composite.fast_grid_scan.is_invalid = lambda: False - # but this is not a solution - # Would be better to use flyscan_xray_centre instead but eiger doesn't work well in S03 - RE(run_gridscan_and_move(flyscan_xray_centre_composite, params, callbacks)) - - # nexus writing - callbacks.nexus_handler.nexus_writer_1.assert_called_once() # type: ignore - callbacks.nexus_handler.nexus_writer_2.assert_called_once() # type: ignore - # ispyb - ispyb_begin_deposition.assert_called_once() - ispyb_end_deposition.assert_called_once() - # zocalo - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once() - callbacks.zocalo_handler.zocalo_interactor.run_start.assert_called_once() - callbacks.zocalo_handler.zocalo_interactor.run_end.assert_called_once() - - def test_callback_collection_list(): callbacks = XrayCentreCallbackCollection.setup() callback_list = list(callbacks) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py index 4d71d72e3..2af74c434 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py @@ -3,6 +3,7 @@ import numpy as np import pytest +from hyperion.experiment_plans.tests.conftest import modified_store_grid_scan_mock from hyperion.external_interaction.callbacks.xray_centre.callback_collection import ( XrayCentreCallbackCollection, ) @@ -30,191 +31,205 @@ def dummy_params(): return GridscanInternalParameters(**default_raw_params()) -def mock_zocalo_functions(callbacks: XrayCentreCallbackCollection): - callbacks.ispyb_handler.start(td.test_start_document) - callbacks.zocalo_handler.start(td.test_start_document) +def init_cbs_with_docs_and_mock_zocalo_and_ispyb( + callbacks: XrayCentreCallbackCollection, dcids=(0, 0), dcgid=4 +): + with patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + lambda _, __: modified_store_grid_scan_mock(dcids=dcids, dcgid=dcgid), + ): + callbacks.ispyb_handler.start(td.test_start_document) + callbacks.zocalo_handler.start(td.test_start_document) callbacks.zocalo_handler.zocalo_interactor.wait_for_result = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_end = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_start = MagicMock() -def test_execution_of_run_gridscan_triggers_zocalo_calls( - mock_ispyb_update_time_and_status: MagicMock, - mock_ispyb_get_time: MagicMock, - mock_ispyb_store_grid_scan: MagicMock, - nexus_writer: MagicMock, - dummy_params, -): - dc_ids = [1, 2] - dcg_id = 4 - - mock_ispyb_store_grid_scan.return_value = [dc_ids, None, dcg_id] - mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING - mock_ispyb_update_time_and_status.return_value = None - - callbacks = XrayCentreCallbackCollection.setup() - mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) - callbacks.zocalo_handler.start(td.test_run_gridscan_start_document) - callbacks.ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) - callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) - callbacks.ispyb_handler.descriptor( - td.test_descriptor_document_during_data_collection - ) - callbacks.ispyb_handler.event(td.test_event_document_during_data_collection) - callbacks.zocalo_handler.start(td.test_do_fgs_start_document) - callbacks.ispyb_handler.stop(td.test_stop_document) - callbacks.zocalo_handler.stop(td.test_stop_document) - - callbacks.zocalo_handler.zocalo_interactor.run_start.assert_has_calls( - [call(x) for x in dc_ids] - ) - assert callbacks.zocalo_handler.zocalo_interactor.run_start.call_count == len( - dc_ids - ) - - callbacks.zocalo_handler.zocalo_interactor.run_end.assert_has_calls( - [call(x) for x in dc_ids] - ) - assert callbacks.zocalo_handler.zocalo_interactor.run_end.call_count == len(dc_ids) - - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_not_called() +class TestXrayCentreZocaloHandler: + def test_execution_of_run_gridscan_triggers_zocalo_calls( + self, + mock_ispyb_update_time_and_status: MagicMock, + mock_ispyb_get_time: MagicMock, + mock_ispyb_store_grid_scan: MagicMock, + nexus_writer: MagicMock, + dummy_params, + ): + dc_ids = [1, 2] + dcg_id = 4 + + mock_ispyb_store_grid_scan.return_value = [dc_ids, None, dcg_id] + mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING + mock_ispyb_update_time_and_status.return_value = None + + callbacks = XrayCentreCallbackCollection.setup() + init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks, dc_ids, dcg_id) + callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) + callbacks.zocalo_handler.start(td.test_run_gridscan_start_document) + callbacks.ispyb_handler.descriptor( + td.test_descriptor_document_pre_data_collection + ) + callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) + callbacks.ispyb_handler.descriptor( + td.test_descriptor_document_during_data_collection + ) + callbacks.ispyb_handler.event(td.test_event_document_during_data_collection) + callbacks.zocalo_handler.start(td.test_do_fgs_start_document) + callbacks.ispyb_handler.stop(td.test_stop_document) + callbacks.zocalo_handler.stop(td.test_stop_document) + callbacks.zocalo_handler.zocalo_interactor.run_start.assert_has_calls( + [call(x) for x in dc_ids] + ) + assert callbacks.zocalo_handler.zocalo_interactor.run_start.call_count == len( + dc_ids + ) -@patch( - "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", - autospec=True, -) -def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_called( - store_3d_grid_scan, - dummy_params: GridscanInternalParameters, -): - callbacks = XrayCentreCallbackCollection.setup() - callbacks.ispyb_handler.descriptor(td.test_descriptor_document_pre_data_collection) - callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) + callbacks.zocalo_handler.zocalo_interactor.run_end.assert_has_calls( + [call(x) for x in dc_ids] + ) + assert callbacks.zocalo_handler.zocalo_interactor.run_end.call_count == len( + dc_ids + ) - callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) - callbacks.ispyb_handler.descriptor( - td.test_descriptor_document_during_data_collection - ) - callbacks.ispyb_handler.event(td.test_event_document_during_data_collection) - - mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) - expected_centre_grid_coords = np.array([1, 2, 3]) - single_crystal_result = [ - { - "max_voxel": [1, 2, 3], - "centre_of_mass": expected_centre_grid_coords, - "bounding_box": [[1, 1, 1], [2, 2, 2]], - "total_count": 192512.0, - } - ] - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( - single_crystal_result - ) - results = callbacks.zocalo_handler.wait_for_results(np.array([0, 0, 0])) + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_not_called() - found_centre = results[0] - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( - 100 + @patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + autospec=True, ) - expected_centre_motor_coords = ( - dummy_params.experiment_params.grid_position_to_motor_position( - expected_centre_grid_coords - 0.5 + def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_called( + self, + store_3d_grid_scan, + dummy_params: GridscanInternalParameters, + ): + callbacks = XrayCentreCallbackCollection.setup() + init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) + callbacks.ispyb_handler.descriptor( + td.test_descriptor_document_pre_data_collection ) - ) - np.testing.assert_array_equal(found_centre, expected_centre_motor_coords) - + callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) -@patch( - "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", - autospec=True, -) -def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_THEN_fallback_centre_used( - store_3d_grid_scan, - dummy_params, -): - callbacks = XrayCentreCallbackCollection.setup() - mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.side_effect = ( - NoDiffractionFound() - ) + callbacks.ispyb_handler.start(td.test_run_gridscan_start_document) + callbacks.ispyb_handler.descriptor( + td.test_descriptor_document_during_data_collection + ) + callbacks.ispyb_handler.event(td.test_event_document_during_data_collection) + + callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) + expected_centre_grid_coords = np.array([1, 2, 3]) + single_crystal_result = [ + { + "max_voxel": [1, 2, 3], + "centre_of_mass": expected_centre_grid_coords, + "bounding_box": [[1, 1, 1], [2, 2, 2]], + "total_count": 192512.0, + } + ] + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( + single_crystal_result + ) + results = callbacks.zocalo_handler.wait_for_results(np.array([0, 0, 0])) - fallback_position = np.array([1, 2, 3]) + found_centre = results[0] + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( + 100 + ) + expected_centre_motor_coords = ( + dummy_params.experiment_params.grid_position_to_motor_position( + expected_centre_grid_coords - 0.5 + ) + ) + np.testing.assert_array_equal(found_centre, expected_centre_motor_coords) - found_centre = callbacks.zocalo_handler.wait_for_results(fallback_position)[0] - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( - 100 + @patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + autospec=True, ) - np.testing.assert_array_equal(found_centre, fallback_position) - - -@patch( - "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", - autospec=True, -) -def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_exception( - store_3d_grid_scan, - dummy_params, -): - callbacks = XrayCentreCallbackCollection.setup() - mock_zocalo_functions(callbacks) + def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_THEN_fallback_centre_used( + self, + store_3d_grid_scan, + dummy_params, + ): + callbacks = XrayCentreCallbackCollection.setup() + init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) + callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.side_effect = ( + NoDiffractionFound() + ) - with pytest.raises(ISPyBDepositionNotMade): - callbacks.zocalo_handler.start(td.test_do_fgs_start_document) + fallback_position = np.array([1, 2, 3]) + found_centre = callbacks.zocalo_handler.wait_for_results(fallback_position)[0] + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( + 100 + ) + np.testing.assert_array_equal(found_centre, fallback_position) -@patch( - "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", - autospec=True, -) -def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_bbox_from_first( - store_3d_grid_scan, - dummy_params: GridscanInternalParameters, -): - callbacks = XrayCentreCallbackCollection.setup() - mock_zocalo_functions(callbacks) - callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) - expected_centre_grid_coords = np.array([4, 6, 2]) - multi_crystal_result = [ - { - "max_voxel": [1, 2, 3], - "centre_of_mass": np.array([3, 11, 11]), - "bounding_box": [[1, 1, 1], [3, 3, 3]], - "n_voxels": 2, - "total_count": 192512.0, - }, - { - "max_voxel": [1, 2, 3], - "centre_of_mass": expected_centre_grid_coords, - "bounding_box": [[2, 2, 2], [8, 8, 7]], - "n_voxels": 65, - "total_count": 6671044.0, - }, - ] - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( - multi_crystal_result + @patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + autospec=True, ) - found_centre, found_bbox = callbacks.zocalo_handler.wait_for_results( - np.array([0, 0, 0]) + def test_GIVEN_ispyb_not_started_WHEN_trigger_zocalo_handler_THEN_raises_exception( + self, + store_3d_grid_scan, + dummy_params, + ): + callbacks = XrayCentreCallbackCollection.setup() + init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) + + with pytest.raises(ISPyBDepositionNotMade): + callbacks.zocalo_handler.start(td.test_do_fgs_start_document) + + @patch( + "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", + autospec=True, ) - callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( - 100 - ) - expected_centre_motor_coords = ( - dummy_params.experiment_params.grid_position_to_motor_position( - np.array( - [ - expected_centre_grid_coords[0] - 0.5, - expected_centre_grid_coords[1] - 0.5, - expected_centre_grid_coords[2] - 0.5, - ] + def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_bbox_from_first( + self, + store_3d_grid_scan, + dummy_params: GridscanInternalParameters, + ): + callbacks = XrayCentreCallbackCollection.setup() + init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) + callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) + expected_centre_grid_coords = np.array([4, 6, 2]) + multi_crystal_result = [ + { + "max_voxel": [1, 2, 3], + "centre_of_mass": np.array([3, 11, 11]), + "bounding_box": [[1, 1, 1], [3, 3, 3]], + "n_voxels": 2, + "total_count": 192512.0, + }, + { + "max_voxel": [1, 2, 3], + "centre_of_mass": expected_centre_grid_coords, + "bounding_box": [[2, 2, 2], [8, 8, 7]], + "n_voxels": 65, + "total_count": 6671044.0, + }, + ] + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.return_value = ( + multi_crystal_result + ) + found_centre, found_bbox = callbacks.zocalo_handler.wait_for_results( + np.array([0, 0, 0]) + ) + callbacks.zocalo_handler.zocalo_interactor.wait_for_result.assert_called_once_with( + 100 + ) + expected_centre_motor_coords = ( + dummy_params.experiment_params.grid_position_to_motor_position( + np.array( + [ + expected_centre_grid_coords[0] - 0.5, + expected_centre_grid_coords[1] - 0.5, + expected_centre_grid_coords[2] - 0.5, + ] + ) ) ) - ) - np.testing.assert_array_equal(found_centre, expected_centre_motor_coords) - assert isinstance(found_bbox, np.ndarray) - expected_bbox_size = np.array([8, 8, 7]) - np.array([2, 2, 2]) - np.testing.assert_array_equal(found_bbox, expected_bbox_size) + np.testing.assert_array_equal(found_centre, expected_centre_motor_coords) + assert isinstance(found_bbox, np.ndarray) + expected_bbox_size = np.array([8, 8, 7]) - np.array([2, 2, 2]) + np.testing.assert_array_equal(found_bbox, expected_bbox_size) diff --git a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py index ee7474d1c..dc622bc76 100755 --- a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py +++ b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py @@ -48,7 +48,7 @@ def __init__(self, ispyb_config: str, experiment_type: str) -> None: self.experiment_type = experiment_type @abstractmethod - def _store_scan_data(self, conn: Connector): + def _store_scan_data(self, conn: Connector) -> tuple: pass @abstractmethod @@ -73,6 +73,7 @@ def append_to_comment( self, data_collection_id: int, comment: str, delimiter: str = " " ) -> None: with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + assert conn is not None, "Failed to connect to ISPyB!" mx_acquisition: MXAcquisition = conn.mx_acquisition mx_acquisition.update_data_collection_append_comments( data_collection_id, comment, delimiter @@ -83,6 +84,9 @@ def get_current_time_string(self): return now.strftime("%Y-%m-%d %H:%M:%S") def get_visit_string(self): + assert ( + self.ispyb_params and self.detector_params + ), "StoreInISPyB didn't acquire params" visit_path_match = self.get_visit_string_from_path(self.ispyb_params.visit_path) if visit_path_match: return visit_path_match @@ -101,12 +105,13 @@ def update_scan_with_end_time_and_status( data_collection_id: int, data_collection_group_id: int, ) -> None: - assert self.ispyb_params is not None - assert self.detector_params is not None + assert self.ispyb_params is not None and self.detector_params is not None if reason is not None and reason != "": self.append_to_comment(data_collection_id, f"{run_status} reason: {reason}") with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + assert conn is not None, "Failed to connect to ISPyB!" + mx_acquisition: MXAcquisition = conn.mx_acquisition params = mx_acquisition.get_data_collection_params() @@ -178,9 +183,11 @@ def _store_data_collection_group_table(self, conn: Connector) -> int: def _store_data_collection_table( self, conn: Connector, data_collection_group_id: int ) -> int: - assert self.ispyb_params is not None - assert self.detector_params is not None - assert self.xtal_snapshots is not None + assert ( + self.ispyb_params is not None + and self.detector_params is not None + and self.xtal_snapshots is not None + ) core: Core = conn.core mx_acquisition: MXAcquisition = conn.mx_acquisition @@ -294,6 +301,7 @@ def _store_scan_data(self, conn: Connector): def begin_deposition(self): with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + assert conn is not None, "Failed to connect to ISPyB" return self._store_scan_data(conn) def end_deposition(self, success: str, reason: str): @@ -311,7 +319,7 @@ def __init__( self, ispyb_config: str, experiment_type: str, - parameters: GridscanInternalParameters = None, + parameters: GridscanInternalParameters | None = None, ) -> None: super().__init__(ispyb_config, experiment_type) self.full_params: GridscanInternalParameters | None = parameters @@ -319,10 +327,13 @@ def __init__( self.data_collection_ids: tuple[int, ...] | None = None self.upper_left: list[int] | None = None self.y_steps: int | None = None - self.y_step_size: int | None = None + self.y_step_size: float | None = None self.grid_ids: tuple[int, ...] | None = None def begin_deposition(self): + assert ( + self.full_params is not None + ), "StoreGridscanInIspyb failed to get parameters." ( self.data_collection_ids, self.grid_ids, @@ -352,12 +363,13 @@ def store_grid_scan(self, full_params: GridscanInternalParameters): self.y_step_size = full_params.experiment_params.y_step_size with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: + assert conn is not None, "Failed to connect to ISPyB" return self._store_scan_data(conn) def _mutate_data_collection_params_for_experiment( self, params: dict[str, Any] ) -> dict[str, Any]: - assert self.full_params is not None + assert self.full_params and self.y_steps params["axis_range"] = 0 params["axis_end"] = self.omega_start params["n_images"] = self.full_params.experiment_params.x_steps * self.y_steps @@ -389,13 +401,16 @@ def _store_grid_info_table( return mx_acquisition.upsert_dc_grid(list(params.values())) def _construct_comment(self) -> str: - assert self.ispyb_params is not None - assert self.full_params is not None - assert self.upper_left is not None - assert self.y_step_size is not None + assert ( + self.ispyb_params is not None + and self.full_params is not None + and self.upper_left is not None + and self.y_step_size is not None + and self.y_steps is not None + ), "StoreGridScanInIspyb failed to get parameters" bottom_right = oav_utils.bottom_right_from_top_left( - self.upper_left, + self.upper_left, # type: ignore self.full_params.experiment_params.x_steps, self.y_steps, self.full_params.experiment_params.x_step_size, @@ -446,6 +461,12 @@ def _store_scan_data(self, conn: Connector): ) def __prepare_second_scan_params(self): + assert ( + self.omega_start is not None + and self.run_number is not None + and self.ispyb_params is not None + and self.full_params is not None + ), "StoreGridscanInIspyb failed to get parameters" self.omega_start += 90 self.run_number += 1 self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_end From 23b586b65a087636c2b86c9106c8e271bce05a1f Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 7 Nov 2023 10:13:16 +0000 Subject: [PATCH 05/15] fix linting --- .../callbacks/rotation/tests/test_rotation_callbacks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index c9e0eb750..30d5318cc 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -16,7 +16,6 @@ from hyperion.external_interaction.callbacks.rotation.callback_collection import ( RotationCallbackCollection, ) -from hyperion.external_interaction.callbacks.xray_centre.tests.conftest import TestData from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade from hyperion.external_interaction.ispyb.store_in_ispyb import ( StoreInIspyb, From c38cdf11767f0b87fccda39802562f0af088146e Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 7 Nov 2023 10:16:30 +0000 Subject: [PATCH 06/15] delete broken tests --- src/hyperion/system_tests/test_main_system.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/src/hyperion/system_tests/test_main_system.py b/src/hyperion/system_tests/test_main_system.py index 5a296420c..7efe608c1 100644 --- a/src/hyperion/system_tests/test_main_system.py +++ b/src/hyperion/system_tests/test_main_system.py @@ -267,15 +267,6 @@ def test_when_started_n_returnstatus_interrupted_bc_RE_aborted_thn_error_reptd( assert response_json["exception_type"] == "Exception" -def test_given_started_when_RE_stops_on_its_own_happily_then_no_error_reported( - test_env: ClientAndRunEngine, -): - test_env.client.put(START_ENDPOINT, data=TEST_PARAMS) - test_env.mock_run_engine.RE_takes_time = False - response_json = wait_for_run_engine_status(test_env.client) - assert response_json["status"] == Status.IDLE.value - - def test_start_with_json_file_gives_success(test_env: ClientAndRunEngine): test_env.mock_run_engine.RE_takes_time = False From 42e7d8834642c3e3f0af2407ac6d2df3db6d4873 Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 9 Nov 2023 19:25:59 +0000 Subject: [PATCH 07/15] zocalo callbacks get parameters themselves --- .../callbacks/rotation/zocalo_callback.py | 12 ++++++++---- .../callbacks/xray_centre/zocalo_callback.py | 13 +++++++++---- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py index bd7ca92c6..c7928f753 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py @@ -9,6 +9,9 @@ from hyperion.external_interaction.zocalo.zocalo_interaction import ZocaloInteractor from hyperion.log import LOGGER from hyperion.parameters.constants import ROTATION_OUTER_PLAN +from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( + RotationInternalParameters, +) class RotationZocaloCallback(CallbackBase): @@ -29,10 +32,11 @@ def start(self, doc: dict): LOGGER.info( "Zocalo callback recieved start document with experiment parameters." ) - assert ( - self.ispyb.params is not None - ), "ISPyB handler attached to Zocalo handler did not recieve parameters" - zocalo_environment = self.ispyb.params.hyperion_params.zocalo_environment + params = RotationInternalParameters.from_json( + doc.get("hyperion_internal_parameters") + ) + zocalo_environment = params.hyperion_params.zocalo_environment + LOGGER.info(f"Zocalo environment set to {zocalo_environment}.") self.zocalo_interactor = ZocaloInteractor(zocalo_environment) if self.run_uid is None: self.run_uid = doc.get("uid") diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py index 9d46e451f..ffe53f66d 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py @@ -17,6 +17,9 @@ ) from hyperion.log import LOGGER from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN +from hyperion.parameters.plan_specific.gridscan_internal_params import ( + GridscanInternalParameters, +) class XrayCentreZocaloCallback(CallbackBase): @@ -53,15 +56,17 @@ def start(self, doc: dict): LOGGER.info( "Zocalo callback recieved start document with experiment parameters." ) - assert ( - self.ispyb.params is not None - ), "ISPyB handler attached to Zocalo handler did not recieve parameters" + params = GridscanInternalParameters.from_json( + doc.get("hyperion_internal_parameters") + ) + zocalo_environment = params.hyperion_params.zocalo_environment + LOGGER.info(f"Zocalo environment set to {zocalo_environment}.") self.zocalo_interactor = ZocaloInteractor( self.ispyb.params.hyperion_params.zocalo_environment ) self.grid_position_to_motor_position: Callable[ [ndarray], ndarray - ] = self.ispyb.params.experiment_params.grid_position_to_motor_position + ] = params.experiment_params.grid_position_to_motor_position LOGGER.info("Zocalo handler received start document.") if doc.get("subplan_name") == "do_fgs": self.do_fgs_uid = doc.get("uid") From 86149a7b3905ed01180d1472e41bb3380075acc7 Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 10 Nov 2023 10:01:54 +0000 Subject: [PATCH 08/15] add IspybIds class --- .../callbacks/ispyb_callback_base.py | 14 ++++--------- .../callbacks/rotation/ispyb_callback.py | 13 +++++++----- .../callbacks/xray_centre/ispyb_callback.py | 10 ++++++---- .../ispyb/store_in_ispyb.py | 20 +++++++++++++++---- 4 files changed, 34 insertions(+), 23 deletions(-) diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index cde6b1ee3..c1248852e 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -5,7 +5,7 @@ from bluesky.callbacks import CallbackBase -from hyperion.external_interaction.ispyb.store_in_ispyb import StoreInIspyb +from hyperion.external_interaction.ispyb.store_in_ispyb import IspybIds, StoreInIspyb from hyperion.log import LOGGER, set_dcgid_tag from hyperion.parameters.constants import ( ISPYB_HARDWARE_READ_PLAN, @@ -20,11 +20,7 @@ ) if TYPE_CHECKING: - from hyperion.external_interaction.ispyb.store_in_ispyb import ( - Store2DGridscanInIspyb, - Store3DGridscanInIspyb, - StoreRotationInIspyb, - ) + from hyperion.external_interaction.ispyb.store_in_ispyb import StoreInIspyb class BaseISPyBCallback(CallbackBase): @@ -35,7 +31,7 @@ def __init__(self) -> None: self.params: GridscanInternalParameters | RotationInternalParameters | None = ( None ) - self.ispyb: StoreRotationInIspyb | Store3DGridscanInIspyb | Store2DGridscanInIspyb + self.ispyb: StoreInIspyb self.descriptors: Dict[str, dict] = {} self.ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG) if self.ispyb_config == SIM_ISPYB_CONFIG: @@ -44,9 +40,7 @@ def __init__(self) -> None: " set the ISPYB_CONFIG_PATH environment variable." ) self.uid_to_finalize_on: Optional[str] = None - self.ispyb_ids: tuple[int, int] | tuple[int, int, int] | tuple[ - None, None, None - ] = (None, None, None) + self.ispyb_ids: IspybIds = IspybIds() def _append_to_comment(self, id: int, comment: str): assert isinstance(self.ispyb, StoreInIspyb) diff --git a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py index 0c15a03ef..ce815a0d6 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py @@ -3,7 +3,10 @@ from hyperion.external_interaction.callbacks.ispyb_callback_base import ( BaseISPyBCallback, ) -from hyperion.external_interaction.ispyb.store_in_ispyb import StoreRotationInIspyb +from hyperion.external_interaction.ispyb.store_in_ispyb import ( + IspybIds, + StoreRotationInIspyb, +) from hyperion.log import LOGGER, set_dcgid_tag from hyperion.parameters.constants import ROTATION_OUTER_PLAN, ROTATION_PLAN_MAIN from hyperion.parameters.plan_specific.rotation_scan_internal_params import ( @@ -29,8 +32,8 @@ class RotationISPyBCallback(BaseISPyBCallback): """ def append_to_comment(self, comment: str): - assert self.ispyb_ids[0] is not None - self._append_to_comment(self.ispyb_ids[0], comment) + assert isinstance(self.ispyb_ids.data_collection_ids, int) + self._append_to_comment(self.ispyb_ids.data_collection_ids, comment) def start(self, doc: dict): if doc.get("subplan_name") == ROTATION_OUTER_PLAN: @@ -43,14 +46,14 @@ def start(self, doc: dict): self.ispyb: StoreRotationInIspyb = StoreRotationInIspyb( self.ispyb_config, self.params ) - self.ispyb_ids: tuple[int, int] | tuple[None, None] = (None, None) + self.ispyb_ids: IspybIds = IspybIds() LOGGER.info("ISPYB handler received start document.") if doc.get("subplan_name") == ROTATION_PLAN_MAIN: self.uid_to_finalize_on = doc.get("uid") def event(self, doc: dict): super().event(doc) - set_dcgid_tag(self.ispyb_ids[1]) + set_dcgid_tag(self.ispyb_ids.data_collection_group_id) def stop(self, doc: dict): if doc.get("run_start") == self.uid_to_finalize_on: 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 3c3a8526b..d998e65af 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -5,6 +5,7 @@ ) from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade from hyperion.external_interaction.ispyb.store_in_ispyb import ( + IspybIds, Store2DGridscanInIspyb, Store3DGridscanInIspyb, StoreGridscanInIspyb, @@ -36,7 +37,7 @@ def __init__(self) -> None: super().__init__() self.params: GridscanInternalParameters self.ispyb: StoreGridscanInIspyb - self.ispyb_ids: tuple = (None, None, None) + self.ispyb_ids: IspybIds = IspybIds() def start(self, doc: dict): if doc.get("subplan_name") == "run_gridscan_move_and_tidy": @@ -54,15 +55,16 @@ def start(self, doc: dict): self.uid_to_finalize_on = doc.get("uid") def append_to_comment(self, comment: str): - for id in self.ispyb_ids[0]: + assert isinstance(self.ispyb_ids.data_collection_ids, tuple) + for id in self.ispyb_ids.data_collection_ids: self._append_to_comment(id, comment) def event(self, doc: dict): super().event(doc) - set_dcgid_tag(self.ispyb_ids[2]) + set_dcgid_tag(self.ispyb_ids.data_collection_group_id) def stop(self, doc: dict): if doc.get("run_start") == self.uid_to_finalize_on: - if self.ispyb_ids == (None, None, None): + if self.ispyb_ids == IspybIds(): raise ISPyBDepositionNotMade("ispyb was not initialised at run start") super().stop(doc) diff --git a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py index dc622bc76..791870696 100755 --- a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py +++ b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py @@ -12,6 +12,7 @@ from ispyb.connector.mysqlsp.main import ISPyBMySQLSPConnector as Connector from ispyb.sp.core import Core from ispyb.sp.mxacquisition import MXAcquisition +from pydantic import BaseModel from hyperion.external_interaction.ispyb.ispyb_dataclass import ( GridscanIspybParams, @@ -35,6 +36,12 @@ VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})(/?$)" +class IspybIds(BaseModel): + data_collection_ids: int | tuple[int, ...] | None = None + data_collection_group_id: int | None = None + grid_id: int | None = None + + class StoreInIspyb(ABC): def __init__(self, ispyb_config: str, experiment_type: str) -> None: self.ispyb_params: IspybParams | None = None @@ -62,7 +69,7 @@ def _mutate_data_collection_params_for_experiment( pass @abstractmethod - def begin_deposition(self, success: str, reason: str): + def begin_deposition(self) -> IspybIds: pass @abstractmethod @@ -299,10 +306,11 @@ def _store_scan_data(self, conn: Connector): return data_collection_id, data_collection_group_id - def begin_deposition(self): + def begin_deposition(self) -> IspybIds: with ispyb.open(self.ISPYB_CONFIG_PATH) as conn: assert conn is not None, "Failed to connect to ISPyB" - return self._store_scan_data(conn) + ids = self._store_scan_data(conn) + return IspybIds(data_collection_ids=ids[0], data_collection_group_id=ids[1]) def end_deposition(self, success: str, reason: str): assert ( @@ -339,7 +347,11 @@ def begin_deposition(self): self.grid_ids, self.data_collection_group_id, ) = self.store_grid_scan(self.full_params) - return self.data_collection_ids, self.grid_ids, self.data_collection_group_id + return IspybIds( + data_collection_ids=self.data_collection_ids, + data_collection_group_id=self.data_collection_group_id, + grid_id=self.grid_ids, + ) def end_deposition(self, success: str, reason: str): assert ( From 7a594dd36443f1576c1ca3184b9b638db86e9ff2 Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 10 Nov 2023 11:01:32 +0000 Subject: [PATCH 09/15] #372 fix dataclass and tests --- .../experiment_plans/tests/conftest.py | 13 ++++++--- .../rotation/tests/test_rotation_callbacks.py | 27 ++++++++++++------- .../callbacks/rotation/zocalo_callback.py | 9 ++++--- .../xray_centre/tests/test_zocalo_handler.py | 19 +++++++++---- .../callbacks/xray_centre/zocalo_callback.py | 19 +++++++------ .../ispyb/store_in_ispyb.py | 4 +-- .../unit_tests/test_store_in_ispyb.py | 7 ++--- 7 files changed, 65 insertions(+), 33 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index 994acfec5..1a2fa995e 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/src/hyperion/experiment_plans/tests/conftest.py @@ -30,7 +30,10 @@ 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 @@ -296,7 +299,9 @@ def modified_interactor_mock(assign_run_end: Callable | None = None): def modified_store_grid_scan_mock(*args, dcids=(0, 0), dcgid=0, **kwargs): mock = MagicMock(spec=Store3DGridscanInIspyb) - mock.begin_deposition.return_value = (dcids, dcgid, 0) + mock.begin_deposition.return_value = IspybIds( + data_collection_ids=dcids, data_collection_group_id=dcgid, grid_ids=(0, 0) + ) return mock @@ -320,7 +325,9 @@ def mock_subscriptions(test_fgs_params): } ) subscriptions.ispyb_handler.ispyb = MagicMock(spec=Store3DGridscanInIspyb) - subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: [[0, 0], 0, 0] + subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: IspybIds( + data_collection_ids=(0, 0), data_collection_group_id=0, grid_ids=(0, 0) + ) return subscriptions diff --git a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index 30d5318cc..a739c3fc0 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -18,6 +18,7 @@ ) from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade from hyperion.external_interaction.ispyb.store_in_ispyb import ( + IspybIds, StoreInIspyb, StoreRotationInIspyb, ) @@ -199,9 +200,11 @@ def test_zocalo_start_and_end_triggered_once( cb.ispyb_handler.params = params def set_ispyb_ids(cbs): - cbs.ispyb_handler.ispyb_ids = (0, 0) + cbs.ispyb_handler.ispyb_ids = IspybIds( + data_collection_ids=0, data_collection_group_id=0 + ) - RE(fake_rotation_scan(params, cb, after_open_do=set_ispyb_ids)) + RE(fake_rotation_scan(params, cb, after_main_do=set_ispyb_ids)) zocalo.assert_called_once() cb.zocalo_handler.zocalo_interactor.run_start.assert_called_once() @@ -248,14 +251,17 @@ def test_zocalo_starts_on_opening_and_ispyb_on_main_so_ispyb_triggered_before_zo cb.ispyb_handler.start(test_start_doc) cb.zocalo_handler.start(test_start_doc) cb.ispyb_handler.ispyb = MagicMock(spec=StoreInIspyb) - cb.ispyb_handler.ispyb_ids = (0, 0) + cb.zocalo_handler.zocalo_interactor.run_start = MagicMock() cb.zocalo_handler.zocalo_interactor.run_end = MagicMock() - def after_open_assert(callbacks: RotationCallbackCollection): + def after_open_do(callbacks: RotationCallbackCollection): callbacks.ispyb_handler.ispyb.begin_deposition.assert_not_called() - def after_main_assert(callbacks: RotationCallbackCollection): + def after_main_do(callbacks: RotationCallbackCollection): + cb.ispyb_handler.ispyb_ids = IspybIds( + data_collection_ids=0, data_collection_group_id=0 + ) callbacks.ispyb_handler.ispyb.begin_deposition.assert_called_once() cb.zocalo_handler.zocalo_interactor.run_end.assert_not_called() @@ -263,7 +269,7 @@ def after_main_assert(callbacks: RotationCallbackCollection): "hyperion.external_interaction.callbacks.rotation.ispyb_callback.StoreRotationInIspyb", autospec=True, ): - RE(fake_rotation_scan(params, cb, after_open_assert, after_main_assert)) + RE(fake_rotation_scan(params, cb, after_open_do, after_main_do)) cb.zocalo_handler.zocalo_interactor.run_end.assert_called_once() @@ -281,11 +287,14 @@ def test_ispyb_handler_grabs_uid_from_main_plan_and_not_first_start_doc( autospec=True, side_effect=cb.ispyb_handler.start ) - def after_open_assert(callbacks: RotationCallbackCollection): + def after_open_do(callbacks: RotationCallbackCollection): callbacks.ispyb_handler.start.assert_called_once() assert callbacks.ispyb_handler.uid_to_finalize_on is None - def after_main_assert(callbacks: RotationCallbackCollection): + def after_main_do(callbacks: RotationCallbackCollection): + cb.ispyb_handler.ispyb_ids = IspybIds( + data_collection_ids=0, data_collection_group_id=0 + ) assert callbacks.ispyb_handler.start.call_count == 2 assert callbacks.ispyb_handler.uid_to_finalize_on is not None @@ -293,4 +302,4 @@ def after_main_assert(callbacks: RotationCallbackCollection): "hyperion.external_interaction.callbacks.rotation.ispyb_callback.StoreRotationInIspyb", autospec=True, ): - RE(fake_rotation_scan(params, cb, after_open_assert, after_main_assert)) + RE(fake_rotation_scan(params, cb, after_open_do, after_main_do)) diff --git a/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py index c7928f753..0692bb5fc 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/zocalo_callback.py @@ -46,8 +46,11 @@ def stop(self, doc: dict): LOGGER.info( f"Zocalo handler received stop document, for run {doc.get('run_start')}." ) - if self.ispyb.ispyb_ids[0] is not None: - self.zocalo_interactor.run_start(self.ispyb.ispyb_ids[0]) - self.zocalo_interactor.run_end(self.ispyb.ispyb_ids[0]) + if self.ispyb.ispyb_ids.data_collection_ids is not None: + assert isinstance(self.ispyb.ispyb_ids.data_collection_ids, int) + self.zocalo_interactor.run_start( + self.ispyb.ispyb_ids.data_collection_ids + ) + self.zocalo_interactor.run_end(self.ispyb.ispyb_ids.data_collection_ids) else: raise ISPyBDepositionNotMade("ISPyB deposition was not initialised!") diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py index 2af74c434..72998f5b3 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py @@ -9,6 +9,7 @@ ) from hyperion.external_interaction.callbacks.xray_centre.tests.conftest import TestData from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade +from hyperion.external_interaction.ispyb.store_in_ispyb import IspybIds from hyperion.external_interaction.zocalo.zocalo_interaction import NoDiffractionFound from hyperion.parameters.external_parameters import from_file as default_raw_params from hyperion.parameters.plan_specific.gridscan_internal_params import ( @@ -54,10 +55,12 @@ def test_execution_of_run_gridscan_triggers_zocalo_calls( nexus_writer: MagicMock, dummy_params, ): - dc_ids = [1, 2] + dc_ids = (1, 2) dcg_id = 4 - mock_ispyb_store_grid_scan.return_value = [dc_ids, None, dcg_id] + mock_ispyb_store_grid_scan.return_value = IspybIds( + data_collection_ids=dc_ids, grid_ids=None, data_collection_group_id=dcg_id + ) mock_ispyb_get_time.return_value = td.DUMMY_TIME_STRING mock_ispyb_update_time_and_status.return_value = None @@ -115,7 +118,9 @@ def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_cal ) callbacks.ispyb_handler.event(td.test_event_document_during_data_collection) - callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) + callbacks.ispyb_handler.ispyb_ids = IspybIds( + data_collection_ids=(0, 0), data_collection_group_id=100, grid_ids=(0, 0) + ) expected_centre_grid_coords = np.array([1, 2, 3]) single_crystal_result = [ { @@ -152,7 +157,9 @@ def test_GIVEN_no_results_from_zocalo_WHEN_communicator_wait_for_results_called_ ): callbacks = XrayCentreCallbackCollection.setup() init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) - callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) + callbacks.ispyb_handler.ispyb_ids = IspybIds( + data_collection_ids=(0, 0), data_collection_group_id=100, grid_ids=(0, 0) + ) callbacks.zocalo_handler.zocalo_interactor.wait_for_result.side_effect = ( NoDiffractionFound() ) @@ -191,7 +198,9 @@ def test_multiple_results_from_zocalo_sorted_by_total_count_returns_centre_and_b ): callbacks = XrayCentreCallbackCollection.setup() init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) - callbacks.ispyb_handler.ispyb_ids = ([0], 0, 100) + callbacks.ispyb_handler.ispyb_ids = IspybIds( + data_collection_ids=(0, 0), data_collection_group_id=100, grid_ids=(0, 0) + ) expected_centre_grid_coords = np.array([4, 6, 2]) multi_crystal_result = [ { diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py index ffe53f66d..e078e2d1a 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/zocalo_callback.py @@ -11,6 +11,7 @@ GridscanISPyBCallback, ) from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade +from hyperion.external_interaction.ispyb.store_in_ispyb import IspybIds from hyperion.external_interaction.zocalo.zocalo_interaction import ( NoDiffractionFound, ZocaloInteractor, @@ -70,9 +71,9 @@ def start(self, doc: dict): LOGGER.info("Zocalo handler received start document.") if doc.get("subplan_name") == "do_fgs": self.do_fgs_uid = doc.get("uid") - if self.ispyb.ispyb_ids[0] is not None: - datacollection_ids = self.ispyb.ispyb_ids[0] - for id in datacollection_ids: + if self.ispyb.ispyb_ids.data_collection_ids is not None: + assert isinstance(self.ispyb.ispyb_ids.data_collection_ids, tuple) + for id in self.ispyb.ispyb_ids.data_collection_ids: self.zocalo_interactor.run_start(id) else: raise ISPyBDepositionNotMade("ISPyB deposition was not initialised!") @@ -82,10 +83,10 @@ def stop(self, doc: dict): LOGGER.info( f"Zocalo handler received stop document, for run {doc.get('run_start')}." ) - if self.ispyb.ispyb_ids == (None, None, None): + if self.ispyb.ispyb_ids == IspybIds(): raise ISPyBDepositionNotMade("ISPyB deposition was not initialised!") - datacollection_ids = self.ispyb.ispyb_ids[0] - for id in datacollection_ids: + assert isinstance(self.ispyb.ispyb_ids.data_collection_ids, tuple) + for id in self.ispyb.ispyb_ids.data_collection_ids: self.zocalo_interactor.run_end(id) self.processing_start_time = time.time() @@ -98,11 +99,13 @@ def wait_for_results(self, fallback_xyz: ndarray) -> tuple[ndarray, Optional[lis Returns: ndarray: The xray centre position to move to """ - datacollection_group_id = self.ispyb.ispyb_ids[2] + assert ( + self.ispyb.ispyb_ids.data_collection_group_id is not None + ), "ISPyB deposition was not initialised!" try: raw_results = self.zocalo_interactor.wait_for_result( - datacollection_group_id + self.ispyb.ispyb_ids.data_collection_group_id ) # Sort from strongest to weakest in case of multiple crystals diff --git a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py index 791870696..040e1a9f8 100755 --- a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py +++ b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py @@ -39,7 +39,7 @@ class IspybIds(BaseModel): data_collection_ids: int | tuple[int, ...] | None = None data_collection_group_id: int | None = None - grid_id: int | None = None + grid_ids: tuple[int, ...] | None = None class StoreInIspyb(ABC): @@ -350,7 +350,7 @@ def begin_deposition(self): return IspybIds( data_collection_ids=self.data_collection_ids, data_collection_group_id=self.data_collection_group_id, - grid_id=self.grid_ids, + grid_ids=self.grid_ids, ) def end_deposition(self, success: str, reason: str): diff --git a/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py b/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py index 386ec4075..45c333cda 100644 --- a/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py +++ b/src/hyperion/external_interaction/unit_tests/test_store_in_ispyb.py @@ -8,6 +8,7 @@ from mockito import mock, when from hyperion.external_interaction.ispyb.store_in_ispyb import ( + IspybIds, Store2DGridscanInIspyb, Store3DGridscanInIspyb, StoreRotationInIspyb, @@ -228,9 +229,9 @@ def test_store_rotation_scan( TEST_DATA_COLLECTION_GROUP_ID, ) - assert dummy_rotation_ispyb.begin_deposition() == ( - TEST_DATA_COLLECTION_IDS[0], - TEST_DATA_COLLECTION_GROUP_ID, + assert dummy_rotation_ispyb.begin_deposition() == IspybIds( + data_collection_ids=TEST_DATA_COLLECTION_IDS[0], + data_collection_group_id=TEST_DATA_COLLECTION_GROUP_ID, ) From 5f8b6c6e0a2fe8f0ef02abd0c3be951472de9380 Mon Sep 17 00:00:00 2001 From: David Perl Date: Fri, 10 Nov 2023 11:10:45 +0000 Subject: [PATCH 10/15] #950 tidy up ispyb callback variables --- .../callbacks/ispyb_callback_base.py | 9 +++------ .../callbacks/rotation/ispyb_callback.py | 1 - .../callbacks/xray_centre/ispyb_callback.py | 1 - 3 files changed, 3 insertions(+), 8 deletions(-) diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index c1248852e..6c02854f7 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -61,13 +61,10 @@ def event(self, doc: dict): hyperion.log""" LOGGER.debug("ISPyB handler received event document.") - assert isinstance( - self.ispyb, StoreInIspyb - ), "ISPyB deposition can't be initialised!" + assert self.ispyb is not None, "ISPyB deposition can't be initialised!" + assert self.params is not None, "ISPyB handler didn't recieve parameters!" + event_descriptor = self.descriptors[doc["descriptor"]] - assert isinstance(self.params, GridscanInternalParameters) or isinstance( - self.params, RotationInternalParameters - ), "ISPyB handler params set with wrong type" if event_descriptor.get("name") == ISPYB_HARDWARE_READ_PLAN: self.params.hyperion_params.ispyb_params.undulator_gap = doc["data"][ "undulator_gap" diff --git a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py index ce815a0d6..d7218b166 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py @@ -37,7 +37,6 @@ def append_to_comment(self, comment: str): def start(self, doc: dict): if doc.get("subplan_name") == ROTATION_OUTER_PLAN: - self.run_uid = doc.get("uid") LOGGER.info( "ISPyB callback recieved start document with experiment parameters." ) 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 d998e65af..15a60f335 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -51,7 +51,6 @@ def start(self, doc: dict): if self.params.experiment_params.is_3d_grid_scan else Store2DGridscanInIspyb(self.ispyb_config, self.params) ) - self.run_start_uid = doc.get("uid") self.uid_to_finalize_on = doc.get("uid") def append_to_comment(self, comment: str): From 4b16ad5ec25746f61920ba59b97cfa49b662e2df Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 14 Nov 2023 09:36:53 +0000 Subject: [PATCH 11/15] #949 fix more typing in StoreInIspyb --- .../experiment_plans/tests/conftest.py | 18 ++++------ .../callbacks/xray_centre/ispyb_callback.py | 5 +-- .../xray_centre/tests/test_ispyb_handler.py | 4 +-- .../ispyb/ispyb_dataclass.py | 6 ++-- .../ispyb/store_in_ispyb.py | 36 ++++++++++--------- 5 files changed, 33 insertions(+), 36 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index 1a2fa995e..ac634d0b0 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/src/hyperion/experiment_plans/tests/conftest.py @@ -312,18 +312,12 @@ def mock_subscriptions(test_fgs_params): modified_interactor_mock, ): subscriptions = XrayCentreCallbackCollection.setup() - subscriptions.ispyb_handler.start( - { - "subplan_name": GRIDSCAN_OUTER_PLAN, - "hyperion_internal_parameters": test_fgs_params.json(), - } - ) - subscriptions.zocalo_handler.start( - { - "subplan_name": GRIDSCAN_OUTER_PLAN, - "hyperion_internal_parameters": test_fgs_params.json(), - } - ) + start_doc = { + "subplan_name": GRIDSCAN_OUTER_PLAN, + "hyperion_internal_parameters": test_fgs_params.json(), + } + subscriptions.ispyb_handler.start(start_doc) + subscriptions.zocalo_handler.start(start_doc) subscriptions.ispyb_handler.ispyb = MagicMock(spec=Store3DGridscanInIspyb) subscriptions.ispyb_handler.ispyb.begin_deposition = lambda: IspybIds( data_collection_ids=(0, 0), data_collection_group_id=0, grid_ids=(0, 0) 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 15a60f335..d37fa447e 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -11,6 +11,7 @@ StoreGridscanInIspyb, ) from hyperion.log import LOGGER, set_dcgid_tag +from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN from hyperion.parameters.plan_specific.gridscan_internal_params import ( GridscanInternalParameters, ) @@ -40,9 +41,9 @@ def __init__(self) -> None: self.ispyb_ids: IspybIds = IspybIds() def start(self, doc: dict): - if doc.get("subplan_name") == "run_gridscan_move_and_tidy": + if doc.get("subplan_name") == GRIDSCAN_OUTER_PLAN: LOGGER.info( - "Nexus writer recieved start document with experiment parameters." + "ISPyB callback recieved start document with experiment parameters." ) json_params = doc.get("hyperion_internal_parameters") self.params = GridscanInternalParameters.from_json(json_params) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py index d42f11ee6..b54327968 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_ispyb_handler.py @@ -40,8 +40,8 @@ def dummy_params(): return GridscanInternalParameters(**default_raw_params()) -def mock_store_in_ispyb(dummy_params, *args, **kwargs) -> Store3DGridscanInIspyb: - mock = Store3DGridscanInIspyb(None, dummy_params) +def mock_store_in_ispyb(config, params, *args, **kwargs) -> Store3DGridscanInIspyb: + mock = Store3DGridscanInIspyb("", params) mock.store_grid_scan = MagicMock(return_value=[DC_IDS, None, DCG_ID]) mock.get_current_time_string = MagicMock(return_value=td.DUMMY_TIME_STRING) mock.update_scan_with_end_time_and_status = MagicMock(return_value=None) diff --git a/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py b/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py index 06eca874d..c098204d2 100644 --- a/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py +++ b/src/hyperion/external_interaction/ispyb/ispyb_dataclass.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional import numpy as np from pydantic import BaseModel, validator @@ -75,8 +75,8 @@ def _parse_position( synchrotron_mode: Optional[str] = None slit_gap_size_x: Optional[float] = None slit_gap_size_y: Optional[float] = None - xtal_snapshots_omega_start: Optional[List[str]] = None - xtal_snapshots_omega_end: Optional[List[str]] = None + xtal_snapshots_omega_start: Optional[list[str]] = None + xtal_snapshots_omega_end: Optional[list[str]] = None @validator("transmission_fraction") def _transmission_not_percentage(cls, transmission_fraction: float): diff --git a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py index 040e1a9f8..a06cd144c 100755 --- a/src/hyperion/external_interaction/ispyb/store_in_ispyb.py +++ b/src/hyperion/external_interaction/ispyb/store_in_ispyb.py @@ -12,6 +12,7 @@ from ispyb.connector.mysqlsp.main import ISPyBMySQLSPConnector as Connector from ispyb.sp.core import Core from ispyb.sp.mxacquisition import MXAcquisition +from numpy import ndarray from pydantic import BaseModel from hyperion.external_interaction.ispyb.ispyb_dataclass import ( @@ -44,15 +45,15 @@ class IspybIds(BaseModel): class StoreInIspyb(ABC): def __init__(self, ispyb_config: str, experiment_type: str) -> None: - self.ispyb_params: IspybParams | None = None - self.detector_params: DetectorParams | None = None - self.run_number: int | None = None - self.omega_start: float | None = None - self.experiment_type: str | None = None - self.xtal_snapshots: list[str] | None = None - self.data_collection_group_id: int | None = None self.ISPYB_CONFIG_PATH: str = ispyb_config self.experiment_type = experiment_type + self.ispyb_params: IspybParams + self.detector_params: DetectorParams + self.run_number: int + self.omega_start: float + self.experiment_type: str + self.xtal_snapshots: list[str] + self.data_collection_group_id: int @abstractmethod def _store_scan_data(self, conn: Connector) -> tuple: @@ -327,15 +328,16 @@ def __init__( self, ispyb_config: str, experiment_type: str, - parameters: GridscanInternalParameters | None = None, + parameters: GridscanInternalParameters, ) -> None: super().__init__(ispyb_config, experiment_type) - self.full_params: GridscanInternalParameters | None = parameters - self.ispyb_params: GridscanIspybParams | None = None + self.full_params: GridscanInternalParameters = parameters + self.ispyb_params: GridscanIspybParams = parameters.hyperion_params.ispyb_params + self.upper_left: list[int] | ndarray = self.ispyb_params.upper_left + self.y_steps: int = self.full_params.experiment_params.y_steps + self.y_step_size: float = self.full_params.experiment_params.y_step_size + self.omega_start = 0 self.data_collection_ids: tuple[int, ...] | None = None - self.upper_left: list[int] | None = None - self.y_steps: int | None = None - self.y_step_size: float | None = None self.grid_ids: tuple[int, ...] | None = None def begin_deposition(self): @@ -366,7 +368,7 @@ def store_grid_scan(self, full_params: GridscanInternalParameters): self.detector_params = full_params.hyperion_params.detector_params self.run_number = self.detector_params.run_number self.omega_start = self.detector_params.omega_start - self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_start + self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_start or [] self.upper_left = [ int(self.ispyb_params.upper_left[0]), int(self.ispyb_params.upper_left[1]), @@ -442,7 +444,7 @@ def _construct_comment(self) -> str: class Store3DGridscanInIspyb(StoreGridscanInIspyb): - def __init__(self, ispyb_config, parameters=None): + def __init__(self, ispyb_config: str, parameters: GridscanInternalParameters): super().__init__(ispyb_config, "Mesh3D", parameters) def _store_scan_data(self, conn: Connector): @@ -481,7 +483,7 @@ def __prepare_second_scan_params(self): ), "StoreGridscanInIspyb failed to get parameters" self.omega_start += 90 self.run_number += 1 - self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_end + self.xtal_snapshots = self.ispyb_params.xtal_snapshots_omega_end or [] self.upper_left = [ int(self.ispyb_params.upper_left[0]), int(self.ispyb_params.upper_left[2]), @@ -491,7 +493,7 @@ def __prepare_second_scan_params(self): class Store2DGridscanInIspyb(StoreGridscanInIspyb): - def __init__(self, ispyb_config, parameters=None): + def __init__(self, ispyb_config: str, parameters: GridscanInternalParameters): super().__init__(ispyb_config, "mesh", parameters) def _store_scan_data(self, conn: Connector): From e84f88c2a198dc5822640a62076765d784e9f706 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 14 Nov 2023 18:35:49 +0000 Subject: [PATCH 12/15] Merge branch 'main' into 950_init_callbacks_w_document --- src/hyperion/experiment_plans/tests/conftest.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index ae7c0b67b..712dfe3b9 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/src/hyperion/experiment_plans/tests/conftest.py @@ -338,10 +338,6 @@ def mock_subscriptions(test_fgs_params): data_collection_ids=(0, 0), data_collection_group_id=0, grid_ids=(0, 0) ) - subscriptions.ispyb_handler.active = True - subscriptions.nexus_handler.active = True - subscriptions.zocalo_handler.active = True - return subscriptions From 0635428dd5e29f611305ae197cb2db97df945fe4 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 21 Nov 2023 10:09:33 +0000 Subject: [PATCH 13/15] #950 fix some errors from merging --- .../experiment_plans/tests/conftest.py | 4 +-- .../tests/test_flyscan_xray_centre_plan.py | 36 +++++++++++-------- .../callbacks/ispyb_callback_base.py | 24 ++++++------- .../callbacks/xray_centre/ispyb_callback.py | 12 +++---- 4 files changed, 41 insertions(+), 35 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/conftest.py b/src/hyperion/experiment_plans/tests/conftest.py index d4083294c..4ea731e13 100644 --- a/src/hyperion/experiment_plans/tests/conftest.py +++ b/src/hyperion/experiment_plans/tests/conftest.py @@ -333,8 +333,8 @@ def mock_subscriptions(test_fgs_params): "subplan_name": GRIDSCAN_OUTER_PLAN, "hyperion_internal_parameters": test_fgs_params.json(), } - subscriptions.ispyb_handler.start(start_doc) - subscriptions.zocalo_handler.start(start_doc) + 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: IspybIds( data_collection_ids=(0, 0), data_collection_group_id=0, grid_ids=(0, 0) diff --git a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py index d45df93cb..c45e3f10f 100644 --- a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py +++ b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py @@ -42,7 +42,10 @@ GridscanISPyBCallback, ) from hyperion.external_interaction.callbacks.xray_centre.tests.conftest import TestData -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, TEST_RESULT_MEDIUM, @@ -70,10 +73,8 @@ def ispyb_plan(test_fgs_params): } ) def standalone_read_hardware_for_ispyb(und, syn, slits, attn, fl): - yield from bps.open_run() yield from read_hardware_for_ispyb_pre_collection(und, syn, slits) yield from read_hardware_for_ispyb_during_collection(attn, fl) - yield from bps.close_run() return standalone_read_hardware_for_ispyb @@ -138,6 +139,9 @@ def test_read_hardware_for_ispyb_updates_from_ophyd_devices( test_ispyb_callback = GridscanISPyBCallback() test_ispyb_callback.active = True test_ispyb_callback.ispyb = MagicMock(spec=Store3DGridscanInIspyb) + test_ispyb_callback.ispyb.begin_deposition.return_value = IspybIds( + data_collection_ids=(2, 3), data_collection_group_id=5, grid_ids=(7, 8, 9) + ) RE.subscribe(test_ispyb_callback) RE( @@ -186,13 +190,13 @@ def test_results_adjusted_and_passed_to_move_xyz( set_up_logging_handlers(logging_level="INFO", dev_mode=True) RE.subscribe(VerbosePlanExecutionLoggingCallback()) - mock_subscriptions.ispyb_handler.start( + mock_subscriptions.ispyb_handler.activity_gated_start( { "subplan_name": GRIDSCAN_OUTER_PLAN, "hyperion_internal_parameters": test_fgs_params.json(), } ) - mock_subscriptions.ispyb_handler.descriptor( + mock_subscriptions.ispyb_handler.activity_gated_descriptor( {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} ) mock_subscriptions.ispyb_handler.activity_gated_event( @@ -206,7 +210,7 @@ def test_results_adjusted_and_passed_to_move_xyz( }, } ) - mock_subscriptions.ispyb_handler.descriptor( + mock_subscriptions.ispyb_handler.activity_gated_descriptor( {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} ) mock_subscriptions.ispyb_handler.activity_gated_event( @@ -336,9 +340,9 @@ def test_individual_plans_triggered_once_and_only_once_in_composite_run( RE: RunEngine, ): td = TestData() - mock_subscriptions.ispyb_handler.start(td.test_start_document) - mock_subscriptions.zocalo_handler.start(td.test_start_document) - mock_subscriptions.ispyb_handler.descriptor( + mock_subscriptions.ispyb_handler.activity_gated_start(td.test_start_document) + mock_subscriptions.zocalo_handler.activity_gated_start(td.test_start_document) + mock_subscriptions.ispyb_handler.activity_gated_descriptor( {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} ) @@ -353,7 +357,7 @@ def test_individual_plans_triggered_once_and_only_once_in_composite_run( }, } ) - mock_subscriptions.ispyb_handler.descriptor( + mock_subscriptions.ispyb_handler.activity_gated_descriptor( {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} ) mock_subscriptions.ispyb_handler.activity_gated_event( @@ -381,7 +385,7 @@ def test_when_gridscan_finished_then_smargon_stub_offsets_are_set( test_fgs_params: GridscanInternalParameters, RE: RunEngine, ): - mock_subscriptions.ispyb_handler.descriptor( + mock_subscriptions.ispyb_handler.activity_gated_descriptor( {"uid": "123abc", "name": ISPYB_HARDWARE_READ_PLAN} ) @@ -396,7 +400,7 @@ def test_when_gridscan_finished_then_smargon_stub_offsets_are_set( }, } ) - mock_subscriptions.ispyb_handler.descriptor( + mock_subscriptions.ispyb_handler.activity_gated_descriptor( {"uid": "abc123", "name": ISPYB_TRANSMISSION_FLUX_READ_PLAN} ) mock_subscriptions.ispyb_handler.activity_gated_event( @@ -411,9 +415,11 @@ def test_when_gridscan_finished_then_smargon_stub_offsets_are_set( set_up_logging_handlers(logging_level="INFO", dev_mode=True) RE.subscribe(VerbosePlanExecutionLoggingCallback()) - mock_subscriptions.zocalo_handler.wait_for_results.return_value = ( - (0, 0, 0), - None, + mock_subscriptions.zocalo_handler.wait_for_results = MagicMock( + return_value=( + (0, 0, 0), + None, + ) ) RE( diff --git a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py index fa4242b3a..6c3895a3b 100644 --- a/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py +++ b/src/hyperion/external_interaction/callbacks/ispyb_callback_base.py @@ -44,22 +44,13 @@ def __init__(self) -> None: self.uid_to_finalize_on: Optional[str] = None self.ispyb_ids: IspybIds = IspybIds() - def _append_to_comment(self, id: int, comment: str): - assert isinstance(self.ispyb, StoreInIspyb) - try: - self.ispyb.append_to_comment(id, comment) - except TypeError: - ISPYB_LOGGER.warning( - "ISPyB deposition not initialised, can't update comment." - ) - - def activity_gated_descriptor(self, doc: dict): - self.descriptors[doc["uid"]] = doc - def activity_gated_start(self, doc: dict): if self.uid_to_finalize_on is None: self.uid_to_finalize_on = doc.get("uid") + def activity_gated_descriptor(self, doc: dict): + self.descriptors[doc["uid"]] = doc + def activity_gated_event(self, doc: dict): """Subclasses should extend this to add a call to set_dcig_tag from hyperion.log""" @@ -114,3 +105,12 @@ def activity_gated_stop(self, doc: dict): ISPYB_LOGGER.warning( f"Failed to finalise ISPyB deposition on stop document: {doc} with exception: {e}" ) + + def _append_to_comment(self, id: int, comment: str): + assert isinstance(self.ispyb, StoreInIspyb) + try: + self.ispyb.append_to_comment(id, comment) + except TypeError: + ISPYB_LOGGER.warning( + "ISPyB deposition not initialised, can't update comment." + ) 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 b42b162fa..cd74f5a2f 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py @@ -40,7 +40,7 @@ def __init__(self) -> None: self.ispyb: StoreGridscanInIspyb self.ispyb_ids: IspybIds = IspybIds() - def start(self, doc: dict): + def activity_gated_start(self, doc: dict): if doc.get("subplan_name") == GRIDSCAN_OUTER_PLAN: self.uid_to_finalize_on = doc.get("uid") ISPYB_LOGGER.info( @@ -55,11 +55,6 @@ def start(self, doc: dict): else Store2DGridscanInIspyb(self.ispyb_config, self.params) ) - def append_to_comment(self, comment: str): - assert isinstance(self.ispyb_ids.data_collection_ids, tuple) - for id in self.ispyb_ids.data_collection_ids: - self._append_to_comment(id, comment) - def activity_gated_event(self, doc: dict): super().activity_gated_event(doc) set_dcgid_tag(self.ispyb_ids.data_collection_group_id) @@ -73,3 +68,8 @@ def activity_gated_stop(self, doc: dict): if self.ispyb_ids == IspybIds(): raise ISPyBDepositionNotMade("ispyb was not initialised at run start") super().activity_gated_stop(doc) + + def append_to_comment(self, comment: str): + assert isinstance(self.ispyb_ids.data_collection_ids, tuple) + for id in self.ispyb_ids.data_collection_ids: + self._append_to_comment(id, comment) From 1ff94c67a8516f17c1be91ba56e5a5b3ab1306c0 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 21 Nov 2023 10:26:31 +0000 Subject: [PATCH 14/15] #950 finish cleaning up merge artefacts --- .../tests/test_grid_detection_plan.py | 2 +- .../rotation/tests/test_rotation_callbacks.py | 1 + .../xray_centre/tests/test_zocalo_handler.py | 14 +++++++++----- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/test_grid_detection_plan.py b/src/hyperion/experiment_plans/tests/test_grid_detection_plan.py index c7d5fc73f..243788189 100644 --- a/src/hyperion/experiment_plans/tests/test_grid_detection_plan.py +++ b/src/hyperion/experiment_plans/tests/test_grid_detection_plan.py @@ -46,7 +46,7 @@ def fake_devices(smargon: Smargon, backlight: Backlight): "dodal.devices.areadetector.plugins.MJPG.Image" ) as mock_image_class: mock_image = MagicMock() - mock_image_class.open.return_value = mock_image + mock_image_class.open.return_value.__enter__.return_value = mock_image composite = OavGridDetectionComposite( backlight=backlight, diff --git a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py index c82b99ff0..10b80908d 100644 --- a/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py +++ b/src/hyperion/external_interaction/callbacks/rotation/tests/test_rotation_callbacks.py @@ -152,6 +152,7 @@ def test_nexus_handler_only_writes_once( cb = RotationCallbackCollection.setup() activate_callbacks(cb) cb.ispyb_handler.activity_gated_start = MagicMock(autospec=True) + cb.ispyb_handler.activity_gated_event = MagicMock(autospec=True) cb.ispyb_handler.activity_gated_stop = MagicMock(autospec=True) RE(fake_rotation_scan(params, cb)) diff --git a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py index 7e41a7aff..c85f167d0 100644 --- a/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py +++ b/src/hyperion/external_interaction/callbacks/xray_centre/tests/test_zocalo_handler.py @@ -39,8 +39,8 @@ def init_cbs_with_docs_and_mock_zocalo_and_ispyb( "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb", lambda _, __: modified_store_grid_scan_mock(dcids=dcids, dcgid=dcgid), ): - callbacks.ispyb_handler.start(td.test_start_document) - callbacks.zocalo_handler.start(td.test_start_document) + callbacks.ispyb_handler.activity_gated_start(td.test_start_document) + callbacks.zocalo_handler.activity_gated_start(td.test_start_document) callbacks.zocalo_handler.zocalo_interactor.wait_for_result = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_end = MagicMock() callbacks.zocalo_handler.zocalo_interactor.run_start = MagicMock() @@ -67,7 +67,9 @@ def test_execution_of_run_gridscan_triggers_zocalo_calls( callbacks = XrayCentreCallbackCollection.setup() init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks, dc_ids, dcg_id) callbacks.ispyb_handler.activity_gated_start(td.test_run_gridscan_start_document) # type: ignore - callbacks.zocalo_handler.start(td.test_run_gridscan_start_document) + callbacks.zocalo_handler.activity_gated_start( + td.test_run_gridscan_start_document + ) callbacks.ispyb_handler.activity_gated_descriptor( td.test_descriptor_document_pre_data_collection ) # type: ignore @@ -111,10 +113,12 @@ def test_zocalo_called_to_wait_on_results_when_communicator_wait_for_results_cal ): callbacks = XrayCentreCallbackCollection.setup() init_cbs_with_docs_and_mock_zocalo_and_ispyb(callbacks) - callbacks.ispyb_handler.descriptor( + callbacks.ispyb_handler.activity_gated_descriptor( td.test_descriptor_document_pre_data_collection ) - callbacks.ispyb_handler.event(td.test_event_document_pre_data_collection) + callbacks.ispyb_handler.activity_gated_event( + td.test_event_document_pre_data_collection + ) callbacks.ispyb_handler.activity_gated_start(td.test_run_gridscan_start_document) # type: ignore callbacks.ispyb_handler.activity_gated_descriptor( From 2e50c4fa338424189cda45af0772ba6ebb525188 Mon Sep 17 00:00:00 2001 From: David Perl Date: Tue, 21 Nov 2023 10:32:34 +0000 Subject: [PATCH 15/15] #950 fix linting --- .../experiment_plans/tests/test_flyscan_xray_centre_plan.py | 1 - src/hyperion/parameters/tests/test_external_parameters.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py index c45e3f10f..31f4b43eb 100644 --- a/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py +++ b/src/hyperion/experiment_plans/tests/test_flyscan_xray_centre_plan.py @@ -1,7 +1,6 @@ import types from unittest.mock import MagicMock, call, patch -import bluesky.plan_stubs as bps import bluesky.preprocessors as bpp import numpy as np import pytest diff --git a/src/hyperion/parameters/tests/test_external_parameters.py b/src/hyperion/parameters/tests/test_external_parameters.py index 0906fe6f3..e26187e57 100644 --- a/src/hyperion/parameters/tests/test_external_parameters.py +++ b/src/hyperion/parameters/tests/test_external_parameters.py @@ -90,7 +90,7 @@ def test_parse_exception_causes_warning(mock_logger): def test_parse_list(): test_data = [([1, 2, 3], "[1, 2, 3]"), ([1, True, 3], "[1, Yes, 3]")] - for (expected, input) in test_data: + for expected, input in test_data: actual = GDABeamlineParameters.parse_value(input) assert expected == actual, f"Actual:{actual}, expected: {expected}\n"