diff --git a/src/ophyd_async/core/_device.py b/src/ophyd_async/core/_device.py index f9937d6055..d659d6625a 100644 --- a/src/ophyd_async/core/_device.py +++ b/src/ophyd_async/core/_device.py @@ -6,7 +6,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, MutableMapping from functools import cached_property from logging import LoggerAdapter, getLogger -from typing import Any, Generic, TypeVar +from typing import Any, Generic, TypeVar, get_args, get_origin from unittest.mock import Mock from bluesky.protocols import HasName @@ -553,11 +553,24 @@ async def _on_exit(self) -> None: await self._process_devices(devices) +def _mocked_device_type(mock_cls: type[DeviceMock]) -> type[Device]: + # Find X in the nearest `DeviceMock[X]` base; unresolved TypeVars or + # forward refs fall back to Device so the mock applies to everything + for cls in mock_cls.__mro__: + for base in cls.__dict__.get("__orig_bases__", ()): + origin = get_origin(base) + if isinstance(origin, type) and issubclass(origin, DeviceMock): + args = get_args(base) + if args and isinstance(args[0], type): + return args[0] + return Device + + def init_devices( set_name: bool = True, child_name_separator: str = "-", connect: bool = True, - mock: bool = False, + mock: bool | type[DeviceMock] | dict[type[Device], bool | type[DeviceMock]] = False, timeout: float = 10.0, ): """Auto initialize top level Device instances: to be used as a context manager. @@ -569,7 +582,15 @@ def init_devices( :param connect: If True, call `device.connect(mock, timeout)` in parallel on all Devices created within the context manager. - :param mock: If True, connect Signals in mock mode. + :param mock: + If True, connect in mock mode using each Device's registered default + mock. If a [](#DeviceMock) subclass, use it for every Device matching + its generic parameter (e.g. `DeviceMock[Motor]` is used for Motors) and + the registered default mock for the rest. If a dict, look up each + Device's class and its bases (most derived first) and use the first + match, where `True` means the registered default and `False` or no + match means connect for real. For example + `{Device: True, Motor: MyMotorMock}` mocks everything, overriding Motors. :param timeout: How long to wait for connect before logging an exception. :raises RuntimeError: If used inside a plan, use [](#ensure_connected) instead. :raises NotConnectedError: If devices could not be connected. @@ -583,6 +604,20 @@ def init_devices( assert t1x.name == "t1x" ``` """ + mocked_type = _mocked_device_type(mock) if isinstance(mock, type) else Device + + def mock_for(device: Device) -> bool | DeviceMock: + if isinstance(mock, dict): + mock_cls = next( + (mock[cls] for cls in type(device).__mro__ if cls in mock), False + ) + elif isinstance(mock, type) and not isinstance(device, mocked_type): + mock_cls = True + else: + mock_cls = mock + if isinstance(mock_cls, type): + return mock_cls() + return bool(mock_cls) async def process_devices(devices: dict[str, Device]): if set_name: @@ -591,7 +626,8 @@ async def process_devices(devices: dict[str, Device]): device.set_name(name, child_name_separator=child_name_separator) if connect: coros = { - name: device.connect(mock, timeout) for name, device in devices.items() + name: device.connect(mock_for(device), timeout) + for name, device in devices.items() } await wait_for_connection(**coros) diff --git a/tests/unit_tests/core/test_auto_init_devices.py b/tests/unit_tests/core/test_auto_init_devices.py index 027d2931b5..a1732afd4e 100644 --- a/tests/unit_tests/core/test_auto_init_devices.py +++ b/tests/unit_tests/core/test_auto_init_devices.py @@ -7,7 +7,9 @@ from ophyd_async.core import ( DEFAULT_TIMEOUT, Device, + DeviceMock, NotConnectedError, + default_mock_class, init_devices, set_mock_value, ) @@ -16,7 +18,10 @@ class FailingDevice(Device): async def connect( - self, mock: bool = False, timeout=DEFAULT_TIMEOUT, force_reconnect=False + self, + mock: bool | DeviceMock = False, + timeout=DEFAULT_TIMEOUT, + force_reconnect=False, ): raise AttributeError() @@ -25,7 +30,10 @@ class WorkingDevice(Device): connected = False async def connect( - self, mock: bool = True, timeout=DEFAULT_TIMEOUT, force_reconnect=False + self, + mock: bool | DeviceMock = True, + timeout=DEFAULT_TIMEOUT, + force_reconnect=False, ): self.connected = True return await super().connect(mock=True) @@ -81,6 +89,61 @@ async def test_init_devices_detects_redeclared_devices(): assert not original_working_device.connected and original_working_device.name == "" +class RecordingMock(DeviceMock): + async def connect(self, device: Device) -> None: + device.connected_mock = self + + +class DefaultMock(RecordingMock): ... + + +class CustomMock(RecordingMock): ... + + +@default_mock_class(DefaultMock) +class DeviceA(Device): ... + + +@default_mock_class(DefaultMock) +class DeviceB(Device): ... + + +class DeviceAOnlyMock(RecordingMock, DeviceMock[DeviceA]): ... + + +class DeviceAOnlySubMock(DeviceAOnlyMock): ... + + +@pytest.mark.parametrize( + "mock, expected_a, expected_b", + [ + (False, None, None), + (True, DefaultMock, DefaultMock), + (CustomMock, CustomMock, CustomMock), + (DeviceAOnlyMock, DeviceAOnlyMock, DefaultMock), + (DeviceAOnlySubMock, DeviceAOnlySubMock, DefaultMock), + ({DeviceA: CustomMock}, CustomMock, None), + ({DeviceA: CustomMock, DeviceB: DefaultMock}, CustomMock, DefaultMock), + ({}, None, None), + ({Device: CustomMock}, CustomMock, CustomMock), + ({Device: True, DeviceA: CustomMock}, CustomMock, DefaultMock), + ({Device: CustomMock, DeviceB: False}, CustomMock, None), + ], +) +async def test_init_devices_mock_options(mock, expected_a, expected_b): + async with init_devices(mock=mock): + device_a = DeviceA() + device_b = DeviceB() + + mock_a = getattr(device_a, "connected_mock", None) + mock_b = getattr(device_b, "connected_mock", None) + assert (type(mock_a) if mock_a else None) is expected_a + assert (type(mock_b) if mock_b else None) is expected_b + # Each device must get its own mock instance, not a shared one + if mock_a and mock_b: + assert mock_a is not mock_b + + def test_connecting_in_plan_raises(RE): def bad_plan(): yield from bps.null()