feat: refactor psi_detector_base class, add tests

ophyd_devices/interfaces/base_classes/psi_detector_base.py

@@ -1,3 +1,9 @@
+"""This module contains the base class for SLS detectors. We follow the approach to integrate
+PSI detectors into the BEC system based on this base class. The base class is used to implement
+certain methods that are expected by BEC, such as stage, unstage, trigger, stop, etc...
+We use composition with a custom prepare class to implement BL specific logic for the detector.
+The beamlines need to inherit from the CustomDetectorMixing for their mixin classes."""
+
 import os
 import time
 
@@ -32,48 +38,39 @@ class CustomDetectorMixin:
     def __init__(self, *_args, parent: Device = None, **_kwargs) -> None:
         self.parent = parent
 
-    def initialize_default_parameter(self) -> None:
+    def on_init(self) -> None:
         """
-        Init parameters for the detector
+        Init sequence for the detector
-
-        Raises (optional):
-            DetectorTimeoutError: if detector cannot be initialized
         """
 
-    def initialize_detector(self) -> None:
+    def on_stage(self) -> None:
         """
-        Init parameters for the detector
+        Specify actions to be executed during stage in preparation for a scan.
+        self.parent.scaninfo already has all current parameters for the upcoming scan.
 
-        Raises (optional):
+        In case the backend service is writing data on disk, this step should include publishing
-            DetectorTimeoutError: if detector cannot be initialized
+        a file_event and file_message to BEC to inform the system where the data is written to.
+
+        IMPORTANT:
+        It must be safe to assume that the device is ready for the scan
+        to start immediately once this function is finished.
         """
 
-    def initialize_detector_backend(self) -> None:
+    def on_unstage(self) -> None:
         """
-        Init parameters for teh detector backend (filewriter)
+        Specify actions to be executed during unstage.
 
-        Raises (optional):
+        This step should include checking if the acqusition was successful,
-            DetectorTimeoutError: if filewriter cannot be initialized
+        and publishing the file location and file event message,
+        with flagged done to BEC.
         """
 
-    def prepare_detector(self) -> None:
+    def on_stop(self) -> None:
-        """
-        Prepare detector for the scan
         """
+        Specify actions to be executed during stop.
+        This must also set self.parent.stopped to True.
 
-    def prepare_detector_backend(self) -> None:
+        This step should include stopping the detector and backend service.
-        """
-        Prepare detector backend for the scan
-        """
-
-    def stop_detector(self) -> None:
-        """
-        Stop the detector
-        """
-
-    def stop_detector_backend(self) -> None:
-        """
-        Stop the detector backend
         """
 
     def on_trigger(self) -> None:
@@ -81,57 +78,36 @@ class CustomDetectorMixin:
         Specify actions to be executed upon receiving trigger signal
         """
 
-    def pre_scan(self) -> None:
+    def on_pre_scan(self) -> None:
         """
-        Specify actions to be executed right before a scan
+        Specify actions to be executed right before a scan starts.
 
-        BEC calls pre_scan just before execution of the scan core.
+        Only use if needed, and it is recommended to keep this function as short/fast as possible.
-        It is convenient to execute time critical features of the detector,
-        e.g. arming it, but it is recommended to keep this function as short/fast as possible.
-        """
-
-    def finished(self) -> None:
-        """
-        Specify actions to be executed during unstage
-
-        This may include checks if acquisition was succesful
-
-        Raises (optional):
-            DetectorTimeoutError: if detector cannot be stopped
-        """
-
-    def check_scan_id(self) -> None:
-        """
-        Check if BEC is running on a new scan_id
-        """
-
-    def publish_file_location(self, done: bool = False, successful: bool = None) -> None:
-        """
-        Publish the designated filepath from data backend to REDIS.
-
-        Typically, the following two message types are published:
-
-        - file_event: event for the filewriter
-        - public_file: event for any secondary service (e.g. radial integ code)
         """
 
     def wait_for_signals(
         self,
-        signal_conditions: list,
+        signal_conditions: list[tuple],
         timeout: float,
         check_stopped: bool = False,
         interval: float = 0.05,
         all_signals: bool = False,
     ) -> bool:
-        """Wait for signals to reach a certain condition
+        """
+        Convenience wrapper to allow waiting for signals to reach a certain condition.
+        For EPICs PVs, an example usage is pasted at the bottom.
 
         Args:
-            signal_conditions (tuple): tuple of (get_current_state, condition) functions
+            signal_conditions (list[tuple]): tuple of executable calls for conditions (get_current_state, condition) to check
             timeout (float): timeout in seconds
             interval (float): interval in seconds
             all_signals (bool): True if all signals should be True, False if any signal should be True
+
         Returns:
             bool: True if all signals are in the desired state, False if timeout is reached
+
+            >>> Example usage for EPICS PVs:
+            >>> self.wait_for_signals(signal_conditions=[(self.acquiring.get, False)], timeout=5, interval=0.05, check_stopped=True, all_signals=True)
         """
         timer = 0
         while True:
@@ -155,7 +131,6 @@ class PSIDetectorBase(Device):
 
     Class attributes:
         custom_prepare_cls (object): class for custom prepare logic (BL specific)
-        Min_readout (float): minimum readout time for detector
 
     Args:
         prefix (str): EPICS PV prefix for component (optional)
@@ -165,69 +140,33 @@ class PSIDetectorBase(Device):
                     normal -> readout for read
                     config -> config parameter for 'ophydobj.read_configuration()'
                     hinted -> which attribute is readout for read
-        read_attrs (list): sequence of attribute names to read
-        configuration_attrs (list): sequence of attribute names via config_parameters
         parent (object): instance of the parent device
         device_manager (object): bec device manager
-        sim_mode (bool): simulation mode, if True, no device manager is required
         **kwargs: keyword arguments
-
-        attributes: lazy_wait_for_connection : bool
     """
 
     custom_prepare_cls = CustomDetectorMixin
 
-    MIN_READOUT = 1e-3
+    def __init__(self, prefix="", *, name, kind=None, parent=None, device_manager=None, **kwargs):
-
+        super().__init__(prefix=prefix, name=name, kind=kind, parent=parent, **kwargs)
-    # Specify which functions are revealed to the user in BEC client
-    USER_ACCESS = ["describe"]
-
-    def __init__(
-        self,
-        prefix="",
-        *,
-        name,
-        kind=None,
-        read_attrs=None,
-        configuration_attrs=None,
-        parent=None,
-        device_manager=None,
-        sim_mode=False,
-        **kwargs,
-    ):
-        super().__init__(
-            prefix=prefix,
-            name=name,
-            kind=kind,
-            read_attrs=read_attrs,
-            configuration_attrs=configuration_attrs,
-            parent=parent,
-            **kwargs,
-        )
-        if device_manager is None and not sim_mode:
-            raise DetectorInitError(
-                f"No device manager for device: {name}, and not started sim_mode: {sim_mode}. Add"
-                " DeviceManager to initialization or init with sim_mode=True"
-            )
-        # Init variables
-        self.sim_mode = sim_mode
         self.stopped = False
         self.name = name
         self.service_cfg = None
         self.scaninfo = None
         self.filewriter = None
-        self.timeout = 5
-        self.wait_for_connection(all_signals=True)
 
-        # Init custom prepare class with BL specific logic
+        if not issubclass(self.custom_prepare_cls, CustomDetectorMixin):
+            raise DetectorInitError("Custom prepare class must be subclass of CustomDetectorMixin")
         self.custom_prepare = self.custom_prepare_cls(parent=self, **kwargs)
-        if not sim_mode:
+
+        if device_manager:
             self._update_service_config()
             self.device_manager = device_manager
         else:
             self.device_manager = bec_utils.DMMock()
-            base_path = kwargs["basepath"] if "basepath" in kwargs else "~/Data10/"
+            base_path = kwargs["basepath"] if "basepath" in kwargs else "."
-            self.service_cfg = {"base_path": os.path.expanduser(base_path)}
+            self.service_cfg = {"base_path": os.path.abspath(base_path)}
+
         self.connector = self.device_manager.connector
         self._update_scaninfo()
         self._update_filewriter()
@@ -241,51 +180,54 @@ class PSIDetectorBase(Device):
         """Update scaninfo from BecScaninfoMixing
         This depends on device manager and operation/sim_mode
         """
-        self.scaninfo = BecScaninfoMixin(self.device_manager, self.sim_mode)
+        self.scaninfo = BecScaninfoMixin(self.device_manager)
         self.scaninfo.load_scan_metadata()
 
     def _update_service_config(self) -> None:
         """Update service config from BEC service config"""
+        # pylint: disable=import-outside-toplevel
         from bec_lib.bec_service import SERVICE_CONFIG
 
         self.service_cfg = SERVICE_CONFIG.config["service_config"]["file_writer"]
 
+    def check_scan_id(self) -> None:
+        """Checks if scan_id has changed and set stopped flagged to True if it has."""
+        old_scan_id = self.scaninfo.scan_id
+        self.scaninfo.load_scan_metadata()
+        if self.scaninfo.scan_id != old_scan_id:
+            self.stopped = True
+
     def _init(self) -> None:
         """Initialize detector, filewriter and set default parameters"""
-        self.custom_prepare.initialize_default_parameter()
+        self.custom_prepare.on_init()
-        self.custom_prepare.initialize_detector()
-        self.custom_prepare.initialize_detector_backend()
 
     def stage(self) -> list[object]:
         """
-         Stage device in preparation for a scan
+        Stage device in preparation for a scan.
-
+        First we check if the device is already staged. Stage is idempotent,
-        Internal Calls:
+        if staged twice it should raise (we let ophyd.Device handle the raise here).
-        - _prep_backend           : prepare detector filewriter for measurement
+        We reset the stopped flag and get the scaninfo from BEC, before calling custom_prepare.on_stage.
-        - _prep_detector              : prepare detector for measurement
 
         Returns:
             list(object): list of objects that were staged
 
         """
-        # Method idempotent, should rais ;obj;'RedudantStaging' if staged twice
         if self._staged != Staged.no:
             return super().stage()
-
-        # Reset flag for detector stopped
         self.stopped = False
-        # Load metadata of the scan
         self.scaninfo.load_scan_metadata()
-        # Prepare detector and file writer
+        self.custom_prepare.on_stage()
-        self.custom_prepare.prepare_detector_backend()
-        self.custom_prepare.prepare_detector()
-        state = False
-        self.custom_prepare.publish_file_location(done=state)
-        # At the moment needed bc signal might not be reliable, BEC too fast.
-        # Consider removing this overhead in future!
-        time.sleep(0.05)
         return super().stage()
 
+    def pre_scan(self) -> None:
+        """Pre-scan logic.
+
+        This function will be called from BEC directly before the scan core starts, and should only implement
+        time-critical actions. Therefore, it should also be kept as short/fast as possible.
+        I.e. Arming a detector in case there is a risk of timing out.
+        """
+        self.custom_prepare.on_pre_scan()
+
     def trigger(self) -> DeviceStatus:
         """Trigger the detector, called from BEC."""
         self.custom_prepare.on_trigger()
@@ -293,26 +235,20 @@ class PSIDetectorBase(Device):
 
     def unstage(self) -> list[object]:
         """
-        Unstage device in preparation for a scan
+        Unstage device after a scan.
 
-        Returns directly if self.stopped,
+        We first check if the scanID has changed, thus, the scan was unexpectedly interrupted but the device was not stopped.
-        otherwise checks with self._finished
+        If that is the case, the stopped flag is set to True, which will immediately unstage the device.
-        if data acquisition on device finished (an was successful)
 
-        Internal Calls:
+        Custom_prepare.on_unstage is called to allow for BL specific logic to be executed.
-        - custom_prepare.check_scan_id          : check if scan_id changed or detector stopped
-        - custom_prepare.finished              : check if device finished acquisition (succesfully)
-        - custom_prepare.publish_file_location : publish file location to bec
 
         Returns:
             list(object): list of objects that were unstaged
         """
-        self.custom_prepare.check_scan_id()
+        self.check_scan_id()
         if self.stopped is True:
             return super().unstage()
-        self.custom_prepare.finished()
+        self.custom_prepare.on_unstage()
-        state = True
-        self.custom_prepare.publish_file_location(done=state, successful=state)
         self.stopped = False
         return super().unstage()
 
@@ -320,11 +256,7 @@ class PSIDetectorBase(Device):
         """
         Stop the scan, with camera and file writer
 
-        Internal Calls:
-        - custom_prepare.stop_detector     : stop detector
-        - custom_prepare.stop_backend : stop detector filewriter
         """
-        self.custom_prepare.stop_detector()
+        self.custom_prepare.on_stop()
-        self.custom_prepare.stop_detector_backend()
         super().stop(success=success)
         self.stopped = True

ophyd_devices/utils/bec_scaninfo_mixin.py

@@ -4,6 +4,8 @@ from bec_lib import bec_logger, messages
 from bec_lib.devicemanager import DeviceManagerBase
 from bec_lib.endpoints import MessageEndpoints
 
+from ophyd_devices.utils.bec_utils import DMMock
+
 logger = bec_logger.logger
 
 
@@ -63,11 +65,9 @@ class BecScaninfoMixin:
         BecScaninfoMixin: BecScaninfoMixin object
     """
 
-    def __init__(
+    def __init__(self, device_manager: DeviceManagerBase = None, bec_info_msg=None) -> None:
-        self, device_manager: DeviceManagerBase = None, sim_mode: bool = False, bec_info_msg=None
+        self.sim_mode = bool(isinstance(device_manager, DMMock))
-    ) -> None:
         self.device_manager = device_manager
-        self.sim_mode = sim_mode
         self.scan_msg = None
         self.scan_id = None
         if bec_info_msg is None:

tests/test_base_classes.py

@@ -0,0 +1,90 @@
+# pylint: skip-file
+from unittest import mock
+
+import pytest
+from ophyd import DeviceStatus, Staged
+from ophyd.utils.errors import RedundantStaging
+
+from ophyd_devices.interfaces.base_classes.psi_detector_base import (
+    CustomDetectorMixin,
+    PSIDetectorBase,
+)
+from ophyd_devices.utils.bec_scaninfo_mixin import BecScaninfoMixin
+
+
+@pytest.fixture
+def detector_base():
+    yield PSIDetectorBase(name="test_detector")
+
+
+def test_detector_base_init(detector_base):
+    assert detector_base.stopped is False
+    assert detector_base.name == "test_detector"
+    assert "base_path" in detector_base.filewriter.service_config
+    assert isinstance(detector_base.scaninfo, BecScaninfoMixin)
+    assert issubclass(detector_base.custom_prepare_cls, CustomDetectorMixin)
+
+
+def test_stage(detector_base):
+    detector_base._staged = Staged.yes
+    with pytest.raises(RedundantStaging):
+        detector_base.stage()
+    assert detector_base.stopped is False
+    detector_base._staged = Staged.no
+    with (
+        mock.patch.object(detector_base.custom_prepare, "on_stage") as mock_on_stage,
+        mock.patch.object(detector_base.scaninfo, "load_scan_metadata") as mock_load_metadata,
+    ):
+        rtr = detector_base.stage()
+        assert isinstance(rtr, list)
+        mock_on_stage.assert_called_once()
+        mock_load_metadata.assert_called_once()
+        assert detector_base.stopped is False
+
+
+def test_pre_scan(detector_base):
+    with mock.patch.object(detector_base.custom_prepare, "on_pre_scan") as mock_on_pre_scan:
+        detector_base.pre_scan()
+        mock_on_pre_scan.assert_called_once()
+
+
+def test_trigger(detector_base):
+    with mock.patch.object(detector_base.custom_prepare, "on_trigger") as mock_on_trigger:
+        rtr = detector_base.trigger()
+        assert isinstance(rtr, DeviceStatus)
+        mock_on_trigger.assert_called_once()
+
+
+def test_unstage(detector_base):
+    detector_base.stopped = True
+    with (
+        mock.patch.object(detector_base.custom_prepare, "on_unstage") as mock_on_unstage,
+        mock.patch.object(detector_base, "check_scan_id") as mock_check_scan_id,
+    ):
+        rtr = detector_base.unstage()
+        assert isinstance(rtr, list)
+        assert mock_check_scan_id.call_count == 1
+        mock_on_unstage.assert_not_called()
+        detector_base.stopped = False
+        rtr = detector_base.unstage()
+        assert isinstance(rtr, list)
+        assert mock_check_scan_id.call_count == 2
+        assert detector_base.stopped is False
+        mock_on_unstage.assert_called_once()
+
+
+def test_stop(detector_base):
+    with mock.patch.object(detector_base.custom_prepare, "on_stop") as mock_on_stop:
+        detector_base.stop()
+        mock_on_stop.assert_called_once()
+        assert detector_base.stopped is True
+
+
+def test_check_scan_id(detector_base):
+    detector_base.scaninfo.scan_id = "abcde"
+    detector_base.stopped = False
+    detector_base.check_scan_id()
+    assert detector_base.stopped is True
+    detector_base.stopped = False
+    detector_base.check_scan_id()
+    assert detector_base.stopped is False