diff --git a/superxas_bec/device_configs/sample_manipulator.yaml b/superxas_bec/device_configs/sample_manipulator.yaml new file mode 100644 index 0000000..d26ebdd --- /dev/null +++ b/superxas_bec/device_configs/sample_manipulator.yaml @@ -0,0 +1,27 @@ +manip_new_trx: + description: Sample Manipulator X-Translation + deviceClass: ophyd.EpicsMotor + deviceConfig: + prefix: X10DA-ES1-MAN:TRX + enabled: true + onFailure: retry + readoutPriority: baseline + softwareTrigger: false +manip_new_try: + description: Sample Manipulator Y-Translation + deviceClass: ophyd.EpicsMotor + deviceConfig: + prefix: X10DA-ES1-MAN:TRY + enabled: true + onFailure: retry + readoutPriority: baseline + softwareTrigger: false +manip_new_trz: + description: Sample Manipulator Z - Along beam + deviceClass: ophyd.EpicsMotor + deviceConfig: + prefix: X10DA-ES1-MAN:TRZ + enabled: true + onFailure: retry + readoutPriority: baseline + softwareTrigger: false \ No newline at end of file diff --git a/superxas_bec/device_configs/timepix.yaml b/superxas_bec/device_configs/timepix.yaml new file mode 100644 index 0000000..2df5f50 --- /dev/null +++ b/superxas_bec/device_configs/timepix.yaml @@ -0,0 +1,40 @@ +# sample_manipulator: +# - !include ./sample_manipulator.yaml + +# pos_x: +# description: Sample Manipulator Y-Translation +# deviceClass: ophyd.EpicsMotor +# deviceConfig: +# prefix: X10DA-ES1-SH1:POSX +# enabled: true +# onFailure: retry +# readoutPriority: baseline +# softwareTrigger: false + +samx: + readoutPriority: baseline + deviceClass: ophyd_devices.SimPositioner + deviceConfig: + delay: 1 + limits: + - -50 + - 50 + tolerance: 0.01 + update_frequency: 400 + deviceTags: + - user motors + enabled: true + readOnly: false + +timepix: + readoutPriority: async + description: ASI Serval Timepix Detector + deviceClass: superxas_bec.devices.timepix.timepix.Timepix + deviceConfig: + prefix: "X10DA-ES-TPX1:" + backend_rest_url: "P6-0008.psi.ch:8452" + hostname: "x10da-bec-001.psi.ch" + onFailure: retry + enabled: true + readOnly: false + softwareTrigger: true diff --git a/superxas_bec/devices/timepix/beamline_test.py b/superxas_bec/devices/timepix/beamline_test.py deleted file mode 100644 index a748ddf..0000000 --- a/superxas_bec/devices/timepix/beamline_test.py +++ /dev/null @@ -1,57 +0,0 @@ -from ophyd_devices.devices.areadetector.cam import ASItpxCam -from ophyd import ADBase -from ophyd import Component as Cpt -from ophyd_devices import TransitionStatus -from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import NetAddresses - -class SuperXasTimePix(ADBase): - cam = Cpt(ASItpxCam, 'cam1:') - -if __name__ == """__main__""": - import time - - from superxas_bec.devices.timepix.timepix import Timepix - from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_mock_server import ( - TimePixFlyMockServer, - ) - - timepix_control = SuperXasTimePix(name='timepix_control', prefix='X10DA-ES-TPX1:') - # timepix_control.wait_for_connection(all_signals=True, timeout=10) - # Create a Timepix object - rest_url = 'p4-0017.psi.ch:8452' - ws_url = 'p4-0017:8452/ws' - timepix = Timepix(name="TimePixDetector", rest_url=rest_url, ws_url=ws_url) - # timepix = Timepix(name="TimePixDetector") - timepix.on_connected() - - ## LOOP for a scan - timepix.stage() - net_addresses:NetAddresses = timepix.timepix_fly_client.get_net_addresses() - print(f"Net addresses: {net_addresses}") - file_path = f"tcp://connect@{net_addresses.address}" - print(f"Using file path: {file_path}") - period = 1 - num_images =5 - timepix_control.cam.trigger_mode.set(0).wait(timeout=10) - timepix_control.cam.acquire_period.set(period).wait(timeout=10) - timepix_control.cam.acquire_time.set(period-2e-3).wait(timeout=10) - timepix_control.cam.num_images.set(num_images).wait(timeout=10) - timepix_control.cam.raw_enable.set(True).wait(timeout=10) - timepix_control.cam.raw_file_path.put(file_path)#.wait(timeout=10) - timepix_control.cam.raw_file_template.put('test_raw')#.wait(timeout=10) - - timepix.pre_scan() # This sends the address to the ASI server from TimePixFly - status = TransitionStatus(timepix_control.cam.acquire, [1, 0]) - timepix_control.cam.acquire.put(1) - status.wait(timeout=20) - # time.sleep(5) - time.sleep(5) - print(timepix._data_buffer) - print(timepix.timepix_fly_client.last_error()) - - # status = timepix.complete() - # status.wait(timeout=10) - # for ii, msg in enumerate(timepix._data_buffer): - # print(f"Received data {ii} with message type {msg['type']}") - timepix.unstage() - timepix.destroy() diff --git a/superxas_bec/devices/timepix/test_script.py b/superxas_bec/devices/timepix/test_script.py deleted file mode 100644 index c11709f..0000000 --- a/superxas_bec/devices/timepix/test_script.py +++ /dev/null @@ -1,80 +0,0 @@ -import enum - -from ophyd import Component as Cpt -from ophyd import Device, EpicsSignal, Kind -from ophyd_devices import CompareStatus -from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase - -# class TimepixServerStatus(int, enum.Enum): -# DONE = 0 -# ACQUIRE = 1 - - -# class TimepixServerControl(Device): -# """Prefix of the detector 'X10DA-ES-TPX1:cam1:' """ - -# acquire = Cpt(EpicsSignal, 'Acquire', kind=Kind.omitted, doc="Detector acquiring signal") -# preview_enable= Cpt(EpicsSignal, 'PreviewEnable', kind=Kind.omitted, doc='Preview enable y/n') - - -# class TimepixServer(PSIDeviceBase, TimepixServerControl): -# """Prefix of the detector 'X10DA-ES-TPX1:cam1:' """ - -# def on_stage(self)-> None: -# self.preview_enable.set(1).wait(timeout=3) -# if self.acquire.get() != TimepixServerStatus.DONE: -# status = CompareStatus(self.acquire, TimepixServerStatus.DONE) -# self.cancel_on_stop(status) -# self.acquire.put(TimepixServerStatus.DONE) -# status.wait(timeout=3) - -# def on_pre_scan(self): -# status = CompareStatus(self.acquire, TimepixServerStatus.ACQUIRE) -# self.cancel_on_stop(status) -# self.acquire.put(1) -# status.wait(timeout=3) - - -if __name__ == """__main__""": - import time - - from superxas_bec.devices.timepix.timepix import Timepix - - # from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_mock_server import ( - # TimePixFlyMockServer, - # ) - # # mock_server = TimePixFlyMockServer() - # server = TimepixServer(name='server', prefix='X10DA-ES-TPX1:cam1:') - # Create a Timepix object - rest_url = "p4-0017.psi.ch:8453" - ws_url = "p4-0017:8453/ws" - # data_server_host = "x10da-bec-001.psi.ch" - # port = 3030 - timepix = Timepix( - name="TimePixDetector", prefix="X10DA-ES-TPX1:", rest_url=rest_url, ws_url=ws_url - ) - timepix.on_connected() - timepix.wait_for_connection(all_signals=True, timeout=10) - # server.wait_for_connection(all_signals=True, timeout=10) - timepix.on_connected() - - ## LOOP for a scan - timepix.stage() - print("Timepix staged.") - # server.stage() - # timepix.pre_scan() - status = timepix.on_pre_scan() - print("Timepix staged in pre scan.") - - timepix.trigger() - print("Timepix triggered.") - # time.sleep(5) - status = timepix.complete() - while not status.done: - time.sleep(0.1) - print(timepix.cam.acquire_busy.get()) - status.wait(timeout=20) - for ii, msg in enumerate(timepix._data_buffer): - print(f"Received data {ii} with message type {msg['type']}") - timepix.unstage() - timepix.destroy() diff --git a/superxas_bec/devices/timepix/test_script_backup.py b/superxas_bec/devices/timepix/test_script_backup.py deleted file mode 100644 index de931ef..0000000 --- a/superxas_bec/devices/timepix/test_script_backup.py +++ /dev/null @@ -1,45 +0,0 @@ -if __name__ == """__main__""": - import threading - import time - from unittest import mock - - import ophyd - from ophyd_devices.tests.utils import MockPV, patch_dual_pvs - - from superxas_bec.devices.timepix.timepix import Timepix - from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_mock_server import ( - TimePixFlyMockServer, - ) - - mock_server = TimePixFlyMockServer() - with mock.patch.object(ophyd, "cl") as mock_cl: - mock_cl.get_pv = MockPV - mock_cl.thread_class = threading.Thread - timepix = Timepix(name="timepix", prefix="") - patch_dual_pvs(timepix) - timepix.timepix_fly_client.on_connected() - timepix._reset_buffers() - timepix.start_data_server() - # timepix.on_connected() - - ## LOOP for a scan - timepix.stage() - status = timepix.pre_scan() - timepix.cam.acquire_busy._read_pv.mock_data = 1 - status.wait(timeout=10) - print(f"Pre-scan status: {status}") - mock_server.start_acquisition() - print("Acquisition started on mock server.") - # time.sleep(5) - status = timepix.complete() - time.sleep(8) - timepix.cam.acquire_busy._read_pv.mock_data = 0 - status.wait(timeout=10) - for ii, msg in enumerate(timepix._data_buffer): - print(f"Message {ii}: {msg.keys()}") - print(f"Received data {ii} with message type {msg['type']}") - if msg["type"] == "EndFrame": - print(f"Received error message: {msg}") - # print(f"Data: {msg}") # Print first 50 characters of data - timepix.unstage() - timepix.destroy() diff --git a/superxas_bec/devices/timepix/timepix.py b/superxas_bec/devices/timepix/timepix.py index c622190..ddf2fe6 100644 --- a/superxas_bec/devices/timepix/timepix.py +++ b/superxas_bec/devices/timepix/timepix.py @@ -1,36 +1,45 @@ """ TimePix Detector class for interfacing with the TimePix detector. The timepix_signals module -implements the HTTP communication to the REST API for the tpx3app app. +implements the HTTP communication to the REST API for the tpx3app app. The implementation +of the backend is stored in the timepix_fly_client module. This is combined with the control +interface in EPICS, which is implemented via the 'ASItpxCam' class. """ -import atexit +from __future__ import annotations + import enum -import json -import signal -import socket import threading import time -from typing import Literal +import traceback +from typing import TYPE_CHECKING, Any, Literal +import numpy as np from bec_lib.logger import bec_logger from ophyd import ADBase from ophyd import Component as Cpt -from ophyd import Device, DeviceStatus, StatusBase -from ophyd_devices import AndStatus, CompareStatus +from ophyd import DeviceStatus, Kind, StatusBase +from ophyd_devices import AsyncSignal, CompareStatus, PreviewSignal, TransitionStatus from ophyd_devices.devices.areadetector.cam import ASItpxCam +from ophyd_devices.devices.areadetector.plugins import HDF5Plugin_V35, ImagePlugin_V35 from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase +from typeguard import typechecked -from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client import ( - TimepixFlyClient, - TimePixFlyStatus, -) +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_backend import TimepixFlyBackend +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client import TimePixFlyStatus from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import ( OtherConfigModel, PixelMap, ) +from superxas_bec.devices.timepix.utils import AndStatusWithList + +if TYPE_CHECKING: + from bec_lib.messages import DevicePreviewMessage + logger = bec_logger.logger +# pylint: disable=redefined-outer-name + class TDCEdge(int, enum.Enum): """TDC Edge enum options for TimePix detector.""" @@ -99,37 +108,63 @@ class EXPOSUREMODE(int, enum.Enum): TRIGGER_WIDTH = 1 +class DATASOURCE(int, enum.Enum): + """Data source for AD Epics backend for Timepix.""" + + NONE = 0 + PREVIEW = 1 + IMAGE = 2 + + # pylint: disable=too-many-instance-attributes, too-many-arguments, too-many-locals class TimePixControl(ADBase): """Interface for the TimePix EPICS control of the TimePix detector.""" cam = Cpt(ASItpxCam, "cam1:") + image = Cpt(ImagePlugin_V35, "image1:") + hdf = Cpt(HDF5Plugin_V35, "HDF1:") + # latest hdf5 plugin + # latest image plugin + + class Timepix(PSIDeviceBase, TimePixControl): """ - TimePix Detector class for interfacing with the TimePix detector. - The Timepix detector REST API service module implements the HTTP communication to the REST API for the tpx3app app - - DATA_SERVER_HOST = "x10da-bec-001.psi.ch" # Default data server host for TimePix detector - DATA_SERVER_PORT = 3015 # Default data server port for TimePix detector - rest_url = 'p4-0017.psi.ch:8452' - ws_url = 'p4-0017:8452/ws' - - Prefix of EPICS Control is 'X10DA-ES-TPX1:' + TimePix class. The IOC is running with the prefix 'X10DA-ES-TPX1:'. + The TimePixFly backend is running on p4-0017.psi.ch. Please check the port from the app + running in headless server mode. The backend_rest url can for instance be 'p4-0017.psi.ch:8452'. + The hostname needs to be set to the name of this machine, e.h. x10da-bec-001.psi.ch. """ + _DETECTOR_SHAPE = (512, 512) # Shape of the TimePix detector + USER_ACCESS = ["troin", "troistep", "get_pixel_map", "set_pixel_map"] + + # TODO Check names with beamline team, async signals current receive a nested data name structure + xes_data = Cpt(AsyncSignal, name="xes_data", ndim=2, max_size=1000) + xes_spectra = Cpt(AsyncSignal, name="xes_spectra", ndim=1, max_size=1000) + tds_period = Cpt(AsyncSignal, name="tds_period", ndim=0, async_update={"type": "add", "max_shape": [None]}, max_size=1000) + tds_total_events = Cpt(AsyncSignal, name="tds_total_events", ndim=0, async_update={"type": "add", "max_shape": [None]}, max_size=1000) + # # Here, we currently inherit a nested name structure --> xes_info_tds_period, xes_info_tds_total_events + + preview = Cpt( + PreviewSignal, + name="preview", + ndim=2, + num_rotation_90=0, # Check this + doc="Preview signal for the Pilatus Detector", + ) + def __init__( self, *, name, prefix: str, - rest_url: str = "localhost:8452", - ws_url: str = "localhost:8452/ws", + backend_rest_url: str, + hostname: str | None = None, + socket_port: int = 0, scan_info=None, device_manager=None, - hostname: str | None = None, - host_port: int | None = None, **kwargs, ): """ @@ -137,51 +172,207 @@ class Timepix(PSIDeviceBase, TimePixControl): Args: name (str): Name of the device. - rest_url (str): Host address of the TimePixFly backend server. - ws_url (str): WebSocket URL for the TimePixFly backend. - scan_info (dict, optional): Scan information. Defaults to None. - device_manager (DeviceManager, optional): Device manager instance. Defaults to None. - data_server_host (str, optional): Host address for the data server. Defaults to DATA - SERVER_HOST. - data_server_port (int, optional): Port for the data server. Defaults to DATA_SERVER_PORT. - **kwargs: Additional keyword arguments for the PSIDeviceBase class. + prefix (str): EPICS prefix for the device. + backend_rest_url (str): URL of the TimePixFly backend REST API. + hostname (str | None): Hostname of the machine running the backend. Defaults to None + which will use the current machine's hostname. + socket_port (int): Port for the socket connection to the backend. Defaults to 0 + which will use the default port from the backend. + scan_info: Scan information object, if available. + device_manager: Device manager instance, if available. + **kwargs: Additional keyword arguments for the base class. """ + + self.backend = TimepixFlyBackend( + backend_rest_url=backend_rest_url, hostname=hostname, socket_port=socket_port + ) + self._pixel_map = PixelMap( + chips=[ + [{"i": 256 ^ 2 - 1, "p": [0, 1], "f": [0.5, 0.5]}], + [{"i": 255 * 256, "p": [0, 1], "f": [0.5, 0.5]}], + [{"i": 255, "p": [1, 2], "f": [0.5, 0.5]}], + [{"i": 0, "p": [1, 2], "f": [0.5, 0.5]}], + [{"i": 256 ^ 2 - 1, "p": [0, 1], "f": [0.5, 0.5]}], + [{"i": 255 * 256, "p": [0, 1], "f": [0.5, 0.5]}], + [{"i": 255, "p": [1, 2], "f": [0.5, 0.5]}], + [{"i": 0, "p": [1, 2], "f": [0.5, 0.5]}], + ] + ) + self._n_energy_points = 3 + self._troistep = 1 + self._troin = 5000 super().__init__( name=name, prefix=prefix, scan_info=scan_info, device_manager=device_manager, **kwargs ) - self._ws_url = ws_url - self._rest_url = rest_url - self.timepix_fly_client = TimepixFlyClient(rest_url=rest_url, ws_url=ws_url, parent=self) - if hostname is None: - hostname = socket.getfqdn() - if not hostname.endswith(".psi.ch"): - logger.info(f"Found host without psi.ch domain {hostname}") + self._poll_thread = threading.Thread( + target=self._poll_array_data, daemon=True, name=f"{self.name}_poll_thread" + ) + self._poll_thread_kill_event = threading.Event() + self._poll_rate = 1 # Poll rate in Hz + self._pv_timeout = 5 + self._readout_time = 2.1e-3 # 2.1ms readout time to ensure readout is >2ms, required from ASI serval server.. + self.r_lock = threading.RLock() # Lock to access the message buffer safely - if host_port is None: - host_port = 3015 - self._hostname = hostname - self._data_server_host: str | None = None - self._data_server_port = host_port - self._rlock = threading.RLock() - # Data server - self._data_server_thread = None - self._data_server_thread_event = None - self._data_server_started = threading.Event() - # Socket server - self._socket_server = None - self._socket_server_timeout = 0.1 - self._socket_server_buffer_size = 4096 - # Data buffers - self._data_buffer = [] - self._global_buffer = "" # Global buffer to store received data - # Decoding - self._decoder = json.JSONDecoder() - # Wait timeout - self._pv_timeout = 5.0 # Default timeout for PV operations - # Register cleanup - atexit.register(self.on_destroy) + def stage(self) -> list[object] | DeviceStatus | StatusBase: # type: ignore + """Stage the device. + + Super stage not safe to call..""" + self.stopped = False + status = self.on_stage() # pylint: disable=assignment-from-no-return + if isinstance(status, StatusBase): + return status + return [] - ### Beamline specifi methods for the TimePix Detector integration ### + def _poll_array_data(self): + """Poll the array data for preview updates.""" + while not self._poll_thread_kill_event.wait(1 / self._poll_rate): + try: + # logger.info(f"Running poll loop for {self.name}..") + value = self.image.array_data.get() + if value is None: + continue + width = self.image.array_size.width.get() + height = self.image.array_size.height.get() + # Geometry correction for the image + data = np.reshape(value, (height, width)) + last_image: DevicePreviewMessage = self.preview.get() + # logger.info(f"Preview image for {self.name} has shape {data.shape}") + if last_image is not None: + if np.array_equal(data, last_image.data): + # No update if image is the same, ~2.5ms on 2400x2400 image (6M) + logger.debug( + f"Pilatus preview image for {self.name} is the same as last one, not updating." + ) + continue + + logger.debug(f"Setting preview datsa for {self.name}") + self.preview.put(data) + except Exception: # pylint: disable=broad-except + content = traceback.format_exc() + logger.error( + f"Error while polling array data for preview of {self.name}: {content}" + ) + + ### + def msg_buffer_callback( + self, + start_frame: dict[ + Literal[ + "type", "Mode", "TRoiStart", "TRoiStep", "TRoiN", "NumEnergyPoints", "save_interval" + ], + Any, + ], + data_frames: list[ + dict[ + Literal["type", "period", "totalEvents", "TDSpectra", "beforeROI", "afterROI"], Any + ] + ], + end_frame: dict[Literal["type", "error"], Any], + ): + """ + Callback method to be attached to the backend to process the message buffer. The callback expects + start_frame, data_frames, and end_frame as arguments. Additionally, one may pass extra kwargs that + will be passed to the callback function. + + Args: + start_frame (dict): The StartFrame. Dictionary representation of detailed structure + described in model .timepix_fly_client.timepix_fly_interface.TimepixStartFrame + data_frames (list): List of XesData frames. Dictionary of structures described in + model .timepix_fly_client.timepix_fly_interface.TimepixDataFrame + end_frame (dict): The EndFrame. Dictionary representation of detailed structure + described in model .timepix_fly_client.timepix_fly_interface.TimepixEndFrame + """ + n_energy_points = start_frame.get("NumEnergyPoints", None) + # if n_energy_points is None: + # logger.error( + # f"NumEnergyPoints not found in start_frame: {start_frame}. Have we received the correct frame?" + # ) + # return + # # TODO What should we do here if n_energy_points and troin do not match with the expected values? + if n_energy_points != self._n_energy_points: + logger.error( + f"Number of energy points {n_energy_points} does not match expected {self._n_energy_points}." + ) + troin = start_frame["TRoiN"] + if troin != self._troin: + logger.error(f"Number of pixels {troin} does not match expected {self._troin}.") + + # Create return data + xes_data = np.zeros((n_energy_points, troin), dtype=np.float32) # dtype from backend code + tds_period = 0 + tds_total_events = 0 + if len(data_frames) == 0: + logger.error( + f"No data frames received in msg_buffer; for start_frame: {start_frame}, end_frame: {end_frame}" + ) + else: + for msg in data_frames: + tds_period += msg["period"] + tds_total_events += msg["totalEvents"] + for ii in range(n_energy_points): + xes_data[ii, :] += msg["TDSpectra"][ii::n_energy_points] + # Put XES data + self.tds_period.put(tds_period) + self.tds_total_events.put(tds_total_events) + self.xes_data.put(xes_data, async_update={"type": "add", "max_shape": [None, n_energy_points, troin]}) + self.xes_spectra.put(xes_data.sum(axis=1), async_update={"type": "add", "max_shape": [None, troin]}) + logger.debug(f"Device data set for Timepix with {tds_period}, {tds_total_events}") + + ### User ACCESS methods + + def get_pixel_map(self) -> dict: + """Get the current pixel map as a dictionary.""" + return self._pixel_map.model_dump() + + def set_pixel_map(self, pixel_map: dict) -> None: + """Set the pixel map from a dictionary.""" + self._pixel_map = PixelMap.model_validate(pixel_map) + + @property + def n_energy_points(self) -> int: + """Energy points for the TimePix detector.""" + return self._n_energy_points + + @property + def pixel_map(self) -> PixelMap: + """Get the current pixel map of the TimePix detector.""" + return self._pixel_map + + @pixel_map.setter + @typechecked + def pixel_map(self, value: PixelMap): + self._pixel_map = value + # TODO set energy points based on pixel map... + + @property + def troistep(self) -> int: + """Get the current ROI step size.""" + return self._troistep + + @troistep.setter + @typechecked + def troistep(self, value: int): + """Set the ROI step size.""" + if value <= 0: + raise ValueError("ROI step size must be a positive integer.") + self._troistep = value + + @property + def troin(self) -> int: + """Get the current ROI number of pixels.""" + return self._troin + + @troin.setter + @typechecked + def troin(self, value: int): + """Set the ROI number of pixels.""" + if value <= 0: + raise ValueError("ROI number of pixels must be a positive integer.") + self._troin = value + + ###################################################################### + ### Beamline specific methods for the TimePix Detector integration ### + ###################################################################### def on_init(self) -> None: """ @@ -190,13 +381,7 @@ class Timepix(PSIDeviceBase, TimePixControl): No signals are connected at this point. If you like to set default values on signals, please use on_connected instead. """ - - def sigint_handler(*args): - """Ensure that the on_destroy method is called when the process is killed.""" - self.on_destroy() - - signal.signal(signal.SIGINT, sigint_handler) - signal.signal(signal.SIGTERM, sigint_handler) + self.backend.on_init() def on_connected(self) -> None: """ @@ -211,11 +396,22 @@ class Timepix(PSIDeviceBase, TimePixControl): self.cam.trigger_mode.set(TRIGGERMODE.INTERNAL).wait(timeout=self._pv_timeout) self.cam.trigger_source.set(TRIGGERSOURCE.HDMI1_1).wait(timeout=self._pv_timeout) self.cam.exposure_mode.set(EXPOSUREMODE.TIMED).wait(timeout=self._pv_timeout) + # self.image.unique_id.set(1).wait(timeout=self._pv_timeout) # Prepare backend for TimePixFly - self.timepix_fly_client.on_connected() - self._reset_buffers() - self.start_data_server() + self.backend.on_connected() + # Register the callback for processing data received by the backend + # TODO add initial callback again once issues are resolved + self.backend.add_callback(self.msg_buffer_callback) + self._msg_dump = [] + + # def _on_msg_received(start_frame, data_frame, end_frame): + # """Callback""" + # self._msg_dump.append( + # {"start_frame": start_frame, "data_frame": data_frame, "end_frame": end_frame} + # ) + + # self.backend.add_callback(_on_msg_received) def on_stage(self) -> DeviceStatus | StatusBase | None: """ @@ -223,84 +419,123 @@ class Timepix(PSIDeviceBase, TimePixControl): Information about the upcoming scan can be accessed from the scan_info (self.scan_info.msg) object. """ - # currently hardcode acquire time.. -> discuss logic here - self.cam.acquire_time.set(0.998).wait(timeout=self._pv_timeout) - self.cam.acquire_period.set(1.0).wait(timeout=self._pv_timeout) - self.cam.num_images.set(10).wait(timeout=self._pv_timeout) + exp_time = self.scan_info.msg.scan_parameters.get("exp_time", 0) + if exp_time - self._readout_time <= 0: + raise ValueError( + f"Exposure time {exp_time} must be greater than readout time {self._readout_time}." + ) + num_images = self.scan_info.msg.scan_parameters.get("frames_per_trigger", 1) + logger.debug(f"Setting exposure time to {exp_time} and number of images to {num_images}") + + self.cam.acquire_time.set(exp_time - self._readout_time).wait(timeout=self._pv_timeout) + self.cam.acquire_period.set(exp_time).wait(timeout=self._pv_timeout) + self.cam.num_images.set(num_images).wait(timeout=self._pv_timeout) + self.cam.raw_enable.set(1).wait(timeout=self._pv_timeout) + self.cam.data_source.set(DATASOURCE.IMAGE).wait(timeout=self._pv_timeout) # ------------------------- # Prepare TimePixFly - self._reset_buffers() - status = DeviceStatus(self) - self.cancel_on_stop(status) - self.timepix_fly_client.add_status_callback( - status, - success=[TimePixFlyStatus.CONFIG], - error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + other_config = OtherConfigModel( + TRoiStep=self.troistep, + TRoiN=self.troin, + output_uri=f"tcp:{self.backend.hostname}:{self.backend.socket_port}", ) - if self._data_server_host is None: - raise RuntimeError(f"Data server host is not set for device {self.name}.") - # Parse scan info for OtherConfig - config = OtherConfigModel( - output_uri=f"tcp:{self._data_server_host}:{self._data_server_port}", - TRoiStep=1, - TRoiN=5000, + logger.debug(f"Current TimePixFly configuration: {other_config}") + pixel_map = self.pixel_map + self.backend.on_stage(other_config=other_config, pixel_map=pixel_map) + + # Fetch the backend socket info + net_add = self.backend.timepix_fly_client.get_net_addresses() + logger.debug(f"Using net_add for timepix_fly backend {net_add}") + self.cam.raw_file_template.set("").wait(timeout=self._pv_timeout) + self.cam.raw_file_path.set(f"tcp://connect@{net_add.address}").wait( + timeout=self._pv_timeout ) - logger.info(config) - # Parse pixel map from scan info if needed, otherwise use some default pixel map. - pixel_map = PixelMap( - chips=[ - [{"i": 256 ^ 2 - 1, "p": [0, 1], "f": [0.5, 0.5]}], - [{"i": 255 * 256, "p": [0, 1], "f": [0.5, 0.5]}], - [{"i": 255, "p": [1, 2], "f": [0.5, 0.5]}], - [{"i": 0, "p": [1, 2], "f": [0.5, 0.5]}], - ] - ) - status.wait(timeout=5.0) - self.timepix_fly_client.set_other_config(config) - self.timepix_fly_client.set_pixel_map(pixel_map) def on_unstage(self) -> None: """Called while unstaging the device.""" - self._reset_buffers() + self.backend.on_unstage() def on_pre_scan(self) -> StatusBase: """Called right before the scan starts on all devices automatically.""" - status_acquire = CompareStatus(self.cam.acquire_busy, ACQUIRESTATUS.DONE) - status_detector_state = CompareStatus(self.cam.detector_state, DETECTORSTATE.IDLE) - status_detector = AndStatus(status_acquire, status_detector_state) - self.cancel_on_stop(status_detector) - status = DeviceStatus(self) - self.cancel_on_stop(status) - self.timepix_fly_client.add_status_callback( - status, - success=[TimePixFlyStatus.AWAIT_CONNECTION], - error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + # TODO check detector error state and raise if we see an error in the detector status + status_acquire = CompareStatus(self.cam.acquire_busy, ACQUIRESTATUS.DONE, timeout=3) + self.cancel_on_stop(status_acquire) + return status_acquire + + def _trigger_callback(self, status: DeviceStatus) -> None: + """Trigger callback to start the acquisition.""" + if status.done and status.success: + status.device.cam.acquire.put(1) + return + logger.error( + f"Status callback from TimePixFly backend trigger failed. Exception {status.exception()}" ) - self.timepix_fly_client.start() - pre_scan_status = AndStatus(status, status_detector) - self.cancel_on_stop(pre_scan_status) # --> State goes to setup - return pre_scan_status def on_trigger(self) -> DeviceStatus | StatusBase | None: """Called when the device is triggered.""" + + # First we make sure that the backend reach 'config' state. This needs to happend before each trigger. + status_backend_config = DeviceStatus(self) + self.cancel_on_stop(status_backend_config) + self.backend.timepix_fly_client.add_status_callback( + status=status_backend_config, + success=[TimePixFlyStatus.CONFIG], + error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + ) + try: + status_backend_config.wait(timeout=5) + except TimeoutError: + # pylint:disable=raise-missing-from + raise TimeoutError( + f"TimePixFly backend of device {self.name} failed to reach 'config' state in trigger" + ) + + # Prepare status objects for coordination of actions + status_camera = TransitionStatus( + self.cam.acquire_busy, [ACQUIRESTATUS.DONE, ACQUIRESTATUS.ACQUIRING, ACQUIRESTATUS.DONE] + ) + self.cancel_on_stop(status_camera) + status = self.backend.on_trigger() + status.wait(timeout=5) # Wait until backend trigger is done + + # Backend ready for connection, now start camera + status = self.backend.on_trigger_finished() + self.cancel_on_stop(status) + return_status = status_camera & status self.cam.acquire.put(1) + return return_status + + # # Add Collect callback + # self.backend.timepix_fly_client.add_status_callback( + # status_backend_collect_started, + # success=[TimePixFlyStatus.COLLECT], + # error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + # ) + + # # Start on trigger on backend + # status_backend_on_trigger = self.backend.on_trigger(status=status_backend_on_trigger) + + # status = AndStatusWithList( + # status_list=[status_camera, status_backend_on_trigger, status_backend_collect_started], + # device=self, + # ) + # self.cancel_on_stop(status) + # return status def on_complete(self) -> DeviceStatus | StatusBase | None: """Called to inquire if a device has completed a scans.""" # Detector Control status_detector = CompareStatus(self.cam.acquire_busy, ACQUIRESTATUS.DONE) - self.cancel_on_stop(status_detector) # TimepixFFly - status = DeviceStatus(self) - logger.info(f"Registering status callback {id(status)} for {self.name} on complete.") - self.cancel_on_stop(status) - self.timepix_fly_client.add_status_callback( - status, - success=[TimePixFlyStatus.CONFIG], - error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + status_backend = DeviceStatus(self) + # Add callback to the backend complete handling + status_backend = self.backend.on_complete(status=status_backend) + # Combine the statuses + complete_status = AndStatusWithList( + status_list=[status_backend, status_detector], device=self ) - complete_status = AndStatus(status, status_detector) + self.cancel_on_stop(complete_status) return complete_status def on_kickoff(self) -> DeviceStatus | StatusBase | None: @@ -308,250 +543,83 @@ class Timepix(PSIDeviceBase, TimePixControl): def on_stop(self) -> None: """Called when the device is stopped.""" + # Camera self.cam.acquire.put(0) - self.timepix_fly_client.stop_running_collection() - self._reset_buffers() - - ### Custom methods for the TimePix Data server ### - - def _wait_for_state_condition( - self, state: Literal["init", "config", "setup", "collect", "shutdown"], timeout: float = 5.0 - ) -> None: - """ - Wait for the TimePixFly backend to reach a specific state. - - Args: - state (Literal["init", "config", "setup", "collect", "shutdown"]): The state to wait for. - timeout (float): The maximum time to wait for the state in seconds. Default is 5.0 seconds. - Raises: - RuntimeError: If the TimePixFly backend does not reach the specified state within the timeout. - """ - - def _check_state(): - return self.timepix_fly_client.state().state == state - - if self.wait_for_condition(_check_state, timeout=timeout, interval=0.25) is False: - raise RuntimeError( - f"TimePix Fly client {self.name} did not reach the '{state}' state in time." - f"Current state: {self.timepix_fly_client.state().state}" - ) - - def _start_data_receiver(self): - """Start the data server thread. If the thread is already running, do nothing.""" - if self._data_server_thread is not None and self._data_server_thread.is_alive(): - return - self._data_server_thread_event = threading.Event() - self._data_server_thread = threading.Thread( - target=self._receive_data_on_socket, name=f"{self.name}_data_server" - ) - self._data_server_thread.start() - - def _stop_data_receiver(self): - if self._data_server_thread is not None and self._data_server_thread.is_alive(): - self._data_server_thread_event.set() - self._data_server_thread.join(timeout=1.0) - if self._data_server_thread is not None and self._data_server_thread.is_alive(): - logger.warning( - f"Data server thread did not stop gracefully for device {self.name}." - ) - else: - logger.warning( - f"Data server thread is not running or has already stopped for device {self.name}." - ) - - def _receive_data_on_socket(self): - """Receive data on socket connection.""" - buffer = "" - # for testing purposes, move logic into separate function for each while loop - while not self._data_server_thread_event.is_set(): - try: - self._data_server_started.set() - # logger.info(self._socket_server) - conn, addr = self._socket_server.accept() - logger.info(f"Accepted connection from {addr}") - with conn: - while not self._data_server_thread_event.is_set(): - # TODO check if recv is blocking try and except socket.timeout - chunk = conn.recv(4096) # Adjust buffer size as needed - if not chunk: - time.sleep(0.1) # No data received, wait a bit before next attempt - continue - logger.info(f"Received data: {len(chunk)} bytes") - buffer += chunk.decode("utf-8") # Trailing byte, i.e. -> "}\n" - self._global_buffer += ( - buffer # Append to global buffer, #TODO check why endfram misses - ) - if buffer.endswith("}\n"): - logger.info( - f"Buffer starts with {buffer[:50]}... and ends with {buffer[-50:]}" - ) - for entry in buffer.split("}\n"): - if entry: - self._decode_received_data(entry + "}") - buffer = "" # Reset buffer after processing - - # Ignore timeout exception on socket aslong as server_thread_event is not set - except socket.timeout: - pass - - def _decode_received_data(self, buffer: str) -> None: - """ - Decode the received data from the socket. - This method should be overridden to implement the actual decoding logic. - """ - try: - obj, _ = self._decoder.raw_decode(buffer) - self._data_buffer.append(obj) - except json.JSONDecodeError: - logger.warning(f"Failed to decode JSON from buffer: {buffer}") - # If decoding fails, append the data to the buffer and wait for more data - - # TD Spectra is organized as TROIN * num_energy_points (send in startframe) + energy_point -> 'p' - # TROIN = 10, energy_points = 2 - # TROIN: (0,0), (0,1), (1,0), (1,1), (2,0), (2,1), (3,0), (3,1), (4,0), (4,1), etc... - - def _reset_buffers(self): - """Reset the data buffers.""" - logger.info(f"Resetting data buffers for {self.name}.") - self._data_buffer = [] + # Backend + self.backend.on_stop() def on_destroy(self): - """Cleanup method to stop the data server thread, clean up the socket server""" - self.timepix_fly_client.shutdown() - with self._rlock: - if self._socket_server: - try: - self._stop_data_receiver() - except Exception as e: - logger.info( - f"Failed to stop data receiver for device {self.name} with exception {e}." - ) - try: - self._socket_server.close() - except Exception as e: # pylint: disable=broad-except - logger.warning(f"Failed to shutdown socket server. Error: {e}") - - def restart_data_receiver(self): - """Restart the data receiver thread.""" - self._stop_data_receiver() - self._reset_buffers() - if not self._socket_server: - logger.warning(f"No socket server found for {self.name}. Starting a new data server.") - self.start_data_server() - else: - self._start_data_receiver() - logger.info("Data receiver thread restarted.") - - def start_data_server(self): - """ - Start the data server for the TimePix detector. - This method should be overridden to implement the actual data server logic. - """ - # AF_INET6 for IPv6, use AF_INET for IPv4; for localhost this may be different depending on the system - # TODO add an os check if self._data_server_host is localhost - # self._socket_server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) #BEAMLINE USAGE! - # addr_info = socket.getaddrinfo( - # self._data_server_host, self._data_server_port, socket.AF_UNSPEC, socket.SOCK_STREAM - # ) - # # Use the first valid address returned by getaddrinfo - # af, socktype, proto, _, _ = addr_info[0] - # self._socket_server = socket.socket(af, socktype, proto) - info = socket.getaddrinfo( - self._hostname, - port=self._data_server_port, - family=socket.AF_INET, - type=socket.SOCK_STREAM, - ) - if len(info) == 0: - raise RuntimeError(f"No socket info found") - if len(info) > 1: - logger.warning(f"Found multiple socket interfaces {info}, using the first one") - af, socktype, proto, _, host_port_info = info[0] - self._data_server_host = host_port_info[0] # Hostname or IP address - self._data_server_port = host_port_info[1] # Port number, should be the same as before - self._socket_server = socket.create_server( - host_port_info, family=af, backlog=1, reuse_port=True - ) - self._data_server_host, self._data_server_port = self._socket_server.getsockname() - logger.info( - f"Starting data server on {self._data_server_host}:{self._data_server_port} for {self.name}." - ) - # self._socket_server = socket.socket(af, socktype, proto) - # self._socket_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - # self._socket_server.bind(host_port_info) - # self._socket_server.listen(1) # Only allow one connection - self._socket_server.settimeout(self._socket_server_timeout) - self._start_data_receiver() + """Cleanup method to stop the device and clean up resources.""" + self.cam.acquire.put(0) + self._poll_thread_kill_event.set() + self.backend.on_stop() + self.backend.on_destroy() # pylint: disable=protected-access -if __name__ == "__main__": - from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_mock_server import ( - TimePixFlyMockServer, +if __name__ == "__main__": # pragma: no cover + + timepix = Timepix( + name="timepix", + prefix="X10DA-ES-TPX1:", + backend_rest_url="P6-0008.psi.ch:8452", # "P4-0017.psi.ch:8452", + hostname="x10da-bec-001.psi.ch", ) + try: + # timepix.wait_for_connection(all_signals=True, timeout=10) + timepix.on_connected() + print("Timepix connected and initialized.") + for exp_time, frames_per_trigger, runs in zip([0.1, 1, 0.2], [20, 5, 1], [10, 5, 30]): - mock_server = TimePixFlyMockServer() + print( + f"Sleeping for 0.5 seconds before starting the scan with exp_time={exp_time} " + f"and frames_per_trigger={frames_per_trigger}. and runs {runs}" + ) + time.sleep(0.5) + timepix.scan_info.msg.scan_parameters.update( + { + "exp_time": exp_time, # Set exposure time to 1 second for testing + "frames_per_trigger": frames_per_trigger, # Set frames per trigger to 5 for testing + } + ) + timepix.stage() + logger.warning(f"Timepix on stage done") + timepix.pre_scan() + logger.warning(f"Timepix on pre_scan done") - # Create a Timepix object - timepix = Timepix(name="TimePixDetector", prefix="") - timepix.on_connected() - timepix.stage() - timepix.pre_scan() - mock_server.start_acquisition() - time.sleep(5) - timepix.complete() - print(timepix._data_buffer) - print(timepix._global_buffer[-100:]) - timepix.unstage() - - # timepix = Timepix(name="TimePixDetector") - # from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_mock_server import ( - # TimePixFlyMockServer, - # ) - - # mock_server = TimePixFlyMockServer() - # timepix.on_connected() - - # infos = socket.getaddrinfo("localhost", None) - # print(infos) - - # print(f"Last error: {timepix.timepix_fly_client.last_error().message}") - # print(f"TimePix version: {timepix.timepix_fly_client.version().version}") - # print(f"TimePix state: {timepix.timepix_fly_client.state().state}") - # timepix.timepix_fly_client.set_other_config( - # OtherConfigModel( - # output_uri=f"tcp:{timepix._data_server_host}:{timepix._data_server_port}", - # TRoiStep=1, - # TRoiN=5000, - # ) - # ) - # print(f"Other config: {timepix.timepix_fly_client.get_other_config()}") - # # print(f"Pixel map from file: {timepix.timepix_fly_client.get_pixel_map()}") - # # PixelMap does not work yet - # pixel_map = PixelMap( - # chips=[ - # [{"i": 256 ^ 2 - 1, "p": [0, 1], "f": [0.5, 0.5]}], - # [{"i": 255 * 256, "p": [0, 1], "f": [0.5, 0.5]}], - # [{"i": 255, "p": [1, 2], "f": [0.5, 0.5]}], - # [{"i": 0, "p": [1, 2], "f": [0.5, 0.5]}], - # ] - # ) - - # timepix.timepix_fly_client.set_pixel_map(pixel_map) - # # test = timepix.timepix_fly_client.get_pixel_map() # TODO Throws an error at the moment - # # print(test) - # print(f"Initialized {timepix.name} with prefix {timepix.prefix}") - # timepix.timepix_fly_client.start() - # mock_server.start_acquisition() # Start the mock server acquisition - - # time.sleep( - # 1 - # ) # Wait for server stats to go to ready to start acquisition# Start acquisition, this will go through EPICS interface probably - - # timepix._data_server_thread_event.wait(timeout=1.0) - # timepix._data_server_thread_event.set() - # print(timepix._data_buffer) - # # raise RuntimeError("Stopping the server thread for testing purposes.") - # data = [el in timepix._data for el in timepix._data] - # # timepix._stop_data_server_thread() - # print("Data server thread stopped.") + msgs = [] + # for ii in range(runs): + for run in range(runs): + logger.warning(f"Starting trigger run {run + 1}/{runs}") + status = timepix.trigger() + logger.warning(f"Timepix triggered") + start_time = time.time() + while not status.done: + try: + status.wait(timeout=1) + except Exception as exc: + logger.warning(f"Trigger not done after ({time.time() - start_time:.2f}s)") + if time.time() - start_time > 20: + logger.warning("Breaking loop manually after 20 seconds of waiting.") + status.set_exception(f"Failed to complete trigger after 20 seconds") + break + n_messages = len(timepix._msg_dump) + logger.warning(f"Messages in Buffer is {n_messages}") + if n_messages > 0: + msg = timepix._msg_dump[-1] + logger.warning( + f"Last message had N start_frame : {msg.get('start_frame')}, N data_frames: {len(msg.get('data_frame'))}, N end_frame : {msg.get('end_frame')}" + ) + status = timepix.complete() + print("Waiting for timepix to complete.") + status.wait(timeout=10) + print("Timepix scan completed.") + timepix.unstage() + # timepix._msg_dump.clear() + print("Timepix unstaged.") + except Exception as e: + content = traceback.format_exc() + logger.error(f"An error occurred: {content}") + finally: + timepix.destroy() + print("Timepix destroyed.") diff --git a/superxas_bec/devices/timepix/timepix_fly_client/__init__.py b/superxas_bec/devices/timepix/timepix_fly_client/__init__.py index e69de29..d97a1b8 100644 --- a/superxas_bec/devices/timepix/timepix_fly_client/__init__.py +++ b/superxas_bec/devices/timepix/timepix_fly_client/__init__.py @@ -0,0 +1 @@ +from .timepix_fly_backend import TimepixFlyBackend diff --git a/superxas_bec/devices/timepix/timepix_fly_client/test_utils/__init__.py b/superxas_bec/devices/timepix/timepix_fly_client/test_utils/__init__.py new file mode 100644 index 0000000..41a1c45 --- /dev/null +++ b/superxas_bec/devices/timepix/timepix_fly_client/test_utils/__init__.py @@ -0,0 +1 @@ +from .timepix_fly_mock_server import TimePixFlyMockServer diff --git a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_mock_server.py b/superxas_bec/devices/timepix/timepix_fly_client/test_utils/timepix_fly_mock_server.py similarity index 85% rename from superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_mock_server.py rename to superxas_bec/devices/timepix/timepix_fly_client/test_utils/timepix_fly_mock_server.py index a1dd1dc..7bbc6d1 100644 --- a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_mock_server.py +++ b/superxas_bec/devices/timepix/timepix_fly_client/test_utils/timepix_fly_mock_server.py @@ -43,5 +43,7 @@ class TimePixFlyMockServer: try: requests.get(f"http://{self.host}:{self.port}/measurement/start", timeout=0.2) except requests.exceptions.RequestException: - pass # Ignore all exceptions as there is currently no return value for the request - self.add_log("Acquisition started on Timepix Fly mock server.") + self.add_log("Failed to start acquisition on Timepix Fly mock server.") + # Ignore all exceptions as there is currently no return value for the request + else: + self.add_log("Acquisition started on Timepix Fly mock server.") diff --git a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_backend.py b/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_backend.py new file mode 100644 index 0000000..9a89b66 --- /dev/null +++ b/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_backend.py @@ -0,0 +1,545 @@ +"""Implementation of the Timepix Fly Backend. It handles the communication +with the TimepixFly backend (https://github.com/paulscherrerinstitute/TimePixFly). +Please be aware that this was developed agains the 'dev' branch (2025/08/15). + +It communicates with the backend through a simple Client (TimepixFlyClient) +that handles the REST and WebSocket communication + callbacks, and provides +hooks for all the relevant ophyd interface, 'on_stage', +'on_trigger', 'on_complete', 'on_stop', etc.""" + +from __future__ import annotations + +import atexit +import json +import signal +import socket +import threading +import time +import traceback +import uuid +from typing import TYPE_CHECKING, Callable, Tuple + +from bec_lib.logger import bec_logger +from ophyd import StatusBase + +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client import ( + TimepixFlyClient, + TimePixFlyStatus, +) +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import ( + OtherConfigModel, + PixelMap, +) + +if TYPE_CHECKING: + from ophyd import DeviceStatus + + from superxas_bec.devices.timepix.timepix import Timepix + +logger = bec_logger.logger + +# pylint: disable=line-too-long +# pylint: disable=redefined-outer-name + + +class TimepixFlyBackendException(Exception): + """Custom exception for Timepix Fly Backend errors.""" + + +class TimepixFlyBackend: + """Timepix Fly Backend Device.""" + + def __init__(self, backend_rest_url: str, hostname: str | None = None, socket_port: int = 0): + """ + Initialize the Timepix Fly Backend device. + + Parameters: + backend_rest_url: The REST URL of the backend. + hostname: The hostname of the device, defaults to None, which means + socket.getfqdn() will be used to fetch hostname. It is recommended to specify + the hostname explicitly with domain name, e.g. 'x10da-bec-001.psi.ch' for use + at the beamline computers of SLS, or localhost for local testing of the backend. + socket_port: The socket port to use. Defaults to 0, + which lets the OS choose an available port. + """ + ws_url = f"{backend_rest_url}/ws" + self.timepix_fly_client = TimepixFlyClient(rest_url=backend_rest_url, ws_url=ws_url) + if hostname is None: + hostname = socket.getfqdn() + self.hostname = hostname + self.socket_port = socket_port # Use 0 as default to let the OS choose an available port + self.__msg_buffer = [] + self.callbacks: dict[str, Tuple[Callable[[dict, list[dict], dict, dict], None], dict]] = {} + self._status_objects: list[StatusBase] = [] + self._decoder = json.JSONDecoder() + self._socket_server: socket.socket | None = None + self._data_thread: threading.Thread | None = None + self._data_thread_shutdown_event = threading.Event() + atexit.register(self.on_destroy) # Ensure cleanup on exit + self.on_init() # TODO is this needed after registering the atexit handler? + + ################################################### + ###### Hooks for the PSIDeviceBase interface ###### + ################################################### + + def on_init(self): + """Called during initialization of the device.""" + try: + + def sigint_handler(*args): + """Hook SIGINT signals to on_destroy.""" + self.on_destroy() + + signal.signal(signal.SIGINT, sigint_handler) + signal.signal(signal.SIGTERM, sigint_handler) + except Exception: # pylint + logger.warning("Could not set signal handlers for SIGINT and SIGTERM.") + + def on_connected(self): + """Called if it is ensured that the device is connected.""" + self.timepix_fly_client.on_connected() + try: + status = self.start_data_server() + status.wait(timeout=5) + except Exception: + content = traceback.format_exc() + logger.error(f"Error starting data server: {content}") + # pylint: disable=raise-missing-from + raise TimepixFlyBackendException( + f"Could not start data server on {self.hostname}:{self.socket_port}. Please check log for detailed error message." + ) + + def on_stage(self, other_config: OtherConfigModel, pixel_map: PixelMap): + """ + Hook for on stage logic. + + Args: + other_config (OtherConfigModel): The configuration for the Timepix Fly detector. + pixel_map (PixelMap): The pixel map for the Timepix Fly detector. + """ + status = StatusBase() + self.cancel_on_stop(status) + self.timepix_fly_client.add_status_callback( + status, + success=[TimePixFlyStatus.CONFIG], + error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + ) + try: + status.wait(timeout=5.0) + except Exception: + content = traceback.format_exc() + logger.error( + f"Error while waiting for Timepix Fly backend to be in config state: {content}" + ) + # pylint: disable=raise-missing-from + raise TimeoutError( + "Timepix Fly backend state did not reach config state. Most likely a timeout error. Please check log for detailed error message." + ) + status = StatusBase() + + self.timepix_fly_client.add_status_callback( + status, + success=[TimePixFlyStatus.CONFIG], + error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + ) + logger.info(f"Setting other config, backend {other_config}") + self.timepix_fly_client.set_other_config(other_config) + self.timepix_fly_client.set_pixel_map(pixel_map) + + def on_pre_scan(self) -> None: + """Hook for on pre_scan logic.""" + + def on_trigger( + self, status: StatusBase | DeviceStatus | None = None + ) -> StatusBase | DeviceStatus: + """ + Hook for on_trigger logic. It adds a status callback based on the TimePixFlyStatus. + The backend needs to get into the AWAIT_CONNECTION state before starting the acquisition. + + Args: + status (StatusBase | DeviceStatus | None): The status object to track the operation. + If None, a new StatusBase object will be created. + Returns: + StatusBase | DeviceStatus: The status object that will be updated with the operation's result + """ + # TODO, could be removed as it's checkd from the top level! + if status is None: + status = StatusBase() + self.cancel_on_stop(status) + self.timepix_fly_client.add_status_callback( + status, + success=[TimePixFlyStatus.AWAIT_CONNECTION], + error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + ) + self.timepix_fly_client.start() + return status + + def on_trigger_finished( + self, status: StatusBase | DeviceStatus | None = None + ) -> StatusBase | DeviceStatus: + """ + Hook for on_trigger_finished logic. It adds a status callback based on the TimePixFlyStatus. + The backend needs to get into the CONFIG state again after a trigger is finished. + In practice, a full scan logic is happening during on trigger. + The status will be marked as finished/successful when the backend state + reaches CONFIG. If an exception state is reached, the status will be marked as failed. + + Args: + status (StatusBase | DeviceStatus | None): The status object to track the operation. + If None, a new StatusBase object will be created. + Returns: + StatusBase | DeviceStatus: The status object that will be updated with the operation's result + """ + if status is None: + status = StatusBase() + self.cancel_on_stop(status) + self.timepix_fly_client.add_status_callback( + status, + success=[TimePixFlyStatus.CONFIG], + error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + ) + return status + + def on_complete( + self, status: StatusBase | DeviceStatus | None = None + ) -> StatusBase | DeviceStatus: + """ + Hook for on_complete logic. It adds a status callback based on the TimePixFlyStatus. + The backend needs to get into the CONFIG state after a single acquisition. + + Args: + status (StatusBase | DeviceStatus | None): The status object to track the operation. + If None, a new StatusBase object will be created. + Returns: + StatusBase | DeviceStatus: The status object that will be updated with the operation's result + """ + if status is None: + status = StatusBase() + self.cancel_on_stop(status) + self.timepix_fly_client.add_status_callback( + status, + success=[TimePixFlyStatus.CONFIG], + error=[TimePixFlyStatus.EXCEPT, TimePixFlyStatus.SHUTDOWN], + ) + return status + + def on_unstage(self): + """Hook for on_unstage logic.""" + + def on_destroy(self): + """Hook for on_destroy logic.""" + self.timepix_fly_client.shutdown() + self._data_thread_shutdown_event.set() + if self._data_thread is not None and self._data_thread.is_alive(): + self._data_thread.join(timeout=1) # Allow the data thread to finish + if self._data_thread.is_alive(): + logger.error( + "Data thread poll loop of timepix_fly_backend did not stop within 1 second." + ) + if self._socket_server is not None: + try: + logger.info(f"Closing socket server on {self.hostname}:{self.socket_port}.") + self._socket_server.close() + # pylint: disable=broad-except + except Exception: + content = traceback.format_exc() + logger.error(f"Error closing socket server: {content}") + + def on_stop(self): + """Hook for on_stop logic.""" + self.stop_all_status_objects() + self.timepix_fly_client.stop_running_collection() + + #################################################### + ########## Custom Methods for the Backend ########## + #################################################### + + def cancel_on_stop(self, status: StatusBase): + """Cancel ongoing operations of a status object when the sto method is called.""" + self._status_objects.append(status) + + def stop_all_status_objects(self): + """Stop all status objects that are currently running.""" + for status in self._status_objects: + with status._lock: + if not status.done: + status.set_exception( + RuntimeError("Stop called on device, all status objects cancelled.") + ) + logger.info(f"Cancelled status object: {status}") + self._status_objects.clear() + + def add_callback(self, callback: callable, kwd: dict | None = None) -> str: + """ + Add a callback that will be executed whenever an acquisition is completed. This is + determind by receiving an EndFrame message from the backend. There will always be + a StartFrame message, follow by optional DataFrame messages, and finally + an EndFrame message. The callback will be called with the StartFrame, all DataFrames + and the EndFrame message as arguments, along with any additional keyword arguments + provided when registering the callback. + The callback signature needs to be: + def callback(start_frame: dict, data_frames: list[dict], end_frame: dict, kwd) -> None: + Args: + - start_frame (dict): The first message received, typically containing metadata. + - data_frames (list[dict]): A list of all data frames received during the acquisition. + - end_frame (dict): The last message received, typically containing the EndFrame type. + - any additional keyword arguments provided when registering the callback. + + Args: + callback (callable): The callback function to be called. The callback signature should be: + kwd (dict | None): Additional keyword arguments to pass to the callback, they will be unpacked + when calling the callback. If None, an empty dictionary will be used. + + Returns: + str: A unique identifier for the callback. + """ + if kwd is None: + kwd = {} + cb_id = uuid.uuid4() + self.callbacks[cb_id] = (callback, kwd) + logger.info(f"Callback {callback.__name__} added with UUID {cb_id}.") + return str(cb_id) + + def remove_callback(self, cb_id: str): + """ + Remove a callback by its unique identifier. + + Args: + cb_id (str): The unique identifier of the callback to remove. + """ + if cb_id in self.callbacks: + self.callbacks.pop(cb_id) + logger.info(f"Callback with UUID {cb_id} removed.") + else: + logger.warning(f"Callback with UUID {cb_id} not found.") + + def start_data_server(self) -> StatusBase: + """ + Start the data server to receive data from the Timepix Fly backend over a socket connection. + It will try to decypher the hostname through socket.getaddrinfo, and if multiple addresses + are found, it will use the first one. Please note that depending on the network configuration, + the hostname might not have the correct domain name attached, so it is recommended to specify + the hostname explicitly with domain name, e.g. 'x10da-bec-001.psi.ch'. + + The method creates a socket server that listens for incoming connections on the specified + hostname and port. It starts a thread that continuously receives data from the socket, + decodes the received JSON data, and processes it. The data is expected to be in JSON format, + with each message ending with a trailing byte "}\n". + + Returns: + StatusBase: A status object that indicates if the data server thread is ready to accept + connections. High level implementation should ensure that the data server is + started (status.wait(timeout=4)) before any data is sent from the backend. + """ + info = socket.getaddrinfo( + self.hostname, port=self.socket_port, family=socket.AF_INET, type=socket.SOCK_STREAM + ) + if len(info) == 0: + raise RuntimeError(f"Could not resolve hostname {self.hostname} for socket server.") + if len(info) > 1: + logger.info( + f"Multiple addresses found for {self.hostname}. Using the first one: {info[0]}" + ) + family, socktype, proto, _, sockaddr = info[0] + + self._socket_server = socket.create_server(sockaddr, family=family, backlog=1) + self._socket_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + # Set the hostname and socket_port to the ones that was picked by the socket.getaddrinfo + self.hostname, self.socket_port = self._socket_server.getsockname() + logger.info( + f"Socket server started on {self.hostname}:{self.socket_port}. Waiting for connections." + ) + + # Create status object to return for the high level implementations + status = StatusBase() + + if self._data_thread is None or not self._data_thread.is_alive(): + self._data_thread_shutdown_event.clear() + self._data_thread = threading.Thread( + target=self._receive_data_on_socket, kwargs={"status": status} + ) + self._data_thread.start() + else: + raise TimepixFlyBackendException( + "Data server thread is already running on timepix_fly_backend." + ) + + return status + + def _receive_data_on_socket(self, status: StatusBase): + """ + Background loop running in a thread, that receives data from the + timepix fly backend over socket_server. The backend reconnects for every acquisition (trigger), + to this socket. Therefore, it is important to handle all connections and disconnections properly. + + The buffer variable stores a string stream of received data. Whenever a trailing byte "}\n" is found, + in the buffer, the buffer is split into chunks of received data and each chunk is decoded + as a JSON object. The decoded objects are then processed, and if an EndFrame message is received, + the registered callbacks are executed with the StartFrame, all DataFrames, and the EndFrame message. + """ + buffer = "" + self._socket_server.settimeout( + 0.1 + ) # Set short socket timeout to avoid blocking the thread loop + status.set_finished() # Indicate that the socket server is ready to accept connections + while not self._data_thread_shutdown_event.is_set(): # Shutdown event + try: + # blocks until connected or timeout reached + conn, addr = self._socket_server.accept() + except socket.timeout: + continue # Timeout is okay, continue + except Exception: # pylint: disable=broad-except + # Log error, check if shutdown event is set. + # Shutdown event should be set before socket_server.close() is called. + content = traceback.format_exc() + logger.error(f"Error accepting connection: {content}") + continue + logger.debug(f"Connection accepted from {addr} for timepix_fly backend.") + # Clear the message buffer before entering the loop. + if self.__msg_buffer: + logger.warning(f"Found messages in msg_buffer: {self.__msg_buffer}") + self.__msg_buffer.clear() + conn.settimeout(0.1) # Set timeout for connection to avoid blocking in recv + with conn: + while not self._data_thread_shutdown_event.is_set(): + try: + # What if we split the chunk + chunk = conn.recv(4096) # Adjust buffer size as needed + except socket.timeout: + # Timeout is okay, continue in loop + continue + except Exception as e: # pylint: disable=broad-except + logger.error(f"Connection error: {e}. Closing connection.") + # conn = None #TODO should we reset conn? + break + if not chunk: + # Receiving an empty chunk means the connection was closed + # conn = None #TODO should we reset conn? + break + buffer += chunk.decode("utf-8") + + # Check if trailing byte "}\n" present in buffer + buffer_chunks = buffer.split("}\n") + for entry in buffer_chunks[:-1]: + # Process all complete JSON objects in the buffer + self._decode_received_data(entry + "}") + # Keep the last incomplete chunk. + # If the buffer ended with "}\n", this will be an empty string. + buffer = buffer_chunks[-1] + + def _decode_received_data(self, buffer: str) -> None: + """ + Decode the received data from the socket. + + Args: + buffer (str): The JSON string received from the socket. + """ + try: + obj, _ = self._decoder.raw_decode(buffer) + except json.JSONDecodeError: + logger.error(f"TimePixFlyBackend: Failed to decode JSON from buffer: {buffer}") + return # TODO should this raise, or only log error as of now? + + self.__msg_buffer.append(obj) + if obj.get("type", "") == "EndFrame": + try: + # If the EndFrame message is received, run the callbacks + logger.info(f"Running callbacks") + self.run_msg_callbacks() + except Exception: # pylint: disable=broad-except + content = traceback.format_exc() + logger.error(f"Error in msg callbacks with error msg: {content}") + msgs_in_buffer = "".join( + [f"{msg['type']} with keys {msg.keys()} \n" for msg in self.__msg_buffer] + ) + logger.debug(f"TimePixFlyBackend: Messages in buffer: {msgs_in_buffer}") + finally: # Make sure to always reset the message buffer after processing + logger.debug( + "TimePixFlyBackend: Resetting message buffer after processing EndFrame message." + ) + logger.debug(f"Messages in buffer: {len(self.__msg_buffer)}") + self.__msg_buffer.clear() + + def run_msg_callbacks(self): + """Run callbacks if EndFrame message is received.""" + # TODO + start_frame = self.__msg_buffer[0] + end_frame = self.__msg_buffer[-1] + data_frames = self.__msg_buffer[1:-1] + logger.info(f"Number of callbacks {len(self.callbacks.keys())}") + for cb, kwd in self.callbacks.values(): + cb(start_frame, data_frames, end_frame, **kwd) + + +if __name__ == "__main__": # pragma: no cover + import time + + from superxas_bec.devices.timepix.timepix_fly_client.test_utils.timepix_fly_mock_server import ( + TimePixFlyMockServer, + ) + from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import ( + TimepixEndFrame, + TimepixStartFrame, + TimepixXESFrame, + ) + + mock_server = TimePixFlyMockServer() + timepix = TimepixFlyBackend( + backend_rest_url="localhost:8452", hostname="localhost", socket_port=3031 + ) + + start_frames = {} + xes_frames = {} + end_frames = {} + + def add_msg_callback(start_frame, data_frames, end_frame, **kwargs): + """Callback to print received messages.""" + counter = len(start_frames) + start_frames[counter] = TimepixStartFrame(**start_frame) + xes_frames[counter] = [TimepixXESFrame(**data_frame) for data_frame in data_frames] + end_frames[counter] = TimepixEndFrame(**end_frame) + + try: + + print("TimepixFlyBackend initialized.") + timepix.on_connected() + print("TimepixFlyBackend connected.") + # Parse scan info for OtherConfig + config = OtherConfigModel( + output_uri=f"tcp:{timepix.hostname}:{timepix.socket_port}", TRoiStep=1, TRoiN=5000 + ) + # Parse pixel map from scan info if needed, otherwise use some default pixel map. + pixel_map = PixelMap( + chips=[ + [{"i": 256 ^ 2 - 1, "p": [0, 1], "f": [0.5, 0.5]}], + [{"i": 255 * 256, "p": [0, 1], "f": [0.5, 0.5]}], + [{"i": 255, "p": [1, 2], "f": [0.5, 0.5]}], + [{"i": 0, "p": [1, 2], "f": [0.5, 0.5]}], + ] + ) + timepix.add_callback(add_msg_callback) + timepix.on_stage(other_config=config, pixel_map=pixel_map) + print("TimepixFlyBackend staged with configuration and pixel map.") + for ii in range(5): + print(f"Starting scan {ii + 1}...;") + time.sleep(1) + status_1 = timepix.on_trigger() + # print("TimepixFlyBackend pre-scan started.") + status_1.wait(timeout=10) + mock_server.start_acquisition() + status_2 = timepix.on_trigger_finished() + status_2.wait(timeout=10) + # print("Acquisition started on mock server.") + + print("TimepixFlyBackend scan completed.") + status = timepix.on_complete() + status.wait(timeout=10) + print( + f"Received {len(start_frames)} start frames, {len(xes_frames)} data frames, and {len(end_frames)} end frames." + ) + except Exception as e: + logger.error(f"Error during TimepixFlyBackend operation: {e}") + finally: + timepix.on_destroy() + print("TimepixFlyBackend destroyed.") diff --git a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_client.py b/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_client.py index b2d1743..2e96604 100644 --- a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_client.py +++ b/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_client.py @@ -1,12 +1,14 @@ """ -This module implements a python client ot the REST interface of the tpx3app. +Module that implements a python client interface to the TimePix Fly tpx3app REST API, +and connects to the TimePix Fly WebSocket server to receive status updates. +It provides methods to start, stop, and configure the TimePix detector, +as well as to retrieve pixel maps, other configuration parameters, and +the current state of the detector. """ from __future__ import annotations -import copy import enum -import json import threading import time import traceback @@ -14,7 +16,7 @@ from typing import Any, Type import requests from bec_lib.logger import bec_logger -from ophyd import Device, StatusBase +from ophyd import StatusBase from websockets import State from websockets.exceptions import WebSocketException from websockets.sync.client import ClientConnection, connect @@ -32,11 +34,8 @@ from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface impor logger = bec_logger.logger -# TODO remove! -SERVER_ADDRESS = "localhost:8452" # Default server address for TimePix REST API - -# pylint: disable=arguments-differ +# pylint: disable=line-too-long class TimePixStatusError(Exception): """Exception raised when the TimePix detector status was in an unexpected state.""" @@ -58,21 +57,22 @@ class TimePixFlyStatus(str, enum.Enum): class TimepixFlyClient: """ - A client for the TimePix detector REST API. - This class is used to interact with the TimePix detector via its REST API. - It provides methods to send GET and PUT requests to the TimePix server. + A client for the TimePix fly backend (tpx3app). + It exposes methods to interact with REST endpoints + and allows to connect callbacks to status objects from ophyd + that allow to dynamically update based on the state of the backend. """ - def __init__(self, rest_url: str, ws_url: str, parent: Device): + def __init__(self, rest_url: str, ws_url: str): """ Initialize the TimePixFlyClient with a server address. Args: - server_address (str): The address of the TimePix REST API: "tpx3app". + rest_url (str): The REST API URL for the TimePix Fly backend, e.g., "localhost:8452". + ws_url (str): The WebSocket URL for the TimePix Fly backend, e.g., "localhost:8452/ws". """ self.rest_url = rest_url self.ws_url = ws_url - self.parent = parent self.ws_client: ClientConnection | None = None self._rlock = threading.RLock() self._timeout = 5 # Default timeout for requests @@ -98,21 +98,28 @@ class TimepixFlyClient: self.connect() self.wait_for_connection(timeout=5) - except Exception as e: + except Exception: + content = traceback.format_exc() logger.error( - f"Error while checking the state of the TimePix server: {e}. " + f"Error while checking the state of the TimePix server: {content}. " f"Please check the server address and ensure the server is running." ) - raise e + # pylint: disable=raise-missing-from + raise ConnectionError( + f"TimePix Fly client failed to connect to {self.rest_url}. Please check logs for detailed error." + ) def stop_running_collection(self): """ - Resets the TimePix backend to the configuration state. + Resets the TimePix backend to the configuration state. We check if the backend + is in the CONFIG state, and if it is not, we stop the current collection. + We check in addition that the client has not been started before via start() + REST API call. stop_collect() will reset the flag _started to False. """ state = self.state() if state.state != TimePixFlyStatus.CONFIG or self._started is True: logger.info( - f"Resetting TimePix backend to config state. Current state: {state.state}, started: {self._started}" + f"Stopping running collection on TimePix backend, current state: {state.state}, was started: {self._started}" ) self.stop_collect() @@ -132,61 +139,90 @@ class TimepixFlyClient: return self._status def add_status_callback( - self, status: StatusBase, success: list[TimePixFlyStatus], error: list[TimePixFlyStatus] + self, + status: StatusBase, + success: list[TimePixFlyStatus], + error: list[TimePixFlyStatus], + run: bool = True, ): """ Add a StatusBase callback for the TimePix detector. The status will be updated when the detector status changes and set to finished when the status matches one of the specified success statuses and to exception when the status matches one of the specified error statuses. + Per default, the callback will immediately check and run if the status is already in success. + Args: status (StatusBase): StatusBase object success (list[StdDaqStatus]): list of statuses that indicate success error (list[StdDaqStatus]): list of statuses that indicate error + run (bool): If True, the callback will be run immediately if the status is already in success. + If False, the callback will not be run immediately. """ - if self.status in success: - status.set_finished() - return + if run is True: + if self.status in success: + status.set_finished() + return + if self.status in error: + last_error = self.last_error() + status.set_exception( + TimePixStatusError( + f"TimePixFly status is '{self.status.value},' last error message: {last_error.message}" + ) + ) + return self._status_callbacks[id(status)] = (status, success, error) def connect(self): - """ - Connect to the TimePix WebSocket server. - """ + """Connect to the TimePix WebSocket server.""" if self._ws_update_thread is not None and self._ws_update_thread.is_alive(): return - self._ws_update_thread = threading.Thread( - target=self._ws_update_loop, name=f"{self.parent.name}_stddaq_ws_loop", daemon=True - ) + self._ws_update_thread = threading.Thread(target=self._ws_update_loop, daemon=True) self._ws_update_thread.start() - def wait_for_connection(self, timeout: float = 10) -> None: + # pylint: disable=raise-missing-from + def wait_for_connection(self, timeout: float = 6) -> None: """ - Wait for the connection to the StdDAQ to be established. + Wait for the connection to the TimepixFly client connection to be established. Args: timeout (float): timeout for the request """ - start_time = time.time() - while True: - if self.ws_client is not None and self.ws_client.state == State.OPEN: - return - try: - self.ws_client = connect(f"ws://{self.ws_url}") - break - except ConnectionRefusedError as exc: - if time.time() - start_time > timeout: - raise TimeoutError("Timeout while waiting for connection to StdDAQ") from exc - time.sleep(0.5) # Try to reconnect every 2 seconds + logger.info( + f"Attempting to connect to TimePixFly WebSocket at {self.ws_url}, with timeout {timeout} seconds." + ) + with self._rlock: + start_time = time.time() + while True: + if self.ws_client is not None and self.ws_client.state == State.OPEN: + return + try: + self.ws_client = connect(f"ws://{self.ws_url}") + break + except ConnectionRefusedError: + if time.time() - start_time > timeout: + content = traceback.format_exc() + logger.error(f"Connection timed out: {content}") + raise TimeoutError( + f"Timeout while waiting for connection to TimePixFly WebSocket server on {self.ws_url}" + ) + except Exception: + content = traceback.format_exc() + logger.error( + f"Failed to connect to TimePixFly WebSocket server on {self.ws_url}: {content}" + ) + raise ConnectionError( + f"Failed to connect to TimePixFly WebSocket server on {self.ws_url} with error: {content}" + ) + time.sleep(0.5) # Try to reconnect every 0.5 seconds def _ws_update_loop(self): - """ - Loop to update the status property of the StdDAQ. - """ + """Websocket update loop, runs in background thread.""" while not self._shutdown_event.is_set(): self._ws_send_and_receive() def _ws_send_and_receive(self): + """Receive messages from the TimePixFly WebSocket server.""" if not self.ws_client: self.wait_for_connection() try: @@ -207,11 +243,11 @@ class TimepixFlyClient: Handle a message received from the StdDAQ. """ try: - logger.info(f"Received message from TimePixFly: {msg}") self._status = TimePixFlyStatus(msg) - except Exception: + logger.info(f"Received TimepixFly status: {self._status.value}") + except Exception: # pylint: disable=broad-except content = traceback.format_exc() - logger.warning(f"Failed to decode websocket message: {content}") + logger.error(f"Failed to decode websocket message: {content}") return self._run_status_callbacks() @@ -222,42 +258,39 @@ class TimepixFlyClient: or exception, respectively and removed from the list of callbacks. """ status = self._status - completed_callbacks = [] - logger.info(f"Running status callbacks for TimePixFly status: {status.value}") + logger.warning(f"Running status callbacks for status: {status.value}") callback_ids = list(self._status_callbacks.keys()) for cb_id in callback_ids: dev_status, success, error = self._status_callbacks[cb_id] - logger.info(f"Checking status callback {cb_id} for TimePixFly status: {status.value}") - if dev_status.done: - logger.warning("Status object already resolved. Skipping Timepix callback.") - self._status_callbacks.pop(cb_id) - continue - if status in success: - dev_status.set_finished() - logger.info(f"Timepix status in succes is {status.value}") - self._status_callbacks.pop(cb_id) - elif status in error: - logger.warning(f"Timepix status in error is {status.value}") - dev_status.set_exception( - TimePixStatusError( - f"TimePixStatus status is '{status.value},' last error: {self.last_error()}" + with dev_status._lock: + if dev_status.done: + self._status_callbacks.pop(cb_id) + continue + if status in success: + dev_status.set_finished() + logger.info(f"Status callback finished in succes: {status.value}") + self._status_callbacks.pop(cb_id) + elif status in error: + last_error = self.last_error() + logger.error( + f"Timepix status in error is {status.value}, with last error: {last_error.message}" ) - ) - self._status_callbacks.pop(cb_id) - # If config is reached, started can be reset to False + dev_status.set_exception( + TimePixStatusError( + f"TimePixStatus status is '{status.value},' last error message: {last_error.message}" + ) + ) + self._status_callbacks.pop(cb_id) + # Reset the _started flag if the status is in CONFIG. if status == TimePixFlyStatus.CONFIG: - self._started = False + self._started = False # Should this be made thread-safe? def shutdown(self): - """ - Shutdown the StdDAQ client. - """ + """Shutdown the TimepixFlyClient client.""" self._shutdown_event.set() - if self._ws_update_thread is not None: - self._ws_update_thread.join() if self.ws_client is not None: self.ws_client.close() - self.ws_client = None + self.ws_client = None ############################ ##### REST API Methods ##### @@ -276,6 +309,7 @@ class TimepixFlyClient: Returns: Any: The parsed response if a model is provided, else the raw response. """ + logger.debug(f"Sending GET request to TimePix server: {get_cmd}") response = requests.get(f"http://{self.rest_url}/{get_cmd}", timeout=self._timeout) response.raise_for_status() # Raise an error for bad responses if get_response_model is not None: @@ -301,6 +335,7 @@ class TimepixFlyClient: Returns: Any: The parsed response if a model is provided, else None. """ + logger.debug(f"Sending PUT request to TimePix server: {put_cmd} with value: {value}") response = requests.put( f"http://{self.rest_url}/{put_cmd}", json=value, timeout=self._timeout ) @@ -313,6 +348,7 @@ class TimepixFlyClient: Start the TimePix detector by sending a GET request to the start endpoint. This method is a wrapper around the REST API call to start the detector. """ + logger.debug(f"Start called from client") self._get(get_cmd="?start=true") self._started = True @@ -449,7 +485,7 @@ class TimepixFlyClient: ) self._put(put_cmd="other-config", value=other_config.model_dump(), put_response_model=None) - def get_net_addresses(self) -> dict[str, str]: + def get_net_addresses(self) -> NetAddresses: """ Get the network addresses of the TimePix detector by sending a GET request to the net-addresses endpoint. diff --git a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_interface.py b/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_interface.py index 0f5debb..3992ce2 100644 --- a/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_interface.py +++ b/superxas_bec/devices/timepix/timepix_fly_client/timepix_fly_interface.py @@ -8,6 +8,8 @@ from typing import Literal from pydantic import BaseModel, Field +# pylint: disable=line-too-long + class TimePixResponse(BaseModel): """Base model for TimePix responses.""" @@ -109,16 +111,37 @@ class PixelMap(TimePixResponse): chips: list[list[dict[Literal["i", "p", "f"], int | float | list[int | float]]]] +# For efficiency, we do not arse the responses into Pydantic models, but use the dict from +# json directly. Nevertheless, we define the models here to have a common interface +# and to be able to use them in the future if needed. class TimepixStartFrame(TimePixResponse): - pass + """TimepixStartFrame is a Pydantic model that represents the start frame of a TimePix acquisition.""" + + type: str = "StartFrame" + Mode: Literal["TOA"] + TRoiStart: int + TRoiStep: int + TRoiN: int + NumEnergyPoints: int + save_interval: int -class TimepixDataFrame(TimePixResponse): - pass +class TimepixXESFrame(TimePixResponse): + """TimepixXESFrame is a Pydantic model that represents a data frame from the TimePix detector.""" + + type: str = "XesData" + period: int + TDSpectra: list[float] + totalEvents: int + beforeROI: int + afterROI: int class TimepixEndFrame(TimePixResponse): - pass + """TimepixEndFrame is a Pydantic model that represents the end frame of a TimePix acquisition.""" + + type: str = "EndFrame" + error: str # Habe ein GET /net-addresses call implementiert diff --git a/superxas_bec/devices/timepix/utils.py b/superxas_bec/devices/timepix/utils.py new file mode 100644 index 0000000..1ec0f49 --- /dev/null +++ b/superxas_bec/devices/timepix/utils.py @@ -0,0 +1,382 @@ +"""Temporary utility module for Status Object implementations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ophyd import Device, DeviceStatus, StatusBase + + +class AndStatusWithList(DeviceStatus): + """ + Custom implementation of the AndStatus that combines the + option to add multiple statuses as a list, and in addition + allows for adding the Device as an object to access its + methods. + + Args""" + + def __init__( + self, + device: Device, + status_list: StatusBase | DeviceStatus | list[StatusBase | DeviceStatus], + **kwargs, + ): + self.all_statuses = status_list if isinstance(status_list, list) else [status_list] + super().__init__(device=device, **kwargs) + self._trace_attributes["all"] = [st._trace_attributes for st in self.all_statuses] + + def inner(status): + with self._lock: + if self._externally_initiated_completion: + return + if self.done: # Return if status is already done.. It must be resolved already + return + + for st in self.all_statuses: + with st._lock: + if st.done and not st.success: + self.set_exception(st.exception()) # st._exception + return + + if all(st.done for st in self.all_statuses) and all( + st.success for st in self.all_statuses + ): + self.set_finished() + + for st in self.all_statuses: + with st._lock: + st.add_callback(inner) + + # TODO improve __repr__ and __str__ + def __repr__(self): + return "".format(self=self) + + def __str__(self): + return "".format(self=self) + + def __contains__(self, status: StatusBase | DeviceStatus) -> bool: + for child in self.all_statuses: + if child == status: + return True + if isinstance(child, AndStatusWithList): + if status in child: + return True + + return False + + # # TODO Check if this actually works.... + # def set_exception(self, exc): + # # Propagate the exception to all sub-statuses that are not done yet. + # + # with self._lock: + # if self._externally_initiated_completion: + # return + # if self.done: # Return if status is already done.. It must be resolved already + # return + # super().set_exception(exc) + # for st in self.all_statuses: + # with st._lock: + # if not st.done: + # st.set_exception(exc) + + def _run_callbacks(self): + """ + Set the Event and run the callbacks. + """ + if self.timeout is None: + timeout = None + else: + timeout = self.timeout + self.settle_time + if not self._settled_event.wait(timeout): + self.log.warning("%r has timed out", self) + with self._externally_initiated_completion_lock: + if self._exception is None: + exc = TimeoutError( + f"AndStatus from device {self.device.name} failed to complete in specified timeout of {self.timeout + self.settle_time}." + ) + self._exception = exc + # Mark this as "settled". + try: + self._settled() + except Exception: + self.log.exception("%r encountered error during _settled()", self) + with self._lock: + self._event.set() + if self._exception is not None: + try: + self._handle_failure() + except Exception: + self.log.exception("%r encountered an error during _handle_failure()", self) + for cb in self._callbacks: + try: + cb(self) + except Exception: + self.log.exception( + "An error was raised on a background thread while " + "running the callback %r(%r).", + cb, + self, + ) + self._callbacks.clear() + + +class AndStatus(StatusBase): + """Custom AndStatus for TimePix detector.""" + + def __init__( + self, + left: StatusBase | DeviceStatus | list[StatusBase | DeviceStatus] | None, + name: str | Device | None = None, + right: StatusBase | DeviceStatus | list[StatusBase | DeviceStatus] | None = None, + **kwargs, + ): + self.left = left if isinstance(left, list) else [left] + if right is not None: + self.right = right if isinstance(right, list) else [right] + else: + self.right = [] + self.all_statuses = self.left + self.right + if name is None: + name = "unname_status" + elif isinstance(name, Device): + name = name.name + else: + name = name + self.name = name + super().__init__(**kwargs) + self._trace_attributes["left"] = [st._trace_attributes for st in self.left] + self._trace_attributes["right"] = [st._trace_attributes for st in self.right] + + def inner(status): + with self._lock: + if self._externally_initiated_completion: + return + if self.done: # Return if status is already done.. It must be resolved already + return + + for st in self.all_statuses: + with st._lock: + if st.done and not st.success: + self.set_exception(st.exception()) # st._exception + return + + if all(st.done for st in self.all_statuses) and all( + st.success for st in self.all_statuses + ): + self.set_finished() + + for st in self.all_statuses: + with st._lock: + st.add_callback(inner) + + def __repr__(self): + return "({self.left!r} & {self.right!r})".format(self=self) + + def __str__(self): + return "{0}(done={1.done}, " "success={1.success})" "".format(self.__class__.__name__, self) + + def __contains__(self, status: StatusBase) -> bool: + for child in [self.left, self.right]: + if child == status: + return True + if isinstance(child, AndStatus): + if status in child: + return True + + return False + + def _run_callbacks(self): + """ + Set the Event and run the callbacks. + """ + if self.timeout is None: + timeout = None + else: + timeout = self.timeout + self.settle_time + if not self._settled_event.wait(timeout): + # We have timed out. It's possible that set_finished() has already + # been called but we got here before the settle_time timer expired. + # And it's possible that in this space be between the above + # statement timing out grabbing the lock just below, + # set_exception(exc) has been called. Both of these possibilties + # are accounted for. + self.log.warning("%r has timed out", self) + with self._externally_initiated_completion_lock: + # Set the exception and mark the Status as done, unless + # set_exception(exc) was called externally before we grabbed + # the lock. + if self._exception is None: + exc = TimeoutError( + f"Status with name {self.name} failed to complete in specified timeout of {self.timeout + self.settle_time}." + ) + self._exception = exc + # Mark this as "settled". + try: + self._settled() + except Exception: + # No alternative but to log this. We can't supersede set_exception, + # and we have to continue and run the callbacks. + self.log.exception("%r encountered error during _settled()", self) + # Now we know whether or not we have succeed or failed, either by + # timeout above or by set_exception(exc), so we can set the Event that + # will mark this Status as done. + with self._lock: + self._event.set() + if self._exception is not None: + try: + self._handle_failure() + except Exception: + self.log.exception("%r encountered an error during _handle_failure()", self) + # The callbacks have access to self, from which they can distinguish + # success or failure. + for cb in self._callbacks: + try: + cb(self) + except Exception: + self.log.exception( + "An error was raised on a background thread while " + "running the callback %r(%r).", + cb, + self, + ) + self._callbacks.clear() + + +# from __future__ import annotations + +# from collections import defaultdict +# from typing import Dict, List, Tuple + +# import numpy as np + +# ROI = List[Tuple[float, float]] + + +# def order_roi_corners_simple(roi: ROI) -> np.ndarray: +# """Order ROI corners as [top-left, top-right, bottom-right, bottom-left].""" +# pts = np.array(roi, dtype=float) +# cx, cy = pts.mean(axis=0) +# angles = np.arctan2(pts[:, 1] - cy, pts[:, 0] - cx) +# idx = np.argsort(angles) +# ordered = pts[idx] +# # Ensure clockwise order +# if np.cross(ordered[1] - ordered[0], ordered[2] - ordered[0]) < 0: +# ordered = ordered[::-1] +# return ordered[:4] + + +# def compute_affine_transform( +# src: np.ndarray, dst: np.ndarray, preserve_scale: bool = True +# ) -> np.ndarray: +# """Compute affine transform mapping src -> dst. Optionally preserve pixel scale.""" +# if preserve_scale: +# # Solve for rotation+translation only +# A = np.array( +# [ +# [src[0, 0], -src[0, 1], 1, 0], +# [src[0, 1], src[0, 0], 0, 1], +# [src[1, 0], -src[1, 1], 1, 0], +# [src[1, 1], src[1, 0], 0, 1], +# ] +# ) +# b = dst[:2].ravel() +# x, residuals, _, _ = np.linalg.lstsq(A, b, rcond=None) +# a, b_, tx, ty = x +# return np.array([[a, -b_, tx], [b_, a, ty]]) +# else: +# # Full affine transform +# src_h = np.hstack([src, np.ones((4, 1))]) +# dst_h = dst +# M, _, _, _ = np.linalg.lstsq(src_h, dst_h, rcond=None) +# return M.T + + +# def apply_affine_transform(coords: np.ndarray, affine: np.ndarray) -> np.ndarray: +# """Apply affine transform to coordinates.""" +# coords_h = np.hstack([coords, np.ones((coords.shape[0], 1))]) +# transformed = coords_h @ affine.T +# return transformed[:, :2] + + +# def roi_pixel_hits( +# image_shape: Tuple[int, int], roi: ROI, start_idx: int = 0, min_fraction_diff: float = 0.0 +# ) -> List[Dict]: +# """ +# For each ROI, return list of hits as dicts {'i': (x,y), 'p': row_idx, 'f': fraction}. +# Supports rotated rectangles using bilinear fraction splitting. +# """ +# hits_dict: Dict[Tuple[int, int], Dict[str, list]] = defaultdict(lambda: {"p": [], "f": []}) + +# corners = order_roi_corners_simple(roi) +# height = int(np.linalg.norm(corners[0] - corners[3])) + 1 +# width = int(np.linalg.norm(corners[0] - corners[1])) + 1 + +# dst = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]], dtype=float) +# affine_mat = compute_affine_transform(corners, dst) # 2x3 + +# # Bounding box in image +# min_x = max(int(np.floor(corners[:, 0].min())), 0) +# max_x = min(int(np.ceil(corners[:, 0].max())), image_shape[1] - 1) +# min_y = max(int(np.floor(corners[:, 1].min())), 0) +# max_y = min(int(np.ceil(corners[:, 1].max())), image_shape[0] - 1) + +# yy, xx = np.meshgrid(np.arange(min_y, max_y + 1), np.arange(min_x, max_x + 1), indexing="ij") +# coords = np.stack([xx.ravel(), yy.ravel()], axis=1) # N x 2 + +# local_coords = apply_affine_transform(coords, affine_mat) +# x_local = local_coords[:, 0] +# y_local = local_coords[:, 1] + +# # Keep pixels inside ROI rectangle +# mask = (x_local >= 0) & (x_local <= width - 1) & (y_local >= 0) & (y_local <= height - 1) +# coords_in = coords[mask] +# x_in = x_local[mask] +# y_in = y_local[mask] + +# # Bilinear fractions +# x0 = np.floor(x_in).astype(int) +# y0 = np.floor(y_in).astype(int) +# dx = x_in - x0 +# dy = y_in - y0 + +# for coord, x0i, y0i, dxv, dyv in zip(coords_in, x0, y0, dx, dy): +# row_base = start_idx + +# # Contributions to 4 neighboring "rows" +# contributions = [ +# (row_base + y0i, (1 - dxv) * (1 - dyv)), +# (row_base + y0i, dxv * (1 - dyv)), +# (row_base + y0i + 1, (1 - dxv) * dyv), +# (row_base + y0i + 1, dxv * dyv), +# ] + +# # Filter negligible contributions +# contributions = [(p, f) for p, f in contributions if f >= min_fraction_diff] + +# # Normalize fractions to sum 1 +# if contributions: +# total_f = sum(f for _, f in contributions) +# contributions = [(p, f / total_f) for p, f in contributions] + +# for p, f in contributions: +# key = (int(coord[0]), int(coord[1])) +# hits_dict[key]["p"].append(int(p)) +# hits_dict[key]["f"].append(float(f)) + +# hits_roi = [{"i": key, "p": value["p"], "f": value["f"]} for key, value in hits_dict.items()] +# return hits_roi + + +# if __name__ == "__main__": +# image_shape = (512, 512) +# rois = [ +# [(25, 25), (50, 50), (50, 25), (75, 50)], # rotated 45 degrees rectangle +# [(300, 300), (400, 300), (400, 400), (300, 400)], # upright rectangle +# ] + +# hits_0 = roi_pixel_hits(image_shape, rois[0], min_fraction_diff=0.1) +# hits_1 = roi_pixel_hits(image_shape, rois[1], start_idx=10, min_fraction_diff=0.2) + +# print(hits_0[:5]) +# print(hits_1[:5]) diff --git a/tests/tests_devices/test_timepix_fly_backend.py b/tests/tests_devices/test_timepix_fly_backend.py new file mode 100644 index 0000000..d1b6671 --- /dev/null +++ b/tests/tests_devices/test_timepix_fly_backend.py @@ -0,0 +1,15 @@ +"""This module tests the Timepix Fly backend functionality.""" + +from __future__ import annotations + +import pytest + +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_backend import TimepixFlyBackend + + +@pytest.fixture(scope="function") +def timepix_fly_backend(): + """Fixture for creating a Timepix Fly backend instance.""" + backend = TimepixFlyBackend(backend_rest_url="http://localhost:8000") + yield backend + backend.on_destroy() diff --git a/tests/tests_devices/test_timepix_fly_client.py b/tests/tests_devices/test_timepix_fly_client.py new file mode 100644 index 0000000..d9b113d --- /dev/null +++ b/tests/tests_devices/test_timepix_fly_client.py @@ -0,0 +1,276 @@ +"""Module to test the Timepix Fly client functionality.""" + +from __future__ import annotations + +from unittest import mock + +import pytest +from ophyd import StatusBase +from websockets import State + +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client import ( + TimepixFlyClient, + TimePixFlyStatus, + TimePixStatusError, +) +from superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_interface import ProgramState + + +@pytest.fixture(scope="function") +def timepix_fly_client(): + """Fixture for creating a Timepix Fly client instance.""" + client = TimepixFlyClient(rest_url="http://localhost:8000", ws_url="ws://localhost:8000/ws") + try: + yield client + finally: + client.shutdown() + + +@pytest.mark.parametrize( + "return_state", + [ + ProgramState(state=TimePixFlyStatus.CONFIG), + ProgramState(state=TimePixFlyStatus.COLLECT), + ProgramState(state=TimePixFlyStatus.SETUP), + ], +) +def test_timepix_fly_client_stop_running_collection(timepix_fly_client, return_state): + """Test the on_connected method of the Timepix Fly client.""" + with ( + mock.patch.object(timepix_fly_client, "stop_collect") as mock_stop_collect, + mock.patch.object(timepix_fly_client, "state", return_value=return_state), + ): + timepix_fly_client.stop_running_collection() + + if return_state.state == TimePixFlyStatus.CONFIG: + assert mock_stop_collect.call_count == 0, "Stop collect should be called once." + timepix_fly_client._started = True + timepix_fly_client.stop_running_collection() + assert mock_stop_collect.call_count == 1, "Stop collect should not be called again." + else: + assert mock_stop_collect.call_count == 1, "Stop collect should be called once." + timepix_fly_client._started = True + timepix_fly_client.stop_running_collection() + assert mock_stop_collect.call_count == 2, "Stop collect should be called in CONFIG." + + +def test_timepix_fly_client_on_connected(timepix_fly_client): + """ + Test timepix fly client connect method. + + This simply ensures that all methods are called. They are tested separately. + """ + with ( + mock.patch.object(timepix_fly_client, "stop_running_collection") as mock_stop_collection, + mock.patch.object(timepix_fly_client, "connect") as mock_connect, + mock.patch.object(timepix_fly_client, "wait_for_connection") as mock_wait_for_connection, + ): + timepix_fly_client.on_connected() + mock_stop_collection.assert_called_once() + mock_connect.assert_called_once() + mock_wait_for_connection.assert_called_once() + + +def test_timepix_fly_client_connect(timepix_fly_client): + """This tests the connect method of timepix fly client.""" + module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client" + # Patch Thread so we don't actually start the background loop + with mock.patch(f"{module_path}.threading.Thread") as mock_thread_cls: + mock_thread = mock.Mock() + mock_thread_cls.return_value = mock_thread + + timepix_fly_client.connect() + + # Thread should be created with the update loop as target and daemon True + mock_thread_cls.assert_called_once() + # start() must be called on the created thread + mock_thread.start.assert_called_once() + + +def test_timepix_fly_client_wait_for_connection(timepix_fly_client): + """This tests the wait_for_connection method of timepix fly client.""" + module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client" + # Case 1: ws_client already present and OPEN + mock_client = mock.Mock() + type(mock_client).state = mock.PropertyMock(return_value=State.OPEN) + timepix_fly_client.ws_client = mock_client + + # Should return immediately + timepix_fly_client.wait_for_connection(timeout=0.1) + + # Case 2: connect() establishes the connection + timepix_fly_client.ws_client = None + mock_client2 = mock.Mock() + type(mock_client2).state = mock.PropertyMock(return_value=State.OPEN) + with mock.patch(f"{module_path}.connect", return_value=mock_client2) as mock_connect: + timepix_fly_client.wait_for_connection(timeout=0.1) + mock_connect.assert_called_once() + + +def test_timepix_fly_client_ws_send_and_received(timepix_fly_client): + """This tests the _ws_send_and_receive method of timepix fly client.""" + # Prepare a mock ws_client where recv returns a message + mock_ws = mock.Mock() + mock_ws.recv.return_value = "init" + timepix_fly_client.ws_client = mock_ws + + # Patch _on_received_ws_message to ensure it's invoked + with mock.patch.object(timepix_fly_client, "_on_received_ws_message") as mock_on_msg: + timepix_fly_client._ws_send_and_receive() + mock_on_msg.assert_called_once_with("init") + + # Now simulate a TimeoutError from recv: should simply return and not call _on_received_ws_message + mock_ws2 = mock.Mock() + mock_ws2.recv.side_effect = TimeoutError + timepix_fly_client.ws_client = mock_ws2 + with mock.patch.object(timepix_fly_client, "_on_received_ws_message") as mock_on_msg2: + timepix_fly_client._ws_send_and_receive() + mock_on_msg2.assert_not_called() + + +def test_timepix_fly_client_on_message_received(timepix_fly_client): + """This tests the _on_received_ws_message method of timepix fly client.""" + with mock.patch.object(timepix_fly_client, "_run_status_callbacks") as mock_run_callbacks: + timepix_fly_client._on_received_ws_message("init") + assert timepix_fly_client._status == TimePixFlyStatus.INIT + mock_run_callbacks.assert_called_once() + + # invalid message should not change status or call callbacks + prev_status = timepix_fly_client._status + with mock.patch.object(timepix_fly_client, "_run_status_callbacks") as mock_run_callbacks2: + timepix_fly_client._on_received_ws_message("invalid_status_string") + # status stays as previous value and callbacks not called + assert timepix_fly_client._status == prev_status + mock_run_callbacks2.assert_not_called() + + +def test_timepix_fly_client_on_status_callbacks(timepix_fly_client): + """This tests the _run_status_callbacks method of timepix fly client.""" + # Immediate run when current status already in success + timepix_fly_client._status = TimePixFlyStatus.INIT + status = StatusBase() + timepix_fly_client.add_status_callback( + status=status, success=[TimePixFlyStatus.INIT], error=[TimePixFlyStatus.EXCEPT], run=True + ) + assert status.done is True and status.success is True + + # Add callback (do not run immediately) and then trigger via _run_status_callbacks + status2 = StatusBase() + timepix_fly_client.add_status_callback( + status=status2, + success=[TimePixFlyStatus.CONFIG], + error=[TimePixFlyStatus.EXCEPT], + run=False, + ) + # Set status to CONFIG and mark started True to check reset + timepix_fly_client._status = TimePixFlyStatus.CONFIG + timepix_fly_client._started = True + timepix_fly_client._run_status_callbacks() + assert status2.done is True and status2.success is True + assert timepix_fly_client._started is False + + # Error path: add with run True when status is EXCEPT + timepix_fly_client._status = TimePixFlyStatus.EXCEPT + status3 = StatusBase() + with mock.patch.object(timepix_fly_client, "last_error") as mock_last_error: + mock_err = mock.Mock() + mock_err.message = "boom" + mock_last_error.return_value = mock_err + + timepix_fly_client.add_status_callback( + status=status3, + success=[TimePixFlyStatus.INIT], + error=[TimePixFlyStatus.EXCEPT], + run=True, + ) + assert status3.done is True and status3.success is False + + +def test_timepix_fly_client_shutdown(timepix_fly_client): + """This tests the shutdown method of timepix fly client.""" + mock_ws = mock.Mock() + timepix_fly_client.ws_client = mock_ws + + timepix_fly_client.shutdown() + + mock_ws.close.assert_called_once() + assert timepix_fly_client.ws_client is None + assert timepix_fly_client._shutdown_event.is_set() + + +def test_timepix_fly_client_start(timepix_fly_client): + """This tests the start method of timepix fly client.""" + # The client._get uses requests.get with f"http://{self.rest_url}/{get_cmd}". + # We mock requests.get and verify it is called with the expected URL and timeout, + # and that the _started flag is set. + module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client" + with mock.patch(f"{module_path}.requests.get") as mock_get: + mock_resp = mock.Mock() + mock_resp.raise_for_status = mock.Mock() + mock_resp.text = "" + mock_get.return_value = mock_resp + + timepix_fly_client.start() + + expected_url = f"http://{timepix_fly_client.rest_url}/?start=true" + mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout) + assert timepix_fly_client._started is True + + +def test_timepix_fly_client_stop_collect(timepix_fly_client): + """This tests the stop_collect method of timepix fly client.""" + module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client" + with mock.patch(f"{module_path}.requests.get") as mock_get: + mock_resp = mock.Mock() + mock_resp.raise_for_status = mock.Mock() + mock_resp.text = "" + mock_get.return_value = mock_resp + + # ensure started is True and then call stop_collect + timepix_fly_client._started = True + timepix_fly_client.stop_collect() + + expected_url = f"http://{timepix_fly_client.rest_url}/?stop_collect=true" + mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout) + assert timepix_fly_client._started is False + + +def test_timepix_fly_client_state(timepix_fly_client): + """This tests the state method of timepix fly client.""" + module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client" + with mock.patch(f"{module_path}.requests.get") as mock_get: + mock_resp = mock.Mock() + mock_resp.raise_for_status = mock.Mock() + # Return a JSON payload compatible with ProgramState + mock_resp.json.return_value = {"type": "ProgramState", "state": "init"} + mock_get.return_value = mock_resp + + program_state = timepix_fly_client.state() + + expected_url = f"http://{timepix_fly_client.rest_url}/state" + mock_get.assert_called_once_with(expected_url, timeout=timepix_fly_client._timeout) + # ProgramState.model defines 'state' as a string literal; ensure we parsed it + assert hasattr(program_state, "state") + assert program_state.state == "init" + + +def test_timepix_fly_client_set_pixel_map(timepix_fly_client): + """This tests the set_pixel_map/_put path by mocking requests.put and checking payload.""" + module_path = "superxas_bec.devices.timepix.timepix_fly_client.timepix_fly_client" + with mock.patch(f"{module_path}.requests.put") as mock_put: + mock_resp = mock.Mock() + mock_resp.raise_for_status = mock.Mock() + mock_put.return_value = mock_resp + + # Minimal valid PixelMap dict (type is optional; model supplies default) + pixel_map = {"chips": [[{"i": 0, "p": 1, "f": 2}]]} + + timepix_fly_client.set_pixel_map(pixel_map) + + expected_url = f"http://{timepix_fly_client.rest_url}/pixel-map" + # Verify requests.put called with expected url, json and timeout + mock_put.assert_called_once() + _, kwargs = mock_put.call_args + assert kwargs.get("timeout") == timepix_fly_client._timeout + assert kwargs.get("json") is not None + assert kwargs.get("json").get("chips") == pixel_map["chips"]