Allow shared parent mock to be passed to Device.connect

src/ophyd_async/core/_device.py

@@ -10,6 +10,7 @@
     Optional,
     TypeVar,
 )
+from unittest.mock import Mock
 
 from bluesky.protocols import HasName
 from bluesky.run_engine import call_in_bluesky_event_loop
@@ -73,7 +74,7 @@ def set_name(self, name: str):
 
     async def connect(
         self,
-        mock: bool = False,
+        mock: bool | Mock = False,
         timeout: float = DEFAULT_TIMEOUT,
         force_reconnect: bool = False,
     ):
@@ -89,29 +90,38 @@ async def connect(
             Time to wait before failing with a TimeoutError.
         """
 
+        if mock is True:  # create a new Mock if one not provided
+            mock = Mock()
+
         if (
             self._previous_connect_was_mock is not None
-            and self._previous_connect_was_mock != mock
+            and self._previous_connect_was_mock != isinstance(mock, Mock)
         ):
             raise RuntimeError(
-                f"`connect(mock={mock})` called on a `Device` where the previous "
-                f"connect was `mock={self._previous_connect_was_mock}`. Changing mock "
-                "value between connects is not permitted."
+                "connect(mock) was called with mock evaluating "
+                f"to {bool(mock)} when previous connect had mock "
+                f" evaluate to {self._previous_connect_was_mock}. "
+                "Changing mock value between connects is not permitted."
             )
-        self._previous_connect_was_mock = mock
+
+        self._previous_connect_was_mock = isinstance(mock, Mock)
 
         # If previous connect with same args has started and not errored, can use it
         can_use_previous_connect = self._connect_task and not (
             self._connect_task.done() and self._connect_task.exception()
         )
         if force_reconnect or not can_use_previous_connect:
             # Kick off a connection
-            coros = {
-                name: child_device.connect(
-                    mock, timeout=timeout, force_reconnect=force_reconnect
+            coros = {}
+
+            for child_name, child_device in self.children():
+                # if mock is Mock instance, get child Mock if exists
+                # else pass down boolean
+                child_mock = getattr(mock, child_name, True) if mock else mock
+                coros[child_name] = child_device.connect(
+                    child_mock, timeout=timeout, force_reconnect=force_reconnect
                 )
-                for name, child_device in self.children()
-            }
+
             self._connect_task = asyncio.create_task(wait_for_connection(**coros))
 
         assert self._connect_task, "Connect task not created, this shouldn't happen"

src/ophyd_async/core/_signal.py

@@ -4,6 +4,7 @@
 import functools
 from collections.abc import AsyncGenerator, Callable, Mapping
 from typing import Any, Generic, TypeVar, cast
+from unittest.mock import Mock
 
 from bluesky.protocols import (
     Locatable,
@@ -62,7 +63,7 @@ def __init__(
 
     async def connect(
         self,
-        mock=False,
+        mock: bool | Mock = False,
         timeout=DEFAULT_TIMEOUT,
         force_reconnect: bool = False,
         backend: SignalBackend[T] | None = None,
@@ -75,6 +76,9 @@ async def connect(
                 raise ValueError("Backend at connection different from previous one.")
 
             self._backend = backend
+
+        mock = bool(mock)  # Device can pass Mock to child Signal
+
         if (
             self._previous_connect_was_mock is not None
             and self._previous_connect_was_mock != mock

src/ophyd_async/fastcs/panda/_hdf_panda.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from unittest.mock import Mock
 
 from ophyd_async.core import DEFAULT_TIMEOUT, PathProvider, SignalR, StandardDetector
 from ophyd_async.epics.pvi import create_children_from_annotations, fill_pvi_entries
@@ -37,7 +38,7 @@ def __init__(
 
     async def connect(
         self,
-        mock: bool = False,
+        mock: bool | Mock = False,
         timeout: float = DEFAULT_TIMEOUT,
         force_reconnect: bool = False,
     ):

src/ophyd_async/plan_stubs/_ensure_connected.py

@@ -1,11 +1,13 @@
+from unittest.mock import Mock
+
 import bluesky.plan_stubs as bps
 
 from ophyd_async.core import DEFAULT_TIMEOUT, Device, wait_for_connection
 
 
 def ensure_connected(
     *devices: Device,
-    mock: bool = False,
+    mock: bool | Mock = False,
     timeout: float = DEFAULT_TIMEOUT,
     force_reconnect=False,
 ):

tests/core/test_device.py

@@ -127,7 +127,7 @@ async def test_device_log_has_correct_name():
     assert device.log.extra["ophyd_async_device_name"] == "device"
 
 
-async def test_device_lazily_connects(RE):
+async def test_device_connect(RE):
     class MockSignalBackendFailingFirst(MockSignalBackend):
         succeed_on_connect = False
 
@@ -178,8 +178,10 @@ async def test_device_refuses_two_connects_differing_on_mock_attribute(RE):
     with pytest.raises(RuntimeError) as exc:
         await motor.connect(mock=True)
     assert str(exc.value) == (
-        "`connect(mock=True)` called on a `Device` where the previous connect was "
-        "`mock=False`. Changing mock value between connects is not permitted."
+        "connect(mock) was called with mock evaluating "
+        "to True when previous connect had mock "
+        " evaluate to False. "
+        "Changing mock value between connects is not permitted."
     )
 
 
@@ -227,8 +229,10 @@ async def test_device_with_device_collector_refuses_to_connect_if_mock_switch():
     with pytest.raises(RuntimeError) as exc:
         await mock_motor.connect(mock=True, timeout=0.01)
     assert str(exc.value) == (
-        "`connect(mock=True)` called on a `Device` where the previous connect was "
-        "`mock=False`. Changing mock value between connects is not permitted."
+        "connect(mock) was called with mock evaluating "
+        "to True when previous connect had mock "
+        " evaluate to False. "
+        "Changing mock value between connects is not permitted."
     )