diff --git a/src/aare/devices/jfjoch.py b/src/aare/devices/jfjoch.py index a0a7de36..f0d2b9c8 100644 --- a/src/aare/devices/jfjoch.py +++ b/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, diff --git a/tests/unit/devices/test_jfjoch.py b/tests/unit/devices/test_jfjoch.py index 8c5a7d13..298aa972 100644 --- a/tests/unit/devices/test_jfjoch.py +++ b/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()