ophyd-async 0.9.0a2__py3-none-any.whl → 0.10.0a2__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,7 +1,10 @@
+"""Utilities for testing devices."""
+
 from . import __pytest_assert_rewrite  # noqa: F401
 from ._assert import (
     ApproxTable,
     MonitorQueue,
+    StatusWatcher,
     approx_value,
     assert_configuration,
     assert_describe_signal,
@@ -14,7 +17,6 @@ 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,
@@ -24,24 +26,36 @@ from ._one_of_everything import (
     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",
@@ -49,4 +63,12 @@ __all__ = [
     "ParentOfEverythingDevice",
     "MonitorQueue",
     "ApproxTable",
+    "StatusWatcher",
+    "int_array_value",
+    "float_array_value",
+    # Derived examples
+    "BeamstopPosition",
+    "Exploder",
+    "MovableBeamstop",
+    "ReadOnlyBeamstop",
 ]

ophyd_async/testing/_assert.py

@@ -1,7 +1,8 @@
 import asyncio
-import time
+from collections.abc import Mapping
 from contextlib import AbstractContextManager
-from typing import Any
+from typing import Any, cast
+from unittest.mock import Mock, call
 
 import pytest
 from bluesky.protocols import Reading
@@ -13,89 +14,89 @@ from ophyd_async.core import (
     SignalDatatypeT,
     SignalR,
     Table,
+    WatchableAsyncStatus,
+    Watcher,
 )
 
+from ._utils import T
+
 
 def approx_value(value: Any):
+    """Allow any value to be compared to another in tests.
+
+    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)
 
 
 async def assert_value(signal: SignalR[SignalDatatypeT], value: Any) -> None:
-    """Assert a signal's value and compare it an expected signal.
-
-    Parameters
-    ----------
-    signal:
-        signal with get_value.
-    value:
-        The expected value from the signal.
-
-    Notes
-    -----
-    Example usage::
-        await assert_value(signal, value)
+    """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 approx_value(value) == actual_value
 
 
 async def assert_reading(
-    readable: AsyncReadable, expected_reading: dict[str, Reading]
+    readable: AsyncReadable,
+    expected_reading: Mapping[str, Mapping[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()
-    approx_expected_reading = {
-        k: dict(v, value=approx_value(expected_reading[k]["value"]))
-        for k, v in expected_reading.items()
+    _assert_readings_approx_equal(expected_reading, actual_reading)
+
+
+def _approx_reading(expected: Mapping[str, Any], actual: Reading) -> Reading:
+    ret = dict(
+        expected,
+        value=approx_value(expected["value"]),
+        timestamp=pytest.approx(expected["timestamp"], rel=0.1)
+        if "timestamp" in expected
+        else actual["timestamp"],
+    )
+    if "alarm_severity" in actual and "alarm_severity" not in expected:
+        ret["alarm_severity"] = actual["alarm_severity"]
+    return cast(Reading, ret)
+
+
+def _assert_readings_approx_equal(
+    expected: Mapping[str, Mapping[str, Any]], actual: Mapping[str, Reading]
+):
+    assert actual == {
+        k: _approx_reading(v, actual[k]) for k, v in expected.items() if k in actual
     }
-    assert actual_reading == approx_expected_reading
 
 
 async def assert_configuration(
     configurable: AsyncConfigurable,
-    configuration: dict[str, Reading],
+    configuration: dict[str, dict[str, Any]],
 ) -> None:
-    """Assert readings from Configurable.
-
-    Parameters
-    ----------
-    configurable:
-        Configurable with Configurable.read function that generate readings.
-
-    configuration:
-        The expected readings from configurable.
-
-    Notes
-    -----
-    Example usage::
-        await assert_configuration(configurable configuration)
+    """Assert that a configurable Device has the given configuration.
 
+    :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()
-    approx_expected_configuration = {
-        k: dict(v, value=approx_value(configuration[k]["value"]))
-        for k, v in configuration.items()
-    }
-    assert actual_configuration == approx_expected_configuration
+    _assert_readings_approx_equal(configuration, actual_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_describe = await signal.describe()
     assert list(actual_describe) == [signal.name]
     (actual_datakey,) = actual_describe.values()
@@ -104,23 +105,18 @@ async def assert_describe_signal(signal: SignalR, /, **metadata):
 
 
 def assert_emitted(docs: dict[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)
+    """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)
     actual_numbers = {name: len(d) for name, d in docs.items()}
@@ -128,6 +124,14 @@ def assert_emitted(docs: dict[str, list[dict]], **numbers: int):
 
 
 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
@@ -144,29 +148,23 @@ class ApproxTable:
 
 
 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()
-        self.signal.subscribe(self.updates.put_nowait)
 
     async def assert_updates(self, expected_value):
         # Get an update, value and reading
-        expected_type = type(expected_value)
-        expected_value = approx_value(expected_value)
-        update = await self.updates.get()
-        value = await self.signal.get_value()
-        reading = await self.signal.read()
-        # Check they match what we expected
-        assert value == expected_value
-        assert type(value) is expected_type
+        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,
-                "timestamp": pytest.approx(time.time(), rel=0.1),
-                "alarm_severity": 0,
             }
         }
-        assert reading == update == expected_reading
+        await assert_reading(self.signal, expected_reading)
+        _assert_readings_approx_equal(expected_reading, update)
 
     def __enter__(self):
         self.signal.subscribe(self.updates.put_nowait)
@@ -174,3 +172,67 @@ class MonitorQueue(AbstractContextManager):
 
     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()