ophyd_devices/tests/test_utils.py

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.interfaces.base_classes.psi_device_base import PSIDeviceBase
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


class TimeoutDevice(PSIDeviceBase, Device):
    sig = Cpt(Signal, value=0)

    def __init__(self, *, timeout=3, **kwargs):
        super().__init__(timeout=timeout, **kwargs)


class TimeoutPositioner(PSIDeviceBase, Positioner):
    def __init__(self, *, timeout=3, **kwargs):
        super().__init__(timeout=timeout, **kwargs)


class PlainTimeoutDevice(Device):
    def __init__(self, *, timeout=3, **kwargs):
        self._timeout = timeout
        super().__init__(**kwargs)


def test_patched_status_types_use_default_timeout_from_object():
    """Patched status objects should use the object's default timeout."""
    dev = TimeoutDevice(name="device", timeout=3)
    pos = TimeoutPositioner(name="positioner", timeout=4)

    statuses = [
        StatusBase(obj=dev),
        Status(obj=dev),
        DeviceStatus(dev),
        MoveStatus(pos, target=10),
        TaskStatus(obj=dev),
        SubscriptionStatus(dev.sig, callback=lambda *args, **kwargs: False, run=False),
        CompareStatus(dev.sig, value=1, run=False),
        ExceptionStatus(dev.sig, value=1, run=False),
        TransitionStatus(dev.sig, transitions=[1], run=False),
    ]

    assert [status.timeout for status in statuses] == [3, 3, 3, 4, 3, 3, 3, None, 3]

    # Finish the statuses so their timeout threads do not fire mid-suite and
    # call device.stop() on the test devices.
    for status in statuses:
        if not status.done:
            status.set_finished()
    dev.destroy()
    pos.destroy()


def test_patched_status_explicit_timeout_overrides_object_default():
    """Explicit status timeouts should take precedence over object defaults."""
    dev = TimeoutDevice(name="device", timeout=3)
    pos = TimeoutPositioner(name="positioner", timeout=4)

    statuses = [
        StatusBase(obj=dev, timeout=7),
        Status(obj=dev, timeout=7),
        DeviceStatus(dev, timeout=7),
        MoveStatus(pos, target=10, timeout=7),
        CompareStatus(dev.sig, value=1, timeout=7, run=False),
    ]

    assert [status.timeout for status in statuses] == [7, 7, 7, 7, 7]

    for status in statuses:
        if not status.done:
            status.set_finished()
    dev.destroy()
    pos.destroy()


def test_patched_status_ignores_signal_internal_timeout():
    """Signal connection timeouts should not become status completion timeouts."""
    sig = Signal(name="signal", value=0)
    sig._timeout = 2

    assert CompareStatus(sig, value=1, run=False).timeout is None


def test_patched_status_ignores_plain_device_internal_timeout():
    """Only PSIDeviceBase instances should provide default status timeouts."""
    dev = PlainTimeoutDevice(name="device", timeout=3)

    assert StatusBase(obj=dev).timeout is None
    assert DeviceStatus(dev).timeout is None


@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