mirror of
https://github.com/bec-project/ophyd_devices.git
synced 2026-07-25 00:22:58 +02:00
1261 lines
40 KiB
Python
1261 lines
40 KiB
Python
import threading
|
|
import time
|
|
from functools import partial
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
import ophyd
|
|
import pytest
|
|
from bec_lib import messages
|
|
from bec_lib.bec_errors import ExceptionWithErrorInfo
|
|
from ophyd import Component as Cpt
|
|
from ophyd import Device, EpicsSignalRO, Signal
|
|
from ophyd.status import DeviceStatus as OphydDeviceStatus
|
|
from ophyd.status import MoveStatus as OphydMoveStatus
|
|
from ophyd.status import Status as OphydStatus
|
|
from ophyd.status import StatusBase as OphydStatusBase
|
|
from ophyd.status import WaitTimeoutError
|
|
from typeguard import TypeCheckError
|
|
|
|
from ophyd_devices.devices.psi_motor import EpicsMotor
|
|
from ophyd_devices.tests.utils import MockPV, patched_device
|
|
from ophyd_devices.utils.bec_signals import (
|
|
AsyncMultiSignal,
|
|
AsyncSignal,
|
|
BECMessageSignal,
|
|
DynamicSignal,
|
|
FileEventSignal,
|
|
PreviewSignal,
|
|
ProgressSignal,
|
|
)
|
|
from ophyd_devices.utils.psi_device_base_utils import (
|
|
AndStatus,
|
|
CompareStatus,
|
|
DeviceStatus,
|
|
ExceptionStatus,
|
|
FileHandler,
|
|
MoveStatus,
|
|
Status,
|
|
StatusBase,
|
|
StatusTimeoutErrorWithErrorInfo,
|
|
SubscriptionStatus,
|
|
TaskHandler,
|
|
TaskKilledError,
|
|
TaskState,
|
|
TaskStatus,
|
|
TransitionStatus,
|
|
)
|
|
|
|
# pylint: disable=protected-access
|
|
# pylint: disable=redefined-outer-name
|
|
|
|
##########################################
|
|
######### Test Task Handler ############
|
|
##########################################
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_epics_signal_ro():
|
|
name = "epics_signal_ro"
|
|
read_pv = "TEST:EPICS_SIGNAL_RO"
|
|
with mock.patch.object(ophyd, "cl") as mock_cl:
|
|
mock_cl.get_pv = MockPV
|
|
mock_cl.thread_class = threading.Thread
|
|
dev = EpicsSignalRO(name=name, read_pv=read_pv)
|
|
yield dev
|
|
|
|
|
|
@pytest.fixture
|
|
def file_handler():
|
|
"""Fixture for FileHandler"""
|
|
yield FileHandler()
|
|
|
|
|
|
@pytest.fixture
|
|
def device():
|
|
"""Fixture for Device"""
|
|
yield Device(name="device")
|
|
|
|
|
|
@pytest.fixture
|
|
def task_handler(device):
|
|
"""Fixture for TaskHandler"""
|
|
yield TaskHandler(parent=device)
|
|
|
|
|
|
def test_utils_file_handler_has_full_path(file_handler):
|
|
"""Ensure that file_handler has a get_full_path method"""
|
|
assert hasattr(file_handler, "get_full_path")
|
|
|
|
|
|
def test_utils_task_status(device):
|
|
"""Test TaskStatus creation"""
|
|
status = TaskStatus(device)
|
|
assert status.obj.name == "device"
|
|
assert status.state == "not_started"
|
|
assert status.task_id == status._task_id
|
|
status.state = "running"
|
|
assert status.state == TaskState.RUNNING
|
|
status.state = TaskState.COMPLETED
|
|
assert status.state == "completed"
|
|
|
|
|
|
def test_utils_task_handler_submit_task_with_args(task_handler):
|
|
"""Ensure that task_handler has a submit_task method"""
|
|
|
|
def my_task(input_arg: bool, input_kwarg: bool = False):
|
|
if input_kwarg is True:
|
|
raise ValueError("input_kwarg is True")
|
|
if input_arg is True:
|
|
return True
|
|
return False
|
|
|
|
# This should fail
|
|
with pytest.raises(TypeError):
|
|
status = task_handler.submit_task(my_task)
|
|
status.wait()
|
|
# This should pass
|
|
|
|
task_stopped = threading.Event()
|
|
|
|
def finished_cb():
|
|
task_stopped.set()
|
|
|
|
status = task_handler.submit_task(
|
|
my_task, task_args=(True,), task_kwargs={"input_kwarg": False}
|
|
)
|
|
status.add_callback(finished_cb)
|
|
task_stopped.wait()
|
|
assert status.done is True
|
|
assert status.state == TaskState.COMPLETED
|
|
# This should fail
|
|
task_stopped = threading.Event()
|
|
status = task_handler.submit_task(my_task, task_args=(True,), task_kwargs={"input_kwarg": True})
|
|
with pytest.raises(ValueError):
|
|
status.wait()
|
|
assert status.state == TaskState.ERROR
|
|
assert status.done is True
|
|
assert status.exception().__class__ == ValueError
|
|
|
|
|
|
@pytest.mark.timeout(100)
|
|
def test_utils_task_handler_task_killed(task_handler):
|
|
"""Ensure that task_handler has a submit_task method"""
|
|
# No tasks should be running
|
|
assert len(task_handler._tasks) == 0
|
|
event = threading.Event()
|
|
task_stopped = threading.Event()
|
|
task_started = threading.Event()
|
|
|
|
def finished_cb():
|
|
task_stopped.set()
|
|
|
|
def my_wait_task():
|
|
task_started.set()
|
|
for _ in range(100):
|
|
event.wait(timeout=0.1)
|
|
|
|
# Create task
|
|
status = task_handler.submit_task(my_wait_task, run=False)
|
|
status.add_callback(finished_cb)
|
|
assert status.state == TaskState.NOT_STARTED
|
|
# Start task
|
|
task_handler.start_task(status)
|
|
task_started.wait()
|
|
assert status.state == TaskState.RUNNING
|
|
# Stop task
|
|
task_handler.kill_task(status)
|
|
task_stopped.wait()
|
|
assert status.state == TaskState.KILLED
|
|
assert status.exception().__class__ == TaskKilledError
|
|
|
|
|
|
@pytest.mark.timeout(100)
|
|
def test_utils_task_handler_task_successful(task_handler):
|
|
"""Ensure that the task handler runs a successful task"""
|
|
assert len(task_handler._tasks) == 0
|
|
event = threading.Event()
|
|
task_stopped = threading.Event()
|
|
task_started = threading.Event()
|
|
|
|
def finished_cb():
|
|
task_stopped.set()
|
|
|
|
def my_wait_task():
|
|
task_started.set()
|
|
for _ in range(100):
|
|
ret = event.wait(timeout=0.1)
|
|
if ret is True:
|
|
break
|
|
|
|
status = task_handler.submit_task(my_wait_task, run=False)
|
|
status.add_callback(finished_cb)
|
|
task_handler.start_task(status)
|
|
task_started.wait()
|
|
assert status.state == TaskState.RUNNING
|
|
event.set()
|
|
task_stopped.wait()
|
|
assert status.state == TaskState.COMPLETED
|
|
|
|
|
|
def test_utils_task_handler_shutdown(task_handler):
|
|
"""Test to shutdown the handler"""
|
|
|
|
task_completed_cb1 = threading.Event()
|
|
task_completed_cb2 = threading.Event()
|
|
|
|
def finished_cb1():
|
|
task_completed_cb1.set()
|
|
|
|
def finished_cb2():
|
|
task_completed_cb2.set()
|
|
|
|
def cb1():
|
|
for _ in range(1000):
|
|
time.sleep(0.2)
|
|
|
|
def cb2():
|
|
for _ in range(1000):
|
|
time.sleep(0.2)
|
|
|
|
status1 = task_handler.submit_task(cb1)
|
|
status1.add_callback(finished_cb1)
|
|
status2 = task_handler.submit_task(cb2)
|
|
status2.add_callback(finished_cb2)
|
|
assert len(task_handler._tasks) == 2
|
|
assert status1.state == TaskState.RUNNING
|
|
assert status2.state == TaskState.RUNNING
|
|
task_handler.shutdown()
|
|
task_completed_cb1.wait()
|
|
task_completed_cb2.wait()
|
|
assert len(task_handler._tasks) == 0
|
|
assert status1.state == TaskState.KILLED
|
|
assert status2.state == TaskState.KILLED
|
|
assert status1.exception().__class__ == TaskKilledError
|
|
|
|
|
|
##########################################
|
|
######### Test PSI cusomt signals ######
|
|
##########################################
|
|
|
|
|
|
def test_utils_bec_message_signal():
|
|
"""Test BECMessageSignal"""
|
|
dev = Device(name="device")
|
|
signal = BECMessageSignal(
|
|
name="bec_message_signal",
|
|
bec_message_type=messages.GUIInstructionMessage,
|
|
value=None,
|
|
parent=dev,
|
|
)
|
|
assert signal.parent == dev
|
|
assert signal._bec_message_type == messages.GUIInstructionMessage
|
|
assert signal._readback is None
|
|
assert signal.name == "bec_message_signal"
|
|
assert signal.describe() == {
|
|
"bec_message_signal": {
|
|
"source": "BECMessageSignal:bec_message_signal",
|
|
"dtype": "GUIInstructionMessage",
|
|
"shape": [],
|
|
"signal_info": {
|
|
"data_type": "raw",
|
|
"saved": True,
|
|
"ndim": 0,
|
|
"scope": "scan",
|
|
"role": "main",
|
|
"enabled": True,
|
|
"rpc_access": False,
|
|
"signals": [("bec_message_signal", 5)],
|
|
"signal_metadata": {},
|
|
"acquisition_group": None,
|
|
},
|
|
}
|
|
}
|
|
# Put works with Message
|
|
msg = messages.GUIInstructionMessage(action="image", parameter={"gui_id": "test"})
|
|
signal.put(msg)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# set works with dict, should call put
|
|
msg_dict = {"action": "image", "parameter": {"gui_id": "test"}}
|
|
status = signal.set(msg_dict)
|
|
assert status.done is True
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Put fails with wrong type
|
|
with pytest.raises(ValueError):
|
|
signal.put("wrong_type")
|
|
# Put fails with wrong dict
|
|
with pytest.raises(ValueError):
|
|
signal.put({"wrong_key": "wrong_value"})
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_msg, output_msg",
|
|
[
|
|
(
|
|
messages.DeviceMessage(
|
|
signals={"sig1": {"value": 1}, "sig2": {"value": 2}}, metadata={"info": "test"}
|
|
),
|
|
messages.DeviceMessage(
|
|
signals={"device_data_sig1": {"value": 1}, "device_data_sig2": {"value": 2}},
|
|
metadata={"info": "test"},
|
|
),
|
|
),
|
|
(
|
|
messages.DeviceMessage(
|
|
signals={"device_data_sig1": {"value": 1}, "device_data_sig2": {"value": 2}},
|
|
metadata={"info": "test"},
|
|
),
|
|
messages.DeviceMessage(
|
|
signals={"device_data_sig1": {"value": 1}, "device_data_sig2": {"value": 2}},
|
|
metadata={"info": "test"},
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_utils_signal_normalization(input_msg, output_msg):
|
|
"""Test signal normalization utility in BECMessageSignal"""
|
|
|
|
class DeviceWithSignal(Device):
|
|
data = Cpt(AsyncMultiSignal, name="data", signals=["sig1", "sig2"], ndim=0, max_size=1000)
|
|
|
|
dev = DeviceWithSignal(name="device")
|
|
dev.data._normalize_signals(input_msg)
|
|
assert input_msg == output_msg
|
|
|
|
|
|
def test_utils_dynamic_signal():
|
|
"""Test DynamicSignal"""
|
|
dev = Device(name="device")
|
|
signal = DynamicSignal(
|
|
name="dynamic_signal", signals=["sig1", "sig2"], value=None, max_size=10, parent=dev
|
|
)
|
|
assert signal.parent == dev
|
|
assert signal._bec_message_type == messages.DeviceMessage
|
|
assert signal._readback is None
|
|
assert signal.name == "dynamic_signal"
|
|
assert signal.signals == [("sig1", 1), ("sig2", 1)]
|
|
assert signal.describe() == {
|
|
"dynamic_signal": {
|
|
"source": "BECMessageSignal:dynamic_signal",
|
|
"dtype": "DeviceMessage",
|
|
"shape": [],
|
|
"signal_info": {
|
|
"data_type": "raw",
|
|
"saved": True,
|
|
"ndim": 1,
|
|
"scope": "scan",
|
|
"role": "main",
|
|
"enabled": True,
|
|
"rpc_access": False,
|
|
"signals": [("sig1", 1), ("sig2", 1)],
|
|
"signal_metadata": {"max_size": 10},
|
|
"acquisition_group": None,
|
|
},
|
|
}
|
|
}
|
|
|
|
# Put works with Message
|
|
msg_dict = {"dynamic_signal_sig1": {"value": 1}, "dynamic_signal_sig2": {"value": 2}}
|
|
with pytest.raises(ValueError):
|
|
# Missing metadata
|
|
signal.put(messages.DeviceMessage(signals=msg_dict))
|
|
metadata = {"async_update": {"type": "add", "max_shape": [None, 1000]}}
|
|
msg = messages.DeviceMessage(signals=msg_dict, metadata=metadata)
|
|
signal.put(msg)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Set works with dict
|
|
status = signal.set(msg_dict, metadata=metadata)
|
|
assert status.done is True
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Put fails with wrong type
|
|
with pytest.raises(TypeCheckError):
|
|
signal.put("wrong_type")
|
|
# Put fails with wrong dict
|
|
with pytest.raises(TypeCheckError):
|
|
signal.put({"wrong_key": "wrong_value"})
|
|
|
|
# Set with acquisition group
|
|
signal.put(msg, acquisition_group="fly-scan")
|
|
reading = signal.read()
|
|
msg.metadata["acquisition_group"] = "fly-scan"
|
|
assert reading[signal.name]["value"] == msg
|
|
|
|
|
|
def test_utils_dynamic_signal_with_defaults():
|
|
"""
|
|
Test DynamicSignal with async_update and acquisition group defaults. If only
|
|
one sub-signal is provided for the dynamic signal, the name of the sub-signal
|
|
will be used in the signals dict and a warning will be issued that the sub-signal
|
|
name is being ignored.
|
|
"""
|
|
dev = Device(name="device")
|
|
create_signal = partial(
|
|
DynamicSignal,
|
|
name="dynamic_signal",
|
|
parent=dev,
|
|
ndim=1,
|
|
value=None,
|
|
max_size=10,
|
|
async_update={"type": "add", "max_shape": [None, 1000]},
|
|
acquisition_group="fly-scanning",
|
|
)
|
|
signal = create_signal(signals=["sig1", "sig2"])
|
|
val = np.random.random(1000)
|
|
msg_dict = {"dynamic_signal_sig1": {"value": val}}
|
|
signal.put(msg_dict)
|
|
reading = signal.read()
|
|
reading_value = reading[signal.name]["value"].model_dump(exclude={"timestamp"})
|
|
assert reading_value["signals"] == msg_dict
|
|
assert reading_value["metadata"]["async_update"] == {"type": "add", "max_shape": [None, 1000]}
|
|
assert reading_value["metadata"]["acquisition_group"] == "fly-scanning"
|
|
|
|
signal.put(msg_dict, acquisition_group="different-group")
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].metadata["acquisition_group"] == "different-group"
|
|
|
|
# Test init variations for single signal
|
|
for signal in [["sig1"], "sig1", None]:
|
|
signal = create_signal(signals=signal)
|
|
assert signal.signals == [(signal.name, ophyd.Kind.hinted.value)]
|
|
|
|
|
|
def test_utils_async_multi_signal():
|
|
"""Test AsyncMultiSignal, which is a DynamicSignal with strict signal validation."""
|
|
device = Device(name="device")
|
|
signal = AsyncMultiSignal(
|
|
name="async_multi_signal",
|
|
ndim=1,
|
|
max_size=1000,
|
|
signals=["sig1", "sig2"],
|
|
async_update={"type": "add", "max_shape": [None, 1000]},
|
|
parent=device,
|
|
)
|
|
val = np.random.random(1000)
|
|
msg_dict = {"async_multi_signal_sig1": {"value": val}}
|
|
with pytest.raises(ValueError):
|
|
# Missing signal
|
|
signal.put(msg_dict)
|
|
msg_dict = {
|
|
"async_multi_signal_sig1": {"value": val},
|
|
"async_multi_signal_sig2": {"value": val},
|
|
}
|
|
signal.put(msg_dict)
|
|
reading = signal.read()
|
|
reading_value = reading[signal.name]["value"].model_dump(exclude={"timestamp"})
|
|
assert reading_value["signals"] == msg_dict
|
|
assert reading_value["metadata"]["async_update"] == {"type": "add", "max_shape": [None, 1000]}
|
|
|
|
|
|
def test_utils_async_signal():
|
|
device = Device(name="device")
|
|
signal = AsyncSignal(
|
|
name="async_signal",
|
|
ndim=1,
|
|
max_size=1000,
|
|
async_update={"type": "add", "max_shape": [None, 200]},
|
|
parent=device,
|
|
)
|
|
val = np.random.random(1000)
|
|
signal.put(
|
|
val,
|
|
async_update={"type": "add_slice", "max_shape": [None, 1000], "index": 1},
|
|
acquisition_group="scan",
|
|
)
|
|
reading = signal.read()
|
|
reading_value = reading[signal.name]["value"].model_dump(exclude={"timestamp"})
|
|
assert np.array_equal(reading_value["signals"][signal.name]["value"], val)
|
|
assert reading_value["metadata"]["async_update"] == {
|
|
"type": "add_slice",
|
|
"max_shape": [None, 1000],
|
|
"index": 1,
|
|
}
|
|
assert reading_value["metadata"]["acquisition_group"] == "scan"
|
|
|
|
|
|
def test_utils_file_event_signal():
|
|
"""Test FileEventSignal"""
|
|
dev = Device(name="device")
|
|
signal = FileEventSignal(name="file_event_signal", value=None, parent=dev)
|
|
assert signal.parent == dev
|
|
assert signal._bec_message_type == messages.FileMessage
|
|
assert signal._readback is None
|
|
assert signal.name == "file_event_signal"
|
|
assert signal.describe() == {
|
|
"file_event_signal": {
|
|
"source": "BECMessageSignal:file_event_signal",
|
|
"dtype": "FileMessage",
|
|
"shape": [],
|
|
"signal_info": {
|
|
"data_type": "raw",
|
|
"saved": False,
|
|
"ndim": 0,
|
|
"scope": "scan",
|
|
"role": "file_event",
|
|
"enabled": True,
|
|
"rpc_access": False,
|
|
"signals": [("file_event_signal", 5)],
|
|
"signal_metadata": {},
|
|
"acquisition_group": None,
|
|
},
|
|
}
|
|
}
|
|
|
|
# Test put works with FileMessage
|
|
msg_dict = {"file_path": "/path/to/another/file.txt", "done": False, "successful": True}
|
|
msg = messages.FileMessage(**msg_dict)
|
|
signal.put(msg)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Test put works with dict
|
|
signal.put(msg_dict)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Test set with kwargs, should call put
|
|
status = signal.set(file_path="/path/to/another/file.txt", done=False, successful=True)
|
|
assert status.done is True
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Test put fails with wrong type
|
|
with pytest.raises(ValueError):
|
|
signal.put(1)
|
|
# Test put fails with wrong dict
|
|
with pytest.raises(ValueError):
|
|
signal.put({"wrong_key": "wrong_value"})
|
|
|
|
|
|
def test_utils_preview_1d_signal():
|
|
"""Test Preview1DSignal"""
|
|
dev = Device(name="device")
|
|
signal = PreviewSignal(name="preview_1d_signal", ndim=1, value=None, parent=dev)
|
|
assert signal.ndim == 1
|
|
assert signal.parent == dev
|
|
assert signal._bec_message_type == messages.DevicePreviewMessage
|
|
assert signal._readback is None
|
|
assert signal.name == "preview_1d_signal"
|
|
assert signal.describe() == {
|
|
"preview_1d_signal": {
|
|
"source": "BECMessageSignal:preview_1d_signal",
|
|
"dtype": "DevicePreviewMessage",
|
|
"shape": [],
|
|
"signal_info": {
|
|
"data_type": "raw",
|
|
"saved": False,
|
|
"ndim": 1,
|
|
"scope": "scan",
|
|
"role": "preview",
|
|
"enabled": True,
|
|
"rpc_access": False,
|
|
"signals": [("preview_1d_signal", 5)],
|
|
"signal_metadata": {"num_rotation_90": 0, "transpose": False},
|
|
"acquisition_group": None,
|
|
},
|
|
}
|
|
}
|
|
# Put works with Message
|
|
msg_dict = {"device": dev.name, "data": np.array([1, 2, 3]), "signal": "preview_1d_signal"}
|
|
msg = messages.DevicePreviewMessage(**msg_dict)
|
|
signal.put(msg)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put works with dict
|
|
signal.put(msg_dict)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put works with value
|
|
status = signal.set(msg_dict["data"])
|
|
assert status.done is True
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put works with value
|
|
signal.put(msg_dict["data"])
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put fails with wrong type
|
|
with pytest.raises(ValueError):
|
|
signal.put(1)
|
|
# Put fails with wrong dict
|
|
with pytest.raises(ValueError):
|
|
signal.put({"wrong_key": "wrong_value"})
|
|
|
|
|
|
def test_utils_preview_2d_signal():
|
|
"""Test Preview2DSignal"""
|
|
dev = Device(name="device")
|
|
signal = PreviewSignal(name="preview_2d_signal", ndim=2, value=None, parent=dev)
|
|
assert signal.ndim == 2
|
|
assert signal.parent == dev
|
|
assert signal._bec_message_type == messages.DevicePreviewMessage
|
|
assert signal._readback is None
|
|
assert signal.name == "preview_2d_signal"
|
|
assert signal.describe() == {
|
|
"preview_2d_signal": {
|
|
"source": "BECMessageSignal:preview_2d_signal",
|
|
"dtype": "DevicePreviewMessage",
|
|
"shape": [],
|
|
"signal_info": {
|
|
"data_type": "raw",
|
|
"saved": False,
|
|
"ndim": 2,
|
|
"scope": "scan",
|
|
"role": "preview",
|
|
"enabled": True,
|
|
"rpc_access": False,
|
|
"signals": [("preview_2d_signal", 5)],
|
|
"signal_metadata": {"num_rotation_90": 0, "transpose": False},
|
|
"acquisition_group": None,
|
|
},
|
|
}
|
|
}
|
|
# Put works with Message
|
|
msg_dict = {
|
|
"device": dev.name,
|
|
"data": np.array([[1, 2, 3], [4, 5, 6]]),
|
|
"signal": "preview_2d_signal",
|
|
}
|
|
msg = messages.DevicePreviewMessage(**msg_dict)
|
|
signal.put(msg)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put works with dict
|
|
signal.put(msg_dict)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put works with value
|
|
status = signal.set(msg_dict["data"])
|
|
assert status.done is True
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put works with value
|
|
signal.put(msg_dict["data"])
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"].model_dump(exclude="timestamp") == msg.model_dump(
|
|
exclude="timestamp"
|
|
)
|
|
# Put fails with wrong type
|
|
with pytest.raises(ValueError):
|
|
signal.put(1)
|
|
# Put fails with wrong dict
|
|
with pytest.raises(ValueError):
|
|
signal.put({"wrong_key": "wrong_value"})
|
|
|
|
|
|
def test_utils_preview_2d_signal_transpose_rgb_swaps_only_spatial_axes():
|
|
"""Transpose must preserve channel axis for RGB-like preview data."""
|
|
dev = Device(name="device")
|
|
signal = PreviewSignal(
|
|
name="preview_rgb_signal", ndim=2, value=None, parent=dev, transpose=True
|
|
)
|
|
data = np.arange(2 * 3 * 4).reshape(2, 3, 4)
|
|
|
|
signal.put(data)
|
|
reading = signal.read()[signal.name]["value"]
|
|
|
|
np.testing.assert_array_equal(reading.data, np.swapaxes(data, 0, 1))
|
|
assert reading.data.shape == (3, 2, 4)
|
|
|
|
|
|
def test_utils_preview_2d_signal_rotation_and_transpose_rgb():
|
|
"""Rotation and transpose should both operate on spatial axes only."""
|
|
dev = Device(name="device")
|
|
signal = PreviewSignal(
|
|
name="preview_rgb_rot_signal",
|
|
ndim=2,
|
|
value=None,
|
|
parent=dev,
|
|
num_rotation_90=1,
|
|
transpose=True,
|
|
)
|
|
data = np.arange(2 * 3 * 4).reshape(2, 3, 4)
|
|
|
|
signal.put(data)
|
|
reading = signal.read()[signal.name]["value"]
|
|
|
|
expected = np.swapaxes(np.rot90(data, k=1, axes=(0, 1)), 0, 1)
|
|
np.testing.assert_array_equal(reading.data, expected)
|
|
|
|
|
|
def test_utils_progress_signal():
|
|
"""Test ProgressSignal"""
|
|
dev = Device(name="device")
|
|
signal = ProgressSignal(name="progress_signal", value=None, parent=dev)
|
|
assert signal.parent == dev
|
|
assert signal._bec_message_type == messages.ProgressMessage
|
|
assert signal._readback is None
|
|
assert signal.name == "progress_signal"
|
|
assert signal.describe() == {
|
|
"progress_signal": {
|
|
"source": "BECMessageSignal:progress_signal",
|
|
"dtype": "ProgressMessage",
|
|
"shape": [],
|
|
"signal_info": {
|
|
"data_type": "raw",
|
|
"saved": False,
|
|
"ndim": 0,
|
|
"scope": "scan",
|
|
"role": "progress",
|
|
"enabled": True,
|
|
"rpc_access": False,
|
|
"signals": [("progress_signal", 5)],
|
|
"signal_metadata": {},
|
|
"acquisition_group": None,
|
|
},
|
|
}
|
|
}
|
|
# Put works with Message
|
|
msg = messages.ProgressMessage(value=1, max_value=10, done=False)
|
|
signal.put(msg)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Put works with dict
|
|
msg_dict = {"value": 1, "max_value": 10, "done": False}
|
|
signal.put(msg_dict)
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Works with kwargs
|
|
status = signal.set(value=1, max_value=10, done=False)
|
|
assert status.done is True
|
|
reading = signal.read()
|
|
assert reading[signal.name]["value"] == msg
|
|
# Put fails with wrong type
|
|
with pytest.raises(ValueError):
|
|
signal.put(1)
|
|
# Put fails with wrong dict
|
|
with pytest.raises(ValueError):
|
|
signal.put({"wrong_key": "wrong_value"})
|
|
|
|
|
|
def test_utils_compare_status_number():
|
|
"""Test CompareStatus with different operations."""
|
|
sig = Signal(name="test_signal", value=0)
|
|
status = CompareStatus(signal=sig, value=5, operation_success="==")
|
|
assert status.done is False
|
|
sig.put(1)
|
|
assert status.done is False
|
|
sig.put(5)
|
|
status.wait(timeout=5)
|
|
assert status.done is True
|
|
|
|
sig.put(5)
|
|
# Test with different operations
|
|
status = CompareStatus(signal=sig, value=5, operation_success="!=")
|
|
assert status.done is False
|
|
sig.put(5)
|
|
assert status.done is False
|
|
sig.put(6)
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
sig.put(0)
|
|
status = CompareStatus(signal=sig, value=5, operation_success=">")
|
|
assert status.done is False
|
|
sig.put(5)
|
|
assert status.done is False
|
|
sig.put(10)
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
# Should raise
|
|
sig.put(0)
|
|
status = CompareStatus(signal=sig, value=5, operation_success="==", failure_value=[10])
|
|
with pytest.raises(ValueError):
|
|
sig.put(10)
|
|
status.wait()
|
|
assert status.done is True
|
|
assert status.success is False
|
|
assert isinstance(status.exception(), ValueError)
|
|
|
|
# failure_operation
|
|
sig.put(0)
|
|
status = CompareStatus(
|
|
signal=sig, value=5, operation_success="==", failure_value=10, operation_failure=">"
|
|
)
|
|
sig.put(10)
|
|
assert status.done is False
|
|
assert status.success is False
|
|
sig.put(11)
|
|
with pytest.raises(ValueError):
|
|
status.wait()
|
|
assert status.done is True
|
|
assert status.success is False
|
|
|
|
# raise if array is returned
|
|
sig.put(0)
|
|
status = CompareStatus(signal=sig, value=5, operation_success="==")
|
|
with pytest.raises(ValueError):
|
|
sig.put([1, 2, 3])
|
|
status.wait(timeout=2)
|
|
assert status.done is True
|
|
assert status.success is False
|
|
|
|
|
|
def test_compare_status_string():
|
|
"""Test CompareStatus with string values"""
|
|
sig = Signal(name="test_signal", value="test")
|
|
status = CompareStatus(signal=sig, value="test", operation_success="==")
|
|
assert status.done is False
|
|
sig.put("test1")
|
|
assert status.done is False
|
|
sig.put("test")
|
|
assert status.done is True
|
|
|
|
sig.put("test")
|
|
# Test with different operations
|
|
status = CompareStatus(signal=sig, value="test", operation_success="!=")
|
|
assert status.done is False
|
|
sig.put("test")
|
|
assert status.done is False
|
|
sig.put("test1")
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
|
|
def test_exception_status_andstatus_does_not_block_success():
|
|
"""ExceptionStatus should fail composites early but not block success while pending."""
|
|
sig_primary_a = Signal(name="primary_a", value=0)
|
|
sig_primary_b = Signal(name="primary_b", value=0)
|
|
sig_watch = Signal(name="watch", value=0)
|
|
|
|
primary_a = CompareStatus(signal=sig_primary_a, value=1, operation_success="==")
|
|
primary_b = CompareStatus(signal=sig_primary_b, value=2, operation_success="==")
|
|
watch = ExceptionStatus(signal=sig_watch, value=0, operation="!=")
|
|
|
|
combined = primary_a & primary_b & watch
|
|
sig_primary_a.put(1)
|
|
assert not combined.done
|
|
sig_primary_b.put(2)
|
|
combined.wait(timeout=1)
|
|
assert combined.done is True
|
|
assert combined.success is True
|
|
assert watch.done is False
|
|
|
|
|
|
def test_exception_status_andstatus_fails_early():
|
|
"""ExceptionStatus should abort a composite status when the watched value is reached."""
|
|
sig_primary = Signal(name="primary", value=0)
|
|
sig_watch = Signal(name="watch", value=0)
|
|
|
|
primary = CompareStatus(signal=sig_primary, value=1, operation_success="==")
|
|
watch = ExceptionStatus(signal=sig_watch, value=0, operation="!=")
|
|
combined = primary & watch
|
|
|
|
sig_watch.put(1)
|
|
with pytest.raises(ValueError):
|
|
combined.wait(timeout=1)
|
|
assert combined.done is True
|
|
assert combined.success is False
|
|
|
|
|
|
def test_exception_status_andstatus_fails_early_with_custom_exception():
|
|
"""ExceptionStatus should abort a composite status with the specified exception when the watched value is reached."""
|
|
sig_primary = Signal(name="primary", value=0)
|
|
sig_watch = Signal(name="watch", value=0)
|
|
|
|
primary = CompareStatus(signal=sig_primary, value=1, operation_success="==")
|
|
watch = ExceptionStatus(
|
|
signal=sig_watch,
|
|
value=0,
|
|
operation="!=",
|
|
exception=RuntimeError("Watch signal reached failure value"),
|
|
)
|
|
combined = primary & watch
|
|
|
|
sig_watch.put(1)
|
|
with pytest.raises(RuntimeError, match="Watch signal reached failure value"):
|
|
combined.wait(timeout=1)
|
|
assert combined.done is True
|
|
assert combined.success is False
|
|
|
|
|
|
def test_exception_status_with_exception():
|
|
"""Test that ExceptionStatus raises the specified exception when the condition is met."""
|
|
sig = Signal(name="test_signal", value=0)
|
|
sig.put(1)
|
|
status = ExceptionStatus(
|
|
signal=sig, value=1, operation="==", exception=RuntimeError("Test signal reached 1")
|
|
)
|
|
assert status.done is True
|
|
assert status.success is False
|
|
with pytest.raises(RuntimeError, match="Test signal reached 1"):
|
|
status.wait(timeout=1)
|
|
|
|
|
|
def test_transition_status():
|
|
"""Test TransitionStatus"""
|
|
sig = Signal(name="test_signal", value=0)
|
|
|
|
# Test strict=True, without intermediate transitions
|
|
sig.put(0)
|
|
status = TransitionStatus(signal=sig, transitions=[1, 2, 3], strict=True)
|
|
|
|
assert status.done is False
|
|
sig.put(1)
|
|
assert status.done is False
|
|
sig.put(2)
|
|
assert status.done is False
|
|
sig.put(3)
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
# Test strict=True, failure_states
|
|
sig.put(1)
|
|
status = TransitionStatus(signal=sig, transitions=[1, 2, 3], strict=True, failure_states=[4])
|
|
assert status.done is False
|
|
sig.put(4)
|
|
with pytest.raises(ValueError):
|
|
status.wait()
|
|
|
|
assert status.done is True
|
|
assert status.success is False
|
|
assert isinstance(status.exception(), ValueError)
|
|
|
|
# Test strict=False, with intermediate transitions
|
|
sig.put(0)
|
|
status = TransitionStatus(signal=sig, transitions=[1, 2, 3], strict=False)
|
|
|
|
assert status.done is False
|
|
sig.put(1) # entering first transition
|
|
sig.put(3)
|
|
sig.put(2) # transision
|
|
assert status.done is False
|
|
sig.put(4)
|
|
sig.put(2)
|
|
sig.put(3) # last transition
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
|
|
def test_transition_status_strings():
|
|
"""Test TransitionStatus with string values"""
|
|
sig = Signal(name="test_signal", value="a")
|
|
|
|
# Test strict=True, without intermediate transitions
|
|
sig.put("a")
|
|
status = TransitionStatus(signal=sig, transitions=["b", "c", "d"], strict=True)
|
|
|
|
assert status.done is False
|
|
sig.put("b")
|
|
assert status.done is False
|
|
sig.put("c")
|
|
assert status.done is False
|
|
sig.put("d")
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
# Test strict=True with additional intermediate transition
|
|
|
|
sig.put("a")
|
|
status = TransitionStatus(signal=sig, transitions=["b", "c", "d"], strict=True)
|
|
|
|
assert status.done is False
|
|
sig.put("b") # first transition
|
|
sig.put("e")
|
|
sig.put("b")
|
|
sig.put("c") # transision
|
|
assert status.done is False
|
|
sig.put("f")
|
|
sig.put("b")
|
|
sig.put("c")
|
|
sig.put("d") # transision
|
|
assert status.done is True
|
|
assert status.success is True
|
|
assert status.exception() is None
|
|
|
|
# Test strict=False, with intermediate transitions
|
|
sig.put("a")
|
|
status = TransitionStatus(signal=sig, transitions=["b", "c", "d"], strict=False)
|
|
|
|
assert status.done is False
|
|
sig.put("b") # entering first transition
|
|
sig.put("d")
|
|
sig.put("c") # transision
|
|
assert status.done is False
|
|
sig.put("e")
|
|
sig.put("c")
|
|
sig.put("d") # last transition
|
|
assert status.done is True
|
|
assert status.success is True
|
|
|
|
|
|
def test_compare_status_with_mock_pv(mock_epics_signal_ro):
|
|
"""Test CompareStatus with EpicsSignalRO, this tests callbacks on EpicsSignals"""
|
|
|
|
signal = mock_epics_signal_ro
|
|
status = CompareStatus(signal=signal, value=5, operation_success="==")
|
|
assert status.done is False
|
|
signal._read_pv.mock_data = 1
|
|
assert status.done is False
|
|
signal._read_pv.mock_data = 5
|
|
status.wait(timeout=1)
|
|
assert status.done is True
|
|
assert status.success is True
|
|
|
|
|
|
def test_compare_status_raises_on_failed_comparison(mock_epics_signal_ro):
|
|
"""Test CompareStatus raises on failed comparison with EpicsSignalRO"""
|
|
|
|
signal = mock_epics_signal_ro
|
|
status = CompareStatus(
|
|
signal=signal, value=5, operation_success="==", failure_value=[np.array([10])]
|
|
)
|
|
assert status.done is False
|
|
signal._read_pv.mock_data = 1
|
|
with pytest.raises(Exception):
|
|
status.wait(timeout=5)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"transitions, expected_done, expected_success",
|
|
[
|
|
([1, 2, 3], True, True), # Transitions completed successfully
|
|
([1, 3, 2], False, False), # Transitions completed with an error
|
|
([5, 4, 2, 1, 2, 3], True, True), # Transitions completed successfully
|
|
],
|
|
)
|
|
def test_transition_status_with_mock_pv(
|
|
mock_epics_signal_ro, transitions, expected_done, expected_success
|
|
):
|
|
"""Test TransitionStatus with EpicsSignalRO, this tests callbacks on EpicsSignals"""
|
|
# Starts immediately with 1
|
|
signal = mock_epics_signal_ro
|
|
signal._read_pv.mock_data = 1
|
|
status = TransitionStatus(signal=signal, transitions=[1, 2, 3], strict=False)
|
|
assert status.done is False
|
|
# Does not have to wait
|
|
signal._read_pv.mock_data = 3
|
|
signal._read_pv.mock_data = 2
|
|
signal._read_pv.mock_data = 3
|
|
status.wait(timeout=1)
|
|
assert status.done is True
|
|
assert status.success is True
|
|
# Test with various transitions
|
|
status = TransitionStatus(signal=signal, transitions=[1, 2, 3], strict=True)
|
|
for transition in transitions:
|
|
signal._read_pv.mock_data = transition
|
|
if expected_done:
|
|
status.wait(timeout=1)
|
|
assert status.done is True
|
|
assert status.success is expected_success
|
|
else:
|
|
with pytest.raises(WaitTimeoutError):
|
|
status.wait(timeout=1)
|
|
assert status.done is False
|
|
assert status.success is False
|
|
|
|
|
|
def test_patched_status_objects():
|
|
"""Test the patched Status objects in ophyd_devices that improve error handling."""
|
|
|
|
# StatusBase & AndStatus
|
|
st = StatusBase()
|
|
st2 = StatusBase()
|
|
and_st = st & st2
|
|
assert st in and_st
|
|
assert isinstance(and_st, AndStatus)
|
|
st.set_exception(ValueError("test error"))
|
|
with pytest.raises(ValueError):
|
|
and_st.wait(timeout=10)
|
|
|
|
# DeviceStatus & Status
|
|
dev = Device(name="device")
|
|
dev_status = DeviceStatus(device=dev)
|
|
|
|
st = Status()
|
|
and_st = st & dev_status
|
|
assert dev_status.device == dev
|
|
dev_status.set_exception(RuntimeError("device error"))
|
|
with pytest.raises(RuntimeError):
|
|
and_st.wait(timeout=10)
|
|
|
|
# Combine DeviceStatus with StatusBase and form AndStatus
|
|
st = StatusBase(obj=dev)
|
|
assert st.obj == dev
|
|
dev_st = DeviceStatus(device=dev)
|
|
combined_st = st & dev_st
|
|
st.set_finished()
|
|
dev_st.set_exception(RuntimeError("combined error"))
|
|
with pytest.raises(RuntimeError):
|
|
combined_st.wait(timeout=10)
|
|
|
|
# SubscriptionStatus
|
|
sig = Signal(name="test_signal", value=0)
|
|
|
|
def _cb(*args, **kwargs):
|
|
pass
|
|
|
|
sub_st = SubscriptionStatus(sig, callback=_cb)
|
|
sub_st.set_exception(ValueError("subscription error"))
|
|
with pytest.raises(ValueError):
|
|
sub_st.wait(timeout=10)
|
|
assert sub_st.done is True
|
|
assert sub_st.success is False
|
|
|
|
# MoveStatus, here the default for call_stop_on_failure is True
|
|
class Positioner(Device):
|
|
SUB_READBACK = "readback"
|
|
setpoint = Signal(name="setpoint", value=0)
|
|
readback = Signal(name="readback", value=0)
|
|
|
|
@property
|
|
def position(self):
|
|
return self.readback.get()
|
|
|
|
def stop(self):
|
|
pass
|
|
|
|
pos = Positioner(name="positioner")
|
|
move_st = MoveStatus(pos, target=10)
|
|
with mock.patch.object(pos, "stop") as mock_stop:
|
|
move_st.set_exception(RuntimeError("move error"))
|
|
mock_stop.assert_called_once()
|
|
with pytest.raises(RuntimeError):
|
|
move_st.wait(timeout=10)
|
|
assert move_st.done is True
|
|
assert move_st.success is False
|
|
|
|
|
|
class Positioner(Device):
|
|
SUB_READBACK = "readback"
|
|
setpoint = Signal(name="setpoint", value=0)
|
|
readback = Signal(name="readback", value=0)
|
|
|
|
@property
|
|
def position(self):
|
|
return self.readback.get()
|
|
|
|
def stop(self):
|
|
pass
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("factory", "expected_compact"),
|
|
[
|
|
pytest.param(
|
|
lambda: StatusBase(timeout=0.01, description="base timed out"),
|
|
"base timed out",
|
|
id="statusbase",
|
|
),
|
|
pytest.param(
|
|
lambda: Status(timeout=0.01, description="status timed out"),
|
|
"status timed out",
|
|
id="status",
|
|
),
|
|
pytest.param(
|
|
lambda: DeviceStatus(
|
|
Device(name="device"), timeout=0.01, description="device timed out"
|
|
),
|
|
"device timed out",
|
|
id="devicestatus",
|
|
),
|
|
pytest.param(
|
|
lambda: MoveStatus(
|
|
Positioner(name="positioner"), target=10, timeout=0.01, description="move timed out"
|
|
),
|
|
"move timed out",
|
|
id="movestatus",
|
|
),
|
|
pytest.param(
|
|
lambda: TaskStatus(Device(name="device"), timeout=0.01, description="task timed out"),
|
|
"task timed out",
|
|
id="taskstatus",
|
|
),
|
|
pytest.param(
|
|
lambda: SubscriptionStatus(
|
|
Signal(name="sig", value=0),
|
|
callback=lambda *args, **kwargs: False,
|
|
timeout=0.01,
|
|
run=False,
|
|
description="subscription timed out",
|
|
),
|
|
"subscription timed out",
|
|
id="subscriptionstatus",
|
|
),
|
|
],
|
|
)
|
|
def test_all_patched_status_types_raise_structured_timeout(factory, expected_compact):
|
|
"""All patched timeout-capable status types should raise structured timeout errors."""
|
|
status = factory()
|
|
|
|
with pytest.raises(StatusTimeoutErrorWithErrorInfo) as exc_info:
|
|
status.wait(timeout=1)
|
|
|
|
assert status.exception() is exc_info.value
|
|
assert isinstance(exc_info.value, ExceptionWithErrorInfo)
|
|
assert "Status initialization traceback" in exc_info.value.error_info.error_message
|
|
assert exc_info.value.error_info.exception_type == "StatusTimeoutError"
|
|
assert expected_compact in exc_info.value.error_info.compact_error_message
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("left_factory", "right_factory"),
|
|
[
|
|
pytest.param(
|
|
lambda: StatusBase(timeout=0.5, description="custom timeout"),
|
|
lambda: StatusBase(),
|
|
id="custom-custom",
|
|
),
|
|
pytest.param(
|
|
lambda: StatusBase(), lambda: OphydStatusBase(timeout=0.5), id="custom-ophydstatusbase"
|
|
),
|
|
pytest.param(
|
|
lambda: StatusBase(), lambda: OphydStatus(timeout=0.5), id="custom-ophydstatus"
|
|
),
|
|
pytest.param(
|
|
lambda: StatusBase(timeout=0.5, description="custom timeout"),
|
|
lambda: OphydDeviceStatus(Device(name="device")),
|
|
id="custom-ophyddevicestatus",
|
|
),
|
|
pytest.param(
|
|
lambda: StatusBase(),
|
|
lambda: OphydMoveStatus(Positioner(name="positioner"), target=10, timeout=0.5),
|
|
id="custom-ophydmovestatus",
|
|
),
|
|
],
|
|
)
|
|
def test_and_status_works_with_plain_ophyd_status_objects(left_factory, right_factory):
|
|
"""Composite statuses should work with both patched and plain ophyd status objects."""
|
|
left = left_factory()
|
|
right = right_factory()
|
|
combined = left & right
|
|
|
|
with pytest.raises(Exception) as exc_info:
|
|
combined.wait(timeout=1)
|
|
|
|
assert combined.done is True
|
|
assert combined.success is False
|
|
assert isinstance(exc_info.value, ExceptionWithErrorInfo)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_device_with_initial_value():
|
|
with patched_device(EpicsMotor, _mock_pv_initial_value=2, name="motor") as mtr:
|
|
yield mtr
|
|
|
|
|
|
def test_mock_device_initial_value(mock_device_with_initial_value: EpicsMotor):
|
|
mtr = mock_device_with_initial_value
|
|
assert mtr.velocity.get() == 2
|