(#1017) Move common parts out of store in ispyb and rename

src/hyperion/external_interaction/callbacks/ispyb_callback_base.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
-import os
 from typing import TYPE_CHECKING, Dict, Optional
 
 from hyperion.external_interaction.callbacks.plan_reactive_callback import (
     PlanReactiveCallback,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import IspybIds, StoreInIspyb
+from hyperion.external_interaction.ispyb.ispyb_utils import get_ispyb_config
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
+    IspybIds,
+    StoreInIspyb,
+)
 from hyperion.log import ISPYB_LOGGER, set_dcgid_tag
 from hyperion.parameters.constants import (
     ISPYB_HARDWARE_READ_PLAN,
@@ -23,7 +26,9 @@
 if TYPE_CHECKING:
     from event_model.documents.event import Event
 
-    from hyperion.external_interaction.ispyb.store_in_ispyb import StoreInIspyb
+    from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
+        StoreInIspyb,
+    )
 
 
 class BaseISPyBCallback(PlanReactiveCallback):
@@ -37,7 +42,7 @@ def __init__(self) -> None:
         )
         self.ispyb: StoreInIspyb
         self.descriptors: Dict[str, dict] = {}
-        self.ispyb_config = os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG)
+        self.ispyb_config = get_ispyb_config()
         if self.ispyb_config == SIM_ISPYB_CONFIG:
             ISPYB_LOGGER.warning(
                 "Using dev ISPyB database. If you want to use the real database, please"

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

@@ -3,7 +3,7 @@
 from hyperion.external_interaction.callbacks.ispyb_callback_base import (
     BaseISPyBCallback,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     IspybIds,
     StoreRotationInIspyb,
 )

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

@@ -10,7 +10,7 @@
     BaseISPyBCallback,
 )
 from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     IspybIds,
     Store2DGridscanInIspyb,
     Store3DGridscanInIspyb,

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

@@ -14,7 +14,7 @@
     GridscanISPyBCallback,
 )
 from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
-from hyperion.external_interaction.ispyb.store_in_ispyb import IspybIds
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import IspybIds
 from hyperion.log import ISPYB_LOGGER
 from hyperion.parameters.constants import GRIDSCAN_OUTER_PLAN
 from hyperion.parameters.plan_specific.gridscan_internal_params import (

src/hyperion/external_interaction/ispyb/ispyb_utils.py

@@ -0,0 +1,35 @@
+import datetime
+import os
+import re
+
+from ispyb import NoResult
+from ispyb.connector.mysqlsp.main import ISPyBMySQLSPConnector as Connector
+from ispyb.sp.core import Core
+
+from hyperion.parameters.constants import (
+    SIM_ISPYB_CONFIG,
+)
+
+VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})(/?$)"
+
+
+def get_current_time_string():
+    now = datetime.datetime.now()
+    return now.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def get_ispyb_config():
+    return os.environ.get("ISPYB_CONFIG_PATH", SIM_ISPYB_CONFIG)
+
+
+def get_visit_string_from_path(path) -> str | None:
+    match = re.search(VISIT_PATH_REGEX, path) if path else None
+    return str(match.group(1)) if match else None
+
+
+def get_session_id_from_visit(conn: Connector, visit: str):
+    try:
+        core: Core = conn.core
+        return core.retrieve_visit_id(visit)
+    except NoResult:
+        raise Exception(f"Not found - session ID for visit {visit}")

src/hyperion/external_interaction/ispyb/store_in_ispyb.py → src/hyperion/external_interaction/ispyb/store_datacollection_in_ispyb.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import datetime
-import re
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
@@ -10,7 +8,6 @@
 import ispyb.sqlalchemy
 from dodal.devices.detector import DetectorParams
 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
@@ -21,6 +18,11 @@
     Orientation,
     RotationIspybParams,
 )
+from hyperion.external_interaction.ispyb.ispyb_utils import (
+    get_current_time_string,
+    get_session_id_from_visit,
+    get_visit_string_from_path,
+)
 from hyperion.log import ISPYB_LOGGER
 from hyperion.tracing import TRACER
 
@@ -34,7 +36,6 @@
 
 I03_EIGER_DETECTOR = 78
 EIGER_FILE_SUFFIX = "h5"
-VISIT_PATH_REGEX = r".+/([a-zA-Z]{2}\d{4,5}-\d{1,3})(/?$)"
 
 
 class IspybIds(BaseModel):
@@ -87,23 +88,19 @@ def append_to_comment(
                 data_collection_id, comment, delimiter
             )
 
-    def get_current_time_string(self):
-        now = datetime.datetime.now()
-        return now.strftime("%Y-%m-%d %H:%M:%S")
-
-    def get_visit_string(self):
+    def get_visit_string(self) -> str:
         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)
+        visit_path_match = get_visit_string_from_path(self.ispyb_params.visit_path)
         if visit_path_match:
             return visit_path_match
-        else:
-            return self.get_visit_string_from_path(self.detector_params.directory)
-
-    def get_visit_string_from_path(self, path):
-        match = re.search(VISIT_PATH_REGEX, path) if path else None
-        return match.group(1) if match else None
+        visit_path_match = get_visit_string_from_path(self.detector_params.directory)
+        if not visit_path_match:
+            raise ValueError(
+                f"Visit not found from {self.ispyb_params.visit_path} or {self.detector_params.directory}"
+            )
+        return visit_path_match
 
     def update_scan_with_end_time_and_status(
         self,
@@ -143,7 +140,7 @@ def _end_deposition(self, dcid: int, success: str, reason: str):
             run_status = "DataCollection Unsuccessful"
         else:
             run_status = "DataCollection Successful"
-        current_time = self.get_current_time_string()
+        current_time = get_current_time_string()
         assert self.data_collection_group_id is not None
         self.update_scan_with_end_time_and_status(
             current_time,
@@ -169,18 +166,10 @@ def _store_position_table(self, conn: Connector, dc_id: int) -> int:
 
     def _store_data_collection_group_table(self, conn: Connector) -> int:
         assert self.ispyb_params is not None
-        core: Core = conn.core
         mx_acquisition: MXAcquisition = conn.mx_acquisition
 
-        try:
-            session_id = core.retrieve_visit_id(self.get_visit_string())
-        except ispyb.NoResult:
-            raise Exception(
-                f"Not found - session ID for visit {self.get_visit_string()} where self.ispyb_params.visit_path is {self.ispyb_params.visit_path}"
-            )
-
         params = mx_acquisition.get_data_collection_group_params()
-        params["parentid"] = session_id
+        params["parentid"] = get_session_id_from_visit(conn, self.get_visit_string())
         params["experimenttype"] = self.experiment_type
         params["sampleid"] = self.ispyb_params.sample_id
         params["sample_barcode"] = self.ispyb_params.sample_barcode
@@ -197,18 +186,11 @@ def _store_data_collection_table(
             and self.xtal_snapshots is not None
         )
 
-        core: Core = conn.core
         mx_acquisition: MXAcquisition = conn.mx_acquisition
-        try:
-            session_id = core.retrieve_visit_id(self.get_visit_string())
-        except ispyb.NoResult:
-            raise Exception(
-                f"Not found - session ID for visit {self.get_visit_string()}"
-            )
 
         params = mx_acquisition.get_data_collection_params()
 
-        params["visitid"] = session_id
+        params["visitid"] = get_session_id_from_visit(conn, self.get_visit_string())
         params["parentid"] = data_collection_group_id
         params["sampleid"] = self.ispyb_params.sample_id
         params["detectorid"] = I03_EIGER_DETECTOR
@@ -251,7 +233,7 @@ def _store_data_collection_table(
             ) = self.xtal_snapshots
         params["synchrotron_mode"] = self.ispyb_params.synchrotron_mode
         params["undulator_gap1"] = self.ispyb_params.undulator_gap
-        params["starttime"] = self.get_current_time_string()
+        params["starttime"] = get_current_time_string()
 
         # temporary file template until nxs filewriting is integrated and we can use
         # that file name

tests/system_tests/external_interaction/conftest.py

@@ -10,7 +10,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     Store2DGridscanInIspyb,
     Store3DGridscanInIspyb,
 )

tests/system_tests/external_interaction/test_ispyb_dev_connection.py

@@ -11,7 +11,7 @@
 from hyperion.external_interaction.callbacks.rotation.callback_collection import (
     RotationCallbackCollection,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     Store2DGridscanInIspyb,
     Store3DGridscanInIspyb,
     StoreGridscanInIspyb,

tests/unit_tests/experiment_plans/conftest.py

@@ -19,7 +19,7 @@
 from hyperion.external_interaction.callbacks.xray_centre.ispyb_callback import (
     GridscanISPyBCallback,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     IspybIds,
     Store3DGridscanInIspyb,
 )

tests/unit_tests/experiment_plans/test_flyscan_xray_centre_plan.py

@@ -37,7 +37,7 @@
 from hyperion.external_interaction.callbacks.xray_centre.ispyb_callback import (
     GridscanISPyBCallback,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     IspybIds,
     Store3DGridscanInIspyb,
 )

tests/unit_tests/external_interaction/callbacks/test_rotation_callbacks.py

@@ -20,7 +20,7 @@
     XrayCentreCallbackCollection,
 )
 from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
-from hyperion.external_interaction.ispyb.store_in_ispyb import (
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
     IspybIds,
     StoreInIspyb,
     StoreRotationInIspyb,

tests/unit_tests/external_interaction/callbacks/xray_centre/conftest.py

@@ -29,7 +29,7 @@ def nexus_writer():
 @pytest.fixture
 def mock_ispyb_get_time():
     with patch(
-        "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb.get_current_time_string"
+        "hyperion.external_interaction.ispyb.ispyb_utils.get_current_time_string"
     ) as p:
         yield p
 

tests/unit_tests/external_interaction/callbacks/xray_centre/test_ispyb_handler.py

@@ -7,7 +7,9 @@
 from hyperion.external_interaction.callbacks.xray_centre.ispyb_callback import (
     GridscanISPyBCallback,
 )
-from hyperion.external_interaction.ispyb.store_in_ispyb import Store3DGridscanInIspyb
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import (
+    Store3DGridscanInIspyb,
+)
 from hyperion.log import (
     ISPYB_LOGGER,
     dc_group_id_filter,
@@ -41,11 +43,14 @@ def mock_emit():
 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)
     return mock
 
 
+@patch(
+    "hyperion.external_interaction.ispyb.store_datacollection_in_ispyb.get_current_time_string",
+    MagicMock(return_value=td.DUMMY_TIME_STRING),
+)
 @patch(
     "hyperion.external_interaction.callbacks.xray_centre.ispyb_callback.Store3DGridscanInIspyb",
     mock_store_in_ispyb,

tests/unit_tests/external_interaction/callbacks/xray_centre/test_zocalo_handler.py

@@ -6,7 +6,7 @@
     XrayCentreCallbackCollection,
 )
 from hyperion.external_interaction.exceptions import ISPyBDepositionNotMade
-from hyperion.external_interaction.ispyb.store_in_ispyb import IspybIds
+from hyperion.external_interaction.ispyb.store_datacollection_in_ispyb import IspybIds
 from hyperion.parameters.external_parameters import from_file as default_raw_params
 from hyperion.parameters.plan_specific.gridscan_internal_params import (
     GridscanInternalParameters,

tests/unit_tests/external_interaction/test_ispyb_utils.py

@@ -0,0 +1,35 @@
+import re
+
+import pytest
+
+from hyperion.external_interaction.ispyb.ispyb_utils import (
+    get_current_time_string,
+    get_visit_string_from_path,
+)
+
+TIME_FORMAT_REGEX = r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}"
+
+
+def test_get_current_time_string():
+    current_time = get_current_time_string()
+
+    assert isinstance(current_time, str)
+    assert re.match(TIME_FORMAT_REGEX, current_time) is not None
+
+
+@pytest.mark.parametrize(
+    "visit_path, expected_match",
+    [
+        ("/dls/i03/data/2022/cm6477-45/", "cm6477-45"),
+        ("/dls/i03/data/2022/cm6477-45", "cm6477-45"),
+        ("/dls/i03/data/2022/mx54663-1/", "mx54663-1"),
+        ("/dls/i03/data/2022/mx54663-1", "mx54663-1"),
+        ("/dls/i03/data/2022/mx53-1/", None),
+        ("/dls/i03/data/2022/mx53-1", None),
+        ("/dls/i03/data/2022/mx5563-1565/", None),
+        ("/dls/i03/data/2022/mx5563-1565", None),
+    ],
+)
+def test_find_visit_in_visit_path(visit_path: str, expected_match: str):
+    test_visit_path = get_visit_string_from_path(visit_path)
+    assert test_visit_path == expected_match