ophyd-async 0.9.0a1__py3-none-any.whl → 0.10.0a1__py3-none-any.whl

ophyd_async/tango/testing/_one_of_everything.py

@@ -0,0 +1,200 @@
+import textwrap
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+import numpy as np
+
+from ophyd_async.core import (
+    Array1D,
+    DTypeScalar_co,
+    StrictEnum,
+)
+from ophyd_async.testing import float_array_value, int_array_value
+from tango import AttrDataFormat, AttrWriteType, DevState
+from tango.server import Device, attribute, command
+
+T = TypeVar("T")
+
+
+class ExampleStrEnum(StrictEnum):
+    A = "AAA"
+    B = "BBB"
+    C = "CCC"
+
+
+def int_image_value(
+    dtype: type[DTypeScalar_co],
+):
+    # how do we type this?
+    array_1d = int_array_value(dtype)
+    return np.vstack((array_1d, array_1d))
+
+
+def float_image_value(
+    dtype: type[DTypeScalar_co],
+):
+    # how do we type this?
+    array_1d = float_array_value(dtype)
+    return np.vstack((array_1d, array_1d))
+
+
+def _valid_command(dformat: AttrDataFormat, dtype: str):
+    if dtype == "DevUChar":
+        return False
+    if dformat != AttrDataFormat.SCALAR and dtype in ["DevState", "DevEnum"]:
+        return False
+    return True
+
+
+@dataclass
+class AttributeData(Generic[T]):
+    name: str
+    tango_type: str
+    initial_scalar: T
+    initial_spectrum: Array1D
+
+
+_all_attribute_definitions = [
+    AttributeData(
+        "str",
+        "DevString",
+        "test_string",
+        np.array(["one", "two", "three"], dtype=str),
+    ),
+    AttributeData(
+        "bool",
+        "DevBoolean",
+        True,
+        np.array([False, True], dtype=bool),
+    ),
+    AttributeData("strenum", "DevEnum", 1, np.array([0, 1, 2])),
+    AttributeData("int8", "DevShort", 1, int_array_value(np.int8)),
+    AttributeData("uint8", "DevUChar", 1, int_array_value(np.uint8)),
+    AttributeData("int16", "DevShort", 1, int_array_value(np.int16)),
+    AttributeData("uint16", "DevUShort", 1, int_array_value(np.uint16)),
+    AttributeData("int32", "DevLong", 1, int_array_value(np.int32)),
+    AttributeData("uint32", "DevULong", 1, int_array_value(np.uint32)),
+    AttributeData("int64", "DevLong64", 1, int_array_value(np.int64)),
+    AttributeData("uint64", "DevULong64", 1, int_array_value(np.uint64)),
+    AttributeData("float32", "DevFloat", 1.234, float_array_value(np.float32)),
+    AttributeData("float64", "DevDouble", 1.234, float_array_value(np.float64)),
+    AttributeData(
+        "my_state",
+        "DevState",
+        DevState.INIT,
+        np.array([DevState.INIT, DevState.ON, DevState.MOVING], dtype=DevState),
+    ),
+]
+
+
+class OneOfEverythingTangoDevice(Device):
+    attr_values = {}
+    initial_values = {}
+
+    def _add_attr(self, attr: attribute, initial_value):
+        self.attr_values[attr.name] = initial_value
+        self.initial_values[attr.name] = initial_value
+        self.add_attribute(attr)
+        self.set_change_event(attr.name, True, False)
+
+    def add_scalar_attr(self, name: str, dtype: str, initial_value: Any):
+        attr = attribute(
+            name=name,
+            dtype=dtype,
+            dformat=AttrDataFormat.SCALAR,
+            access=AttrWriteType.READ_WRITE,
+            fget=self.read,
+            fset=self.write,
+            enum_labels=[e.value for e in ExampleStrEnum],
+        )
+        self._add_attr(attr, initial_value)
+
+    def add_array_attrs(self, name: str, dtype: str, initial_value: np.ndarray):
+        spectrum_name = f"{name}_spectrum"
+        spectrum_attr = attribute(
+            name=spectrum_name,
+            dtype=dtype,
+            dformat=AttrDataFormat.SPECTRUM,
+            access=AttrWriteType.READ_WRITE,
+            fget=self.read,
+            fset=self.write,
+            max_dim_x=initial_value.shape[-1],
+            enum_labels=[e.value for e in ExampleStrEnum],
+        )
+        image_name = f"{name}_image"
+        image_attr = attribute(
+            name=image_name,
+            dtype=dtype,
+            dformat=AttrDataFormat.IMAGE,
+            access=AttrWriteType.READ_WRITE,
+            fget=self.read,
+            fset=self.write,
+            max_dim_x=initial_value.shape[-1],
+            max_dim_y=2,
+            enum_labels=[e.value for e in ExampleStrEnum],
+        )
+        self._add_attr(spectrum_attr, initial_value)
+        # have image just be 2 of the initial spectrum stacked
+        self._add_attr(image_attr, np.vstack((initial_value, initial_value)))
+
+    def add_scalar_command(self, name: str, dtype: str):
+        if _valid_command(AttrDataFormat.SCALAR, dtype):
+            self.add_command(
+                command(
+                    f=getattr(self, f"{name}_cmd"),
+                    dtype_in=dtype,
+                    dtype_out=dtype,
+                    dformat_in=AttrDataFormat.SCALAR,
+                    dformat_out=AttrDataFormat.SCALAR,
+                ),
+            )
+
+    def add_spectrum_command(self, name: str, dtype: str):
+        if _valid_command(AttrDataFormat.SPECTRUM, dtype):
+            self.add_command(
+                command(
+                    f=getattr(self, f"{name}_spectrum_cmd"),
+                    dtype_in=dtype,
+                    dtype_out=dtype,
+                    dformat_in=AttrDataFormat.SPECTRUM,
+                    dformat_out=AttrDataFormat.SPECTRUM,
+                ),
+            )
+
+    def initialize_dynamic_attributes(self):
+        for attr_data in _all_attribute_definitions:
+            self.add_scalar_attr(
+                attr_data.name, attr_data.tango_type, attr_data.initial_scalar
+            )
+            self.add_array_attrs(
+                attr_data.name, attr_data.tango_type, attr_data.initial_spectrum
+            )
+            self.add_scalar_command(attr_data.name, attr_data.tango_type)
+            self.add_spectrum_command(attr_data.name, attr_data.tango_type)
+
+    @command
+    def reset_values(self):
+        for attr_name in self.attr_values:
+            self.attr_values[attr_name] = self.initial_values[attr_name]
+
+    def read(self, attr):
+        value = self.attr_values[attr.get_name()]
+        attr.set_value(value)
+
+    def write(self, attr):
+        new_value = attr.get_write_value()
+        self.attr_values[attr.get_name()] = new_value
+        self.push_change_event(attr.get_name(), new_value)
+
+    echo_command_code = textwrap.dedent(
+        """\
+            def {}(self, arg):
+                return arg
+            """
+    )
+
+    for attr_data in _all_attribute_definitions:
+        if _valid_command(AttrDataFormat.SCALAR, attr_data.tango_type):
+            exec(echo_command_code.format(f"{attr_data.name}_cmd"))
+        if _valid_command(AttrDataFormat.SPECTRUM, attr_data.tango_type):
+            exec(echo_command_code.format(f"{attr_data.name}_spectrum_cmd"))

ophyd_async/testing/__init__.py

@@ -1,5 +1,13 @@
+"""Utilities for testing devices."""
+
+from . import __pytest_assert_rewrite  # noqa: F401
 from ._assert import (
+    ApproxTable,
+    MonitorQueue,
+    StatusWatcher,
+    approx_value,
     assert_configuration,
+    assert_describe_signal,
     assert_emitted,
     assert_reading,
     assert_value,
@@ -9,25 +17,58 @@ from ._mock_signal_utils import (
     get_mock,
     get_mock_put,
     mock_puts_blocked,
-    reset_mock_put_calls,
     set_mock_put_proceeds,
     set_mock_value,
     set_mock_values,
 )
+from ._one_of_everything import (
+    ExampleEnum,
+    ExampleTable,
+    OneOfEverythingDevice,
+    ParentOfEverythingDevice,
+    float_array_value,
+    int_array_value,
+)
+from ._single_derived import (
+    BeamstopPosition,
+    Exploder,
+    MovableBeamstop,
+    ReadOnlyBeamstop,
+)
 from ._wait_for_pending import wait_for_pending_wakeups
 
+# The order of this list determines the order of the documentation,
+# so does not match the alphabetical order of the imports
 __all__ = [
+    "approx_value",
+    # Assert functions
+    "assert_value",
+    "assert_reading",
     "assert_configuration",
+    "assert_describe_signal",
     "assert_emitted",
-    "assert_reading",
-    "assert_value",
-    "callback_on_mock_put",
+    # Mocking utilities
     "get_mock",
+    "set_mock_value",
+    "set_mock_values",
     "get_mock_put",
+    "callback_on_mock_put",
     "mock_puts_blocked",
-    "reset_mock_put_calls",
     "set_mock_put_proceeds",
-    "set_mock_value",
-    "set_mock_values",
+    # Wait for pending wakeups
     "wait_for_pending_wakeups",
+    "ExampleEnum",
+    "ExampleTable",
+    "OneOfEverythingDevice",
+    "ParentOfEverythingDevice",
+    "MonitorQueue",
+    "ApproxTable",
+    "StatusWatcher",
+    "int_array_value",
+    "float_array_value",
+    # Derived examples
+    "BeamstopPosition",
+    "Exploder",
+    "MovableBeamstop",
+    "ReadOnlyBeamstop",
 ]

ophyd_async/testing/__pytest_assert_rewrite.py

@@ -0,0 +1,4 @@
+import pytest
+
+# So that bare asserts give a nice pytest traceback
+pytest.register_assert_rewrite("ophyd_async.testing._assert")

ophyd_async/testing/_assert.py

@@ -1,128 +1,232 @@
-from collections.abc import Mapping
+import asyncio
+from contextlib import AbstractContextManager
 from typing import Any
+from unittest.mock import Mock, call
 
+import pytest
 from bluesky.protocols import Reading
+from event_model import DataKey
 
-from ophyd_async.core import AsyncConfigurable, AsyncReadable, SignalDatatypeT, SignalR
+from ophyd_async.core import (
+    AsyncConfigurable,
+    AsyncReadable,
+    SignalDatatypeT,
+    SignalR,
+    Table,
+    WatchableAsyncStatus,
+    Watcher,
+)
 
+from ._utils import T
 
-def _generate_assert_error_msg(name: str, expected_result, actual_result) -> str:
-    WARNING = "\033[93m"
-    FAIL = "\033[91m"
-    ENDC = "\033[0m"
-    return (
-        f"Expected {WARNING}{name}{ENDC} to produce"
-        + f"\n{FAIL}{expected_result}{ENDC}"
-        + f"\nbut actually got \n{FAIL}{actual_result}{ENDC}"
-    )
 
+def approx_value(value: Any):
+    """Allow any value to be compared to another in tests.
 
-async def assert_value(signal: SignalR[SignalDatatypeT], value: Any) -> None:
-    """Assert a signal's value and compare it an expected signal.
+    This is needed because numpy arrays give a numpy array back when compared,
+    not a bool. This means that you can't ``assert array1==array2``. Numpy
+    arrays can be wrapped with `pytest.approx`, but this doesn't work for
+    `Table` instances: in this case we use `ApproxTable`.
+    """
+    return ApproxTable(value) if isinstance(value, Table) else pytest.approx(value)
 
-    Parameters
-    ----------
-    signal:
-        signal with get_value.
-    value:
-        The expected value from the signal.
 
-    Notes
-    -----
-    Example usage::
-        await assert_value(signal, value)
+async def assert_value(signal: SignalR[SignalDatatypeT], value: Any) -> None:
+    """Assert that a Signal has the given value.
 
+    :param signal: Signal with get_value.
+    :param value: The expected value from the signal.
     """
     actual_value = await signal.get_value()
-    assert actual_value == value, _generate_assert_error_msg(
-        name=signal.name,
-        expected_result=value,
-        actual_result=actual_value,
-    )
+    assert approx_value(value) == actual_value
 
 
 async def assert_reading(
-    readable: AsyncReadable, expected_reading: Mapping[str, Reading]
+    readable: AsyncReadable,
+    expected_reading: dict[str, dict[str, Any]],
 ) -> None:
-    """Assert readings from readable.
-
-    Parameters
-    ----------
-    readable:
-        Callable with readable.read function that generate readings.
-
-    reading:
-        The expected readings from the readable.
-
-    Notes
-    -----
-    Example usage::
-        await assert_reading(readable, reading)
+    """Assert that a readable Device has the given reading.
 
+    :param readable: Device with an async ``read()`` method to get the reading from.
+    :param expected_reading: The expected reading from the readable.
     """
     actual_reading = await readable.read()
-    assert expected_reading == actual_reading, _generate_assert_error_msg(
-        name=readable.name,
-        expected_result=expected_reading,
-        actual_result=actual_reading,
-    )
+    _assert_readings_approx_equal(expected_reading, actual_reading)
+
+
+def _assert_readings_approx_equal(expected, actual):
+    assert expected.keys() == actual.keys()
+    approx_expected_reading = {
+        k: dict(
+            v,
+            value=approx_value(v["value"]),
+            timestamp=pytest.approx(v["timestamp"], rel=0.1)
+            if "timestamp" in v
+            else actual[k]["timestamp"],
+            alarm_severity=v.get("alarm_severity", actual[k]["alarm_severity"]),
+        )
+        for k, v in expected.items()
+    }
+    assert actual == approx_expected_reading
 
 
 async def assert_configuration(
     configurable: AsyncConfigurable,
-    configuration: Mapping[str, Reading],
+    configuration: dict[str, dict[str, Any]],
 ) -> None:
-    """Assert readings from Configurable.
+    """Assert that a configurable Device has the given configuration.
 
-    Parameters
-    ----------
-    configurable:
-        Configurable with Configurable.read function that generate readings.
+    :param configurable:
+        Device with an async ``read_configuration()`` method to get the
+        configuration from.
+    :param configuration: The expected configuration from the configurable.
+    """
+    actual_configuration = await configurable.read_configuration()
+    _assert_readings_approx_equal(configuration, actual_configuration)
 
-    configuration:
-        The expected readings from configurable.
 
-    Notes
-    -----
-    Example usage::
-        await assert_configuration(configurable configuration)
+async def assert_describe_signal(signal: SignalR, /, **metadata):
+    """Assert the describe of a signal matches the expected metadata.
 
+    :param signal: The signal to describe.
+    :param metadata: The expected metadata.
     """
-    actual_configurable = await configurable.read_configuration()
-    assert configuration == actual_configurable, _generate_assert_error_msg(
-        name=configurable.name,
-        expected_result=configuration,
-        actual_result=actual_configurable,
-    )
-
-
-def assert_emitted(docs: Mapping[str, list[dict]], **numbers: int):
-    """Assert emitted document generated by running a Bluesky plan
-
-    Parameters
-    ----------
-    Doc:
-        A dictionary
-
-    numbers:
-        expected emission in kwarg from
-
-    Notes
-    -----
-    Example usage::
-        docs = defaultdict(list)
-        RE.subscribe(lambda name, doc: docs[name].append(doc))
-        RE(my_plan())
-        assert_emitted(docs, start=1, descriptor=1, event=1, stop=1)
+    actual_describe = await signal.describe()
+    assert list(actual_describe) == [signal.name]
+    (actual_datakey,) = actual_describe.values()
+    expected_datakey = DataKey(source=signal.source, **metadata)
+    assert actual_datakey == expected_datakey
+
+
+def assert_emitted(docs: dict[str, list[dict]], **numbers: int):
+    """Assert emitted document generated by running a Bluesky plan.
+
+    :param docs: A mapping of document type -> list of documents that have been emitted.
+    :param numbers: The number of each document type expected.
+
+    :example:
+    ```python
+    docs = defaultdict(list)
+    RE.subscribe(lambda name, doc: docs[name].append(doc))
+    RE(my_plan())
+    assert_emitted(docs, start=1, descriptor=1, event=1, stop=1)
+    ```
     """
-    assert list(docs) == list(numbers), _generate_assert_error_msg(
-        name="documents",
-        expected_result=list(numbers),
-        actual_result=list(docs),
-    )
+    assert list(docs) == list(numbers)
     actual_numbers = {name: len(d) for name, d in docs.items()}
-    assert actual_numbers == numbers, _generate_assert_error_msg(
-        name="emitted",
-        expected_result=numbers,
-        actual_result=actual_numbers,
-    )
+    assert actual_numbers == numbers
+
+
+class ApproxTable:
+    """For approximating two tables are equivalent.
+
+    :param expected: The expected table.
+    :param rel: The relative tolerance.
+    :param abs: The absolute tolerance.
+    :param nan_ok: Whether NaNs are allowed.
+    """
+
+    def __init__(self, expected: Table, rel=None, abs=None, nan_ok: bool = False):
+        self.expected = expected
+        self.rel = rel
+        self.abs = abs
+        self.nan_ok = nan_ok
+
+    def __eq__(self, value):
+        approx_fields = {
+            k: pytest.approx(v, self.rel, self.abs, self.nan_ok)
+            for k, v in self.expected
+        }
+        expected = type(self.expected).model_construct(**approx_fields)  # type: ignore
+        return expected == value
+
+
+class MonitorQueue(AbstractContextManager):
+    """Monitors a `Signal` and stores its updates."""
+
+    def __init__(self, signal: SignalR):
+        self.signal = signal
+        self.updates: asyncio.Queue[dict[str, Reading]] = asyncio.Queue()
+
+    async def assert_updates(self, expected_value):
+        # Get an update, value and reading
+        update = await asyncio.wait_for(self.updates.get(), timeout=5)
+        await assert_value(self.signal, expected_value)
+        expected_reading = {
+            self.signal.name: {
+                "value": expected_value,
+            }
+        }
+        await assert_reading(self.signal, expected_reading)
+        _assert_readings_approx_equal(expected_reading, update)
+
+    def __enter__(self):
+        self.signal.subscribe(self.updates.put_nowait)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.signal.clear_sub(self.updates.put_nowait)
+
+
+class StatusWatcher(Watcher[T]):
+    """Watches an `AsyncStatus`, storing the calls within."""
+
+    def __init__(self, status: WatchableAsyncStatus) -> None:
+        self._event = asyncio.Event()
+        self.mock = Mock()
+        """Mock that stores watcher updates from the status."""
+        status.watch(self)
+
+    def __call__(
+        self,
+        current: T | None = None,
+        initial: T | None = None,
+        target: T | None = None,
+        name: str | None = None,
+        unit: str | None = None,
+        precision: int | None = None,
+        fraction: float | None = None,
+        time_elapsed: float | None = None,
+        time_remaining: float | None = None,
+    ) -> Any:
+        self.mock(
+            current=current,
+            initial=initial,
+            target=target,
+            name=name,
+            unit=unit,
+            precision=precision,
+            fraction=fraction,
+            time_elapsed=time_elapsed,
+            time_remaining=time_remaining,
+        )
+        self._event.set()
+
+    async def wait_for_call(
+        self,
+        current: T | None = None,
+        initial: T | None = None,
+        target: T | None = None,
+        name: str | None = None,
+        unit: str | None = None,
+        precision: int | None = None,
+        fraction: float | None = None,
+        # Any so we can use pytest.approx
+        time_elapsed: float | Any = None,
+        time_remaining: float | Any = None,
+    ):
+        await asyncio.wait_for(self._event.wait(), timeout=1)
+        assert self.mock.call_count == 1
+        assert self.mock.call_args == call(
+            current=current,
+            initial=initial,
+            target=target,
+            name=name,
+            unit=unit,
+            precision=precision,
+            fraction=fraction,
+            time_elapsed=time_elapsed,
+            time_remaining=time_remaining,
+        )
+        self.mock.reset_mock()
+        self._event.clear()