diff --git a/src/dodal/devices/zocalo/__init__.py b/src/dodal/devices/zocalo/__init__.py index edd2579d2e..3644917422 100644 --- a/src/dodal/devices/zocalo/__init__.py +++ b/src/dodal/devices/zocalo/__init__.py @@ -1,6 +1,4 @@ -from dodal.devices.zocalo.zocalo_interaction import ( - ZocaloTrigger, -) +from dodal.devices.zocalo.zocalo_interaction import ZocaloStartInfo, ZocaloTrigger from dodal.devices.zocalo.zocalo_results import ( NoResultsFromZocalo, NoZocaloSubscription, @@ -17,6 +15,7 @@ "ZOCALO_READING_PLAN_NAME", "NoResultsFromZocalo", "NoZocaloSubscription", + "ZocaloStartInfo", ] ZOCALO_READING_PLAN_NAME = "zocalo reading" diff --git a/src/dodal/devices/zocalo/zocalo_interaction.py b/src/dodal/devices/zocalo/zocalo_interaction.py index 217da22d41..48303e5f7d 100644 --- a/src/dodal/devices/zocalo/zocalo_interaction.py +++ b/src/dodal/devices/zocalo/zocalo_interaction.py @@ -1,5 +1,8 @@ +import dataclasses import getpass import socket +from dataclasses import dataclass +from typing import Optional import zocalo.configuration from workflows.transport import lookup @@ -16,6 +19,22 @@ def _get_zocalo_connection(environment): return transport +@dataclass +class ZocaloStartInfo: + """ + ispyb_dcid (int): The ID of the data collection in ISPyB + filename (str): The name of the file that the detector will store into dev/shm + number_of_frames (int): The number of frames in this collection. + start_index (int): The index of the first image of this collection within the file + written by the detector + """ + + ispyb_dcid: int + filename: Optional[str] + start_index: int + number_of_frames: int + + class ZocaloTrigger: """This class just sends 'run_start' and 'run_end' messages to zocalo, it is intended to be used in bluesky callback classes. To get results from zocalo back @@ -42,16 +61,20 @@ def _send_to_zocalo(self, parameters: dict): finally: transport.disconnect() - def run_start(self, data_collection_id: int): + def run_start( + self, + start_data: ZocaloStartInfo, + ): """Tells the data analysis pipeline we have started a run. Assumes that appropriate data has already been put into ISPyB Args: - data_collection_id (int): The ID of the data collection representing the - gridscan in ISPyB + start_data (ZocaloStartInfo): Data about the collection to send to zocalo """ - LOGGER.info(f"Starting Zocalo job with ispyb id {data_collection_id}") - self._send_to_zocalo({"event": "start", "ispyb_dcid": data_collection_id}) + LOGGER.info(f"Starting Zocalo job {start_data}") + data = dataclasses.asdict(start_data) + data["event"] = "start" + self._send_to_zocalo(data) def run_end(self, data_collection_id: int): """Tells the data analysis pipeline we have finished a run. diff --git a/tests/devices/system_tests/test_zocalo_results.py b/tests/devices/system_tests/test_zocalo_results.py index 18849f0c46..3b40d20980 100644 --- a/tests/devices/system_tests/test_zocalo_results.py +++ b/tests/devices/system_tests/test_zocalo_results.py @@ -39,7 +39,7 @@ async def zocalo_device(): async def test_read_results_from_fake_zocalo(zocalo_device: ZocaloResults): zocalo_device._subscribe_to_results() zc = ZocaloTrigger("dev_artemis") - zc.run_start(0) + zc.run_start(0, 0, 100) zc.run_end(0) zocalo_device.timeout_s = 5 @@ -66,7 +66,7 @@ async def test_stage_unstage_controls_read_results_from_fake_zocalo( def plan(): yield from bps.open_run() - zc.run_start(0) + zc.run_start(0, 0, 100) zc.run_end(0) yield from bps.sleep(0.15) yield from bps.trigger_and_read([zocalo_device]) diff --git a/tests/devices/unit_tests/test_zocalo_interaction.py b/tests/devices/unit_tests/test_zocalo_interaction.py index ea43abfafe..5767fcd1bf 100644 --- a/tests/devices/unit_tests/test_zocalo_interaction.py +++ b/tests/devices/unit_tests/test_zocalo_interaction.py @@ -9,11 +9,19 @@ from dodal.devices.zocalo import ( ZocaloTrigger, ) +from dodal.devices.zocalo.zocalo_interaction import ZocaloStartInfo SIM_ZOCALO_ENV = "dev_artemis" EXPECTED_DCID = 100 -EXPECTED_RUN_START_MESSAGE = {"event": "start", "ispyb_dcid": EXPECTED_DCID} +EXPECTED_FILENAME = "test/file" +EXPECTED_RUN_START_MESSAGE = { + "ispyb_dcid": EXPECTED_DCID, + "filename": EXPECTED_FILENAME, + "number_of_frames": 100, + "start_index": 0, + "event": "start", +} EXPECTED_RUN_END_MESSAGE = { "event": "end", "ispyb_dcid": EXPECTED_DCID, @@ -65,27 +73,40 @@ def with_exception(function_to_run, mock_transport): @mark.parametrize( - "function_to_test,function_wrapper,expected_message", + "function_wrapper,expected_message", [ - (zc.run_start, normally, EXPECTED_RUN_START_MESSAGE), + (normally, EXPECTED_RUN_START_MESSAGE), ( - zc.run_start, with_exception, EXPECTED_RUN_START_MESSAGE, ), - (zc.run_end, normally, EXPECTED_RUN_END_MESSAGE), - (zc.run_end, with_exception, EXPECTED_RUN_END_MESSAGE), ], ) -def test__run_start_and_end( - function_to_test: Callable, function_wrapper: Callable, expected_message: Dict -): +def test_run_start(function_wrapper: Callable, expected_message: Dict): + """ + Args: + function_wrapper (Callable): A wrapper used to test for expected exceptions + expected_message (Dict): The expected dictionary sent to zocalo + """ + data = ZocaloStartInfo(EXPECTED_DCID, EXPECTED_FILENAME, 0, 100) + function_to_run = partial(zc.run_start, data) + function_to_run = partial(function_wrapper, function_to_run) + _test_zocalo(function_to_run, expected_message) + + +@mark.parametrize( + "function_wrapper,expected_message", + [ + (normally, EXPECTED_RUN_END_MESSAGE), + (with_exception, EXPECTED_RUN_END_MESSAGE), + ], +) +def test__run_start_and_end(function_wrapper: Callable, expected_message: Dict): """ Args: - function_to_test (Callable): The function to test e.g. start/stop zocalo - function_wrapper (Callable): A wrapper around the function, used to test for expected exceptions + function_wrapper (Callable): A wrapper used to test for expected exceptions expected_message (Dict): The expected dictionary sent to zocalo """ - function_to_run = partial(function_to_test, EXPECTED_DCID) + function_to_run = partial(zc.run_end, EXPECTED_DCID) function_to_run = partial(function_wrapper, function_to_run) _test_zocalo(function_to_run, expected_message)