store Mocks for Devices and MockSignalBackends for Signals in diction…

…aries

src/ophyd_async/core/_device.py

@@ -13,6 +13,8 @@
 from ._protocol import Connectable
 from ._utils import DEFAULT_TIMEOUT, NotConnected, wait_for_connection
 
+_device_mocks: dict[Device, Mock] = {}
+
 
 class DeviceConnector:
     async def connect(
@@ -56,7 +58,6 @@ def __init__(
         self, name: str = "", connector: DeviceConnector | None = None
     ) -> None:
         self._connector = connector or DeviceConnector()
-        self.mock = None
         self.set_name(name)
 
     @property
@@ -107,16 +108,18 @@ async def connect(
         timeout:
             Time to wait before failing with a TimeoutError.
         """
+        uses_mock = bool(mock)
         can_use_previous_connect = (
-            bool(mock) is self._connect_mock_arg
+            uses_mock is self._connect_mock_arg
             and self._connect_task
             and not (self._connect_task.done() and self._connect_task.exception())
         )
         if mock is True:
             mock = Mock()  # create a new Mock if one not provided
-        self.mock = mock
         if force_reconnect or not can_use_previous_connect:
-            self._connect_mock_arg = bool(mock)
+            self._connect_mock_arg = uses_mock
+            if self._connect_mock_arg:
+                _device_mocks[self] = mock
             coro = self._connector.connect(
                 device=self, mock=mock, timeout=timeout, force_reconnect=force_reconnect
             )
@@ -126,6 +129,10 @@ async def connect(
         # Wait for it to complete
         await self._connect_task
 
+    def __hash__(self):
+        # to allow Devices to be used as dict keys
+        return hash(id(self))
+
 
 DeviceT = TypeVar("DeviceT", bound=Device)
 
@@ -180,6 +187,9 @@ def children(self) -> Iterator[tuple[str, Device]]:
         for key, child in self._children.items():
             yield str(key), child
 
+    def __hash__(self):  # to allow DeviceVector to be used as dict keys
+        return hash(id(self))
+
 
 class DeviceCollector:
     """Collector of top level Device instances to be used as a context manager

src/ophyd_async/core/_mock_signal_backend.py

@@ -16,7 +16,7 @@ class MockSignalBackend(SignalBackend[SignalDatatypeT]):
     def __init__(
         self,
         initial_backend: SignalBackend[SignalDatatypeT],
-        mock: bool | Mock = True,
+        mock: Mock,
     ) -> None:
         if isinstance(initial_backend, MockSignalBackend):
             raise ValueError("Cannot make a MockSignalBackend for a MockSignalBackend")
@@ -33,7 +33,7 @@ def __init__(
             )
 
         # use existing Mock if provided
-        self.mock = Mock() if isinstance(mock, bool) else mock
+        self.mock = mock
         self.mock.attach_mock(AsyncMock(name="put", spec=Callable), "put")
 
         super().__init__(datatype=self.initial_backend.datatype)

src/ophyd_async/core/_mock_signal_utils.py

@@ -3,19 +3,17 @@
 from typing import Any
 from unittest.mock import AsyncMock, Mock
 
-from ._device import Device
+from ._device import Device, _device_mocks
 from ._mock_signal_backend import MockSignalBackend
-from ._signal import Signal
+from ._signal import Signal, _mock_signal_backends
 from ._soft_signal_backend import SignalDatatypeT
 
 
 def _get_mock_signal_backend(signal: Signal) -> MockSignalBackend:
-    backend = signal._connector.backend  # noqa:SLF001
-    assert isinstance(backend, MockSignalBackend), (
-        "Expected to receive a `MockSignalBackend`, instead "
-        f" received {type(backend)}. "
-    )
-    return backend
+    assert (
+        signal in _mock_signal_backends
+    ), "No `MockSignalBackend` registered for signal."
+    return _mock_signal_backends[signal]
 
 
 def set_mock_value(signal: Signal[SignalDatatypeT], value: SignalDatatypeT):
@@ -51,7 +49,7 @@ def get_mock_put(signal: Signal) -> AsyncMock:
 def get_mock(device: Device | Signal) -> Mock:
     if isinstance(device, Signal):
         return _get_mock_signal_backend(device).mock
-    return device.mock
+    return _device_mocks[device]
 
 
 def reset_mock_put_calls(signal: Signal):

src/ophyd_async/core/_signal.py

@@ -32,6 +32,8 @@
 from ._status import AsyncStatus
 from ._utils import CALCULATE_TIMEOUT, DEFAULT_TIMEOUT, CalculatableTimeout, Callback, T
 
+_mock_signal_backends: dict[Device, MockSignalBackend] = {}
+
 
 async def _wait_for(coro: Awaitable[T], timeout: float | None, source: str) -> T:
     try:
@@ -61,6 +63,7 @@ async def connect(
     ):
         if mock:
             self.backend = MockSignalBackend(self._init_backend, mock)
+            _mock_signal_backends[device] = self.backend
         else:
             self.backend = self._init_backend
         device.log.debug(f"Connecting to {self.backend.source(device.name, read=True)}")

tests/core/test_mock_signal_backend.py

@@ -135,10 +135,7 @@ async def test_mock_utils_throw_error_if_backend_isnt_mock_signal_backend():
     exc_msgs.append(str(exc.value))
 
     for msg in exc_msgs:
-        assert msg == (
-            "Expected to receive a `MockSignalBackend`, instead "
-            f" received {SoftSignalBackend}. "
-        )
+        assert msg == "No `MockSignalBackend` registered for signal."
 
 
 async def test_get_mock_put():

tests/core/test_signal.py

@@ -3,7 +3,7 @@
 import re
 import time
 from asyncio import Event
-from unittest.mock import ANY
+from unittest.mock import ANY, Mock
 
 import pytest
 from bluesky.protocols import Reading
@@ -42,7 +42,7 @@ def num_occurrences(substring: str, string: str) -> int:
 
 async def test_signal_connects_to_previous_backend(caplog):
     caplog.set_level(logging.DEBUG)
-    int_mock_backend = MockSignalBackend(SoftSignalBackend(int))
+    int_mock_backend = MockSignalBackend(SoftSignalBackend(int), Mock())
     original_connect = int_mock_backend.connect
     times_backend_connect_called = 0
 
@@ -61,7 +61,7 @@ async def new_connect(timeout=1):
 
 async def test_signal_connects_with_force_reconnect(caplog):
     caplog.set_level(logging.DEBUG)
-    signal = Signal(MockSignalBackend(SoftSignalBackend(int)))
+    signal = Signal(MockSignalBackend(SoftSignalBackend(int), Mock()))
     await signal.connect()
     assert num_occurrences(f"Connecting to {signal.source}", caplog.text) == 1
     await signal.connect(force_reconnect=True)
@@ -80,7 +80,7 @@ async def connect(self, timeout=DEFAULT_TIMEOUT):
                 self.succeed_on_connect = True
                 raise RuntimeError("connect fail")
 
-    signal = SignalRW(MockSignalBackendFailingFirst(SoftSignalBackend(int)))
+    signal = SignalRW(MockSignalBackendFailingFirst(SoftSignalBackend(int), Mock()))
 
     with pytest.raises(RuntimeError, match="connect fail"):
         await signal.connect(mock=False)