Update init_devices context manager to allow for granular selection of device mocks

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)
 

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()