src/aare/devices/jfjoch.py

@@ -1,6 +1,9 @@
+import functools
 import math
 import time
+from collections.abc import Callable
 from enum import Enum
+from typing import Concatenate, ParamSpec, TypeVar, final
 
 import jfjoch_client
 from aarecommon.config.beamline import get_jfjoch_url
@@ -12,9 +15,14 @@ from aarecommon.models.raster_grid import RasterGridRequest
 from aarecommon.models.rotation_scan import RotationScanRequest
 from jfjoch_client.api.default_api import DefaultApi
 from jfjoch_client.api_client import ApiClient
+from jfjoch_client.models.detector_list_element import DetectorListElement
+from jfjoch_client.models.scan_result import ScanResult
 
 logger = setup_logger("aareDAQ")
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
 
 class ScanTypeEnum(Enum):
     RASTER = "Raster"
@@ -25,23 +33,50 @@ class ScanTypeEnum(Enum):
     UNKNOWN = "Unknown"
 
 
+def needs_init(
+    method: Callable[Concatenate["JFJochWrapper", P], R],
+) -> Callable[Concatenate["JFJochWrapper", P], R]:
+    """Make sure the JFJoch client/api exist before the method touches them."""
+
+    @functools.wraps(method)
+    def wrapper(self: "JFJochWrapper", *args: P.args, **kwargs: P.kwargs) -> R:
+        self.connect()
+        return method(self, *args, **kwargs)
+
+    return wrapper
+
+
+@final
 class JFJochWrapper:
     def __init__(self, bl: MXBeamline):
         self._simulated = bl == MXBeamline.SIMULATED
         self._url = get_jfjoch_url(bl)
+        self._connected: bool = False
+        _client: ApiClient
+        _api: DefaultApi
+
+    def connect(self) -> None:
+        """Connect to JFJoch on first use; a no-op once the client exists."""
+        if self._connected:
+            return
 
         if self._url == "simulated":
             from unittest.mock import MagicMock
 
             self._client = MagicMock(spec=ApiClient)
             self._api = MagicMock(spec=DefaultApi)
+            self._connected = True
             return
 
         self._client = jfjoch_client.ApiClient(jfjoch_client.Configuration(host=self._url))
         self._api = jfjoch_client.DefaultApi(self._client)
-        self.cancel()
-        # if not self.is_idle():
-        #     self.initialize()
+        # set before cancel(), which is itself decorated and would otherwise recurse
+        self._connected = True
+        try:
+            self.cancel(on_init=True)
+        except:
+            self._connected = False
+            raise
 
     @staticmethod
     def _extract_status_code(error: Exception) -> int | None:
@@ -68,6 +103,7 @@ class JFJochWrapper:
             status_code=self._extract_status_code(error),
         )
 
+    @needs_init
     def initialize(self):
         try:
             self._api.initialize_post()
@@ -76,14 +112,19 @@ class JFJochWrapper:
                 "JFJoch initialize failed", error=e, operation="POST", endpoint="initialize_post"
             ) from e
 
-    def cancel(self):
+    @needs_init
+    def cancel(self, on_init: bool = False):
         try:
             self._api.cancel_post()
         except Exception as e:
             raise self._jfjoch_error(
-                "JFJoch cancel failed", error=e, operation="POST", endpoint="cancel_post"
+                f"JFJoch cancel failed{' on initialization' if on_init else ''}",
+                error=e,
+                operation="POST",
+                endpoint="cancel_post",
             ) from e
 
+    @needs_init
     def is_idle(self) -> bool:
         status = self._api.status_get()
         return status.state == "Idle"
@@ -199,6 +240,7 @@ class JFJochWrapper:
 
         return dataset_settings
 
+    @needs_init
     def _start_scan(
         self,
         scan_type: ScanTypeEnum,
@@ -233,6 +275,7 @@ class JFJochWrapper:
     def measure_raster(self, r: RasterGridRequest, s: DAQStatusModel, async_start: bool = True):
         self._start_scan(ScanTypeEnum.RASTER, r, s, async_start=async_start)
 
+    @needs_init
     def wait_till_running(self, timeout: float = 60):
         if self._simulated:
             return None
@@ -247,7 +290,8 @@ class JFJochWrapper:
                 endpoint="wait_until_running_post",
             ) from e
 
-    def wait_till_done(self, timeout: float) -> jfjoch_client.models.ScanResult | None:
+    @needs_init
+    def wait_till_done(self, timeout: float) -> ScanResult | None:
         if self._simulated:
             return None
         try:
@@ -261,7 +305,8 @@ class JFJochWrapper:
                 endpoint="wait_till_done_post / result_scan_get",
             ) from e
 
-    def detector(self) -> jfjoch_client.models.DetectorListElement:
+    @needs_init
+    def detector(self) -> DetectorListElement:
         try:
             detector_list = self._api.config_select_detector_get()
         except Exception as e:
@@ -296,6 +341,7 @@ class JFJochWrapper:
             "take_pedestal is not implemented in DAQ through the JFJoch API yet"
         )
 
+    @needs_init
     def get_diffraction_image(
         self,
         image_id: int,

tests/unit/devices/test_jfjoch.py

@@ -34,7 +34,10 @@ def mock_jfjoch_client():
 
 @pytest.fixture
 def jfjoch_wrapper(mock_jfjoch_client):
-    return JFJochWrapper(MXBeamline.X10SA)
+    wrapper = JFJochWrapper(MXBeamline.X10SA)
+    # the client/api are created lazily, connect up front so tests start from a live wrapper
+    wrapper.connect()
+    return wrapper
 
 
 def create_mock_daq_status():
@@ -107,6 +110,19 @@ def test_init_unknown():
         JFJochWrapper("UNKNOWN")
 
 
+def test_client_is_created_lazily(mock_jfjoch_client):
+    wrapper = JFJochWrapper(MXBeamline.X10SA)
+    assert wrapper._connected is False
+    mock_jfjoch_client.status_get.assert_not_called()
+
+    # a decorated method connects on demand, further calls reuse the same api
+    wrapper.is_idle()
+    assert wrapper._connected is True
+    api = wrapper._api
+    wrapper.is_idle()
+    assert wrapper._api is api
+
+
 def test_initialize(jfjoch_wrapper, mock_jfjoch_client):
     mock_jfjoch_client.initialize_post.reset_mock()
     jfjoch_wrapper.initialize()