Allow per-type mock overrides anywhere in a device tree

CLAUDE.md

@@ -39,6 +39,7 @@ tox -p                  # all envs in parallel (CI equivalent)
 - **No private-attribute access in tests** (`det._trigger_logic`, …) — call public methods and assert on output (e.g. `trigger()` then `describe()`), unless no public equivalent exists.
 - **Happy path in one public-interface test; unhappy paths small.** For a multi-step lifecycle (e.g. `prepare` → `kickoff` → `complete`), write a *single* "happy path" test that drives the whole sequence through the public interface and asserts the observable end state — don't split one test per step. Then add small "unhappy path" tests (ordering errors, limit violations, injected failures); these stay majority-public-interface but may use mocks to shorten sequences or inject errors. The smell this avoids is a step-scoped test that has to call the *other* steps to set itself up (e.g. a `kickoff` test that also calls `prepare`+`complete`) — fold that into the happy path instead.
 - **Parametrize normal + edge cases together** in one `@pytest.mark.parametrize`, not two functions.
+- **Keep CI test stages to about a minute.** Each "Run tests" stage for unit tests and for system tests should take about a minute in CI. Take care to trim tests down in size, scale and sleeping to stay in that range.
 - `set_mock_value(signal, value)` injects state; `init_devices(mock=True)` (async CM) builds devices; `assert_has_calls(device, [...])` from `ophyd_async.testing` checks PV writes in order.
 
 ## Docs & docstring conventions

docs/explanations/devices-signals-backends.md

@@ -26,12 +26,13 @@ The `Device` class is the base of all ophyd-async objects that are published to
 - a [](#Device.set_name) method to set its name and also set the names of its children using the parent name as a prefix, called at init and also when a new child is attached to an already named Device
 - a [](#Device.connect) method that connects it and its children
 
-All the above methods are concrete, but `connect()` calls out to a [](#DeviceConnector) to actually do the connection, only handling caching itself. This enables plug-in behavior on connect (like the introspection of child Attributes in Tango or PVI, or the special case for `Signal` we will see later).
+All the above methods are concrete, but `connect()` calls out to a [](#DeviceConnector) to actually do the connection, only handling caching itself. This enables plug-in behavior on connect (like the introspection of child Attributes in Tango or PVI, or the special case for `Signal` we will see later). A connect is cached: if the last connect succeeded, calling `connect()` again with the same mock choice (the `DeviceMock` class, `mocked` and `mock_types`) returns without connecting again, in mock mode as well as real. Concurrent calls are not merged: each one that is not cached connects. A connect that is not cached, because the mock choice changed, the last connect failed or has not finished, or `force_reconnect=True` was passed, connects everything below the Device again, so the tree under a re-connected Device is always rebuilt as a whole.
 
 A `DeviceConnector` provides the ability to:
 - [](#DeviceConnector.create_children_from_annotations) that is called during `__init__` to turn annotations into concrete child Devices
-- [](#DeviceConnector.connect_mock) that is called if `connect(mock=True)` is called, and should connect the child Devices in mock mode for testing without a control system
-- [](#DeviceConnector.connect_real) that is called if `connect(mock=False)` is called, and should connect the child Devices to the control system in parallel
+- [](#DeviceConnector.connect) that is called on every connect that is not cached, and should connect the Device to the control system or in mock mode as its `DeviceMock` says, then connect its children with the same timeout and `force_reconnect`: in parallel, or serially when the `DeviceMock` leaves nothing real below the Device. Each child is given a `DeviceMock` that says whether it is mocked, made from the `mock_types` of its parent's, which is how a mapping passed to `init_devices` reaches the whole Device tree
+
+A connect that is not cached, in either mode, may drop subscriptions made before it, so subscribe again afterwards. In mock mode it also resets mock values and the `Mock` call record, and drops staging, so `get_value(cached=True)` raises until the Signal is staged again. Calling `child.connect(mock=True)` directly on a child with a different mock choice from its last connect rebuilds the child's mock as a root, detached from the parent's mock tree and without the parent's `mock_types`, so re-connect from the parent instead.
 
 The base `DeviceConnector` provides suitable methods for use with non-introspected Devices, but there are various control system specific connectors that handle filling annotations in [declarative Devices](./declarative-vs-procedural.md).
 

docs/tutorials/writing-tests-for-devices.md

@@ -25,7 +25,7 @@ We will be writing a test using the pytest framework which encourages fixtures t
 :pyobject: mock_motor
 ```
 
-This fixture opts out of the [automatic mock behaviour](../explanations/when-to-extend-movable.md) by connecting with a plain [](#LazyMock), giving the tests control over when the readback updates mid-move. [](#set_mock_units) and [](#set_mock_precision) inject units and precision metadata directly on the readback signal, without needing dedicated child signals on the device.
+This fixture opts out of the [automatic mock behaviour](../explanations/when-to-extend-movable.md) by connecting with a plain [](#DeviceMock), giving the tests control over when the readback updates mid-move. [](#set_mock_units) and [](#set_mock_precision) inject units and precision metadata directly on the readback signal, without needing dedicated child signals on the device.
 
 If we had any cleanup to do, we would do that after the yield statement.
 
@@ -53,9 +53,47 @@ used when connected in mock mode:
 ```
 
 Now whenever a `Motor` is connected using [](#init_devices)`(mock=True)`, it will
-automatically use `InstantMotorMock` without any fixture setup. You can still override
-the automatic mock for specific tests by passing an explicit [](#DeviceMock) instance
-or a plain [](#LazyMock) directly to `connect()`, as the `mock_motor` fixture above does.
+automatically use `InstantMotorMock` without any fixture setup.
+
+A default mock class should make the Device's verbs work in the most minimal way
+possible, e.g. a move completes instantly. Mocks that model real behavior, like velocity
+or timing, should be opt-in rather than registered as the default.
+
+You can still override the automatic mock for specific tests by passing a
+[](#DeviceMock) instance to `connect()`, as the `mock_motor` fixture above does.
+
+:::{note}
+A second `connect(mock=True)` with the same mock choice is cached, so it keeps the mock values and the call record. `connect(mock=True, force_reconnect=True)`, or a different mock choice, makes a new mock for the Device and all of its children, resetting mock values and the call record and dropping subscriptions and staging made before it. Subscribe or stage again afterwards.
+:::
+
+(override-mock-by-type)=
+### Mocking by Device type
+
+To mock only some Devices, or to override their mocks, including Devices deep in a composite Device, pass a mapping of Device type to `True`, `False` or a `DeviceMock` class to [](#init_devices) or [](#ensure_connected):
+
+```python
+# Real, except Motors and everything under them, which are mocked
+with init_devices(mock={Motor: True}):
+    stage = Stage("PREFIX:")
+
+# Mocked, except Motors, which are real all the way down, signals included
+with init_devices(mock={Device: True, Motor: False}):
+    stage = Stage("PREFIX:")
+
+# Real, except Motors, which use VeloAndAcclRespectingMotorMock
+with init_devices(mock={Motor: VeloAndAcclRespectingMotorMock}):
+    stage = Stage("PREFIX:")
+
+# The whole detector is mocked, plugins included
+with init_devices(mock={AreaDetector: True}):
+    det = AreaDetector(...)
+```
+
+Each Device looks itself up by the nearest class in its MRO, and a `DeviceMock` class means mocked with that class. A Device with no entry follows its parent, so mocking a Device mocks its whole subtree. An entry under `Device` is the default for the first level only, the top level Devices, or the direct children of a `DeviceMock` passed to `connect()`. Below that a Device with no entry follows its parent, so `{Device: True}` mocks the whole tree, and `{Device: XMock}` gives `XMock` to the top level Devices and each child its own default mock class. `Signal` and `Command` types cannot be keys, passing one raises `TypeError`: they always follow their parent Device.
+
+Every connected Device holds the [](#DeviceMock) it was connected with, real ones included, and the mocks of mocked children are nested under the mock of their parent, real or not. So in the examples above [](#get_mock) on `stage` shows the calls made on the mocked motors, as `get_mock(stage).mtr.user_setpoint.put`. Real Signals record nothing, so a fully real tree has no calls to see.
+
+`Device.connect()` takes a `bool` or a `DeviceMock` instance instead. The instance is mocked unless made with `mocked=False`, and its children follow it unless its `mock_types` say otherwise, with a `Device` entry applying to its direct children only. Only `init_devices` and `ensure_connected` check a mapping, so pass one through them.
 
 ## Simulating a signal that stops updating
 
@@ -96,7 +134,9 @@ timing — the [](#Motor) module also ships `VeloAndAcclRespectingMotorMock`, wh
 drives the readback towards the setpoint over time while respecting the motor's
 velocity and acceleration.
 
-To use it, pass an instance directly to `connect()` instead of relying on the default:
+To use it for every Device of a type, rather than one Device, pass `init_devices(mock={DeviceType: MockType})`, as explained in [](#override-mock-by-type).
+
+To use it for a single Device, pass an instance directly to `connect()` instead of relying on the default:
 
 ```{literalinclude} ../../tests/unit_tests/epics/test_motor.py
 :language: python
@@ -146,7 +186,7 @@ Now let's call some verbs and check that they do the right thing. We want to che
 :pyobject: test_motor_stopped
 ```
 
-This time we use [](#get_mock_execute) to get an [](#unittest.mock.AsyncMock) that will be called every time `stop_.trigger()` is called. We check it hasn't been called, then call our method, then check it has been called with no arguments. We also show that we can call [](#get_mock) on the parent to see all of the mock calls that have been made on all its children, useful to check ordering.
+This time we use [](#get_mock_execute) to get an [](#unittest.mock.AsyncMock) that will be called every time `stop_.trigger()` is called. We check it hasn't been called, then call our method, then check it has been called with no arguments. We also show that we can call [](#get_mock) on the parent to see all of the mock calls that have been made on all its children, useful to check ordering. [](#get_mock) works on any connected Device, real or mocked.
 
 For [](#Signal)s, the equivalent is [](#get_mock_put), which returns an `AsyncMock` that records every `Signal.set()` / `put()` call.
 

src/ophyd_async/core/__init__.py

@@ -36,12 +36,8 @@
     DeviceConnector,
     DeviceMap,
     DeviceMock,
-    DeviceProcessor,
     DeviceVector,
-    LazyMock,
     default_mock_class,
-    init_devices,
-    print_connection_status,
 )
 from ._device_filler import DeviceAnnotation, DeviceFiller
 from ._enums import (
@@ -52,6 +48,12 @@
     YesNo,
 )
 from ._flyable import FlyableLogic, FlyMotorInfo, StandardFlyable
+from ._init_devices import (
+    DeviceProcessor,
+    connect_devices,
+    init_devices,
+    print_connection_status,
+)
 from ._log import config_ophyd_async_logging, logger
 from ._mock_signal_backend import MockSignalBackend
 from ._mock_signal_utils import (
@@ -172,6 +174,7 @@ def __getattr__(name):
 
     renames = {
         "NotConnected": NotConnectedError,
+        "LazyMock": DeviceMock,
     }
     rename = renames.get(name)
     if rename is not None:
@@ -194,6 +197,7 @@ def __getattr__(name):
     "DeviceMap",
     "DeviceVector",
     "DeviceProcessor",
+    "connect_devices",
     "init_devices",
     "print_connection_status",
     # Movable
@@ -240,7 +244,6 @@ def __getattr__(name):
     "soft_signal_rw",
     # Mock signal
     "DeviceMock",
-    "LazyMock",
     "MockSignalBackend",
     "default_mock_class",
     # Mocking utilities

src/ophyd_async/core/_command.py

@@ -8,7 +8,7 @@
 from typing import Generic, cast
 from unittest.mock import AsyncMock
 
-from ._device import Device, DeviceConnector, LazyMock
+from ._device import Device, DeviceConnector, DeviceMock
 from ._soft_signal_backend import SoftConverter, make_converter
 from ._status import AsyncStatus
 from ._utils import (
@@ -61,15 +61,20 @@ class CommandConnector(DeviceConnector, Generic[P, T]):
     def __init__(self, backend: CommandBackend[P, T]):
         self.backend = self._init_backend = backend
 
-    async def connect_mock(self, device: Device, mock: LazyMock):
-        """Connect the backend in mock mode."""
-        self.backend = MockCommandBackend(self._init_backend, mock)
-
-    async def connect_real(self, device: Device, timeout: float, force_reconnect: bool):
-        """Connect the backend to real hardware."""
-        self.backend = self._init_backend
-        device.log.debug(f"Connecting to {self.backend.source(device.name)}")
-        await self.backend.connect(timeout)
+    async def connect(
+        self,
+        device: Device,
+        mock: DeviceMock,
+        timeout: float,
+        force_reconnect: bool,
+    ):
+        """Connect the backend in mock mode or to real hardware."""
+        if mock.mocked:
+            self.backend = MockCommandBackend(self._init_backend, mock)
+        else:
+            self.backend = self._init_backend
+            device.log.debug(f"Connecting to {self.backend.source(device.name)}")
+            await self.backend.connect(timeout)
 
 
 class Command(Device, Generic[P, T]):
@@ -265,7 +270,7 @@ async def execute(self, *args: P.args, **kwargs: P.kwargs) -> T:
 class MockCommandBackend(CommandBackend[P, T]):
     """A backend for a Command that uses a mock for testing."""
 
-    def __init__(self, initial_backend: CommandBackend[P, T], mock: LazyMock):
+    def __init__(self, initial_backend: CommandBackend[P, T], mock: DeviceMock):
         self._initial_backend = initial_backend
         self._mock = mock
         self._mock_execute_callback: MockExecuteCallback[P, T] | None = None