From ff016cea6f2f00331b151ef6cdeff5be7d50bcdc Mon Sep 17 00:00:00 2001
From: Tom Cobb <tom.cobb@diamond.ac.uk>
Date: Fri, 8 Nov 2024 14:27:31 +0000
Subject: [PATCH] An idea...

---
 src/ophyd_async/core/_device.py | 12 ++++++------
 src/ophyd_async/core/_signal.py |  2 +-
 tests/core/test_device.py       |  4 ++--
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/ophyd_async/core/_device.py b/src/ophyd_async/core/_device.py
index 365e5714a..2d05fc9e8 100644
--- a/src/ophyd_async/core/_device.py
+++ b/src/ophyd_async/core/_device.py
@@ -35,10 +35,10 @@ def create_children_from_annotations(self, device: Device):
         during ``__init__``.
         """
 
-    async def connect_mock(self, device: Device, mock: LazyMock):
+    def connect_mock(self, device: Device, mock: LazyMock):
         # Connect serially, no errors to gather up as in mock mode
         for name, child_device in device.children():
-            await child_device.connect(mock=mock.child(name))
+            assert not child_device.connect(mock=mock.child(name))
 
     async def connect_real(self, device: Device, timeout: float, force_reconnect: bool):
         """Used during ``Device.connect``.
@@ -123,12 +123,12 @@ def __setattr__(self, name: str, value: Any) -> None:
         # Avoid the super call as this happens a lot
         return object.__setattr__(self, name, value)
 
-    async def connect(
+    def connect(
         self,
         mock: bool | LazyMock = False,
         timeout: float = DEFAULT_TIMEOUT,
         force_reconnect: bool = False,
-    ) -> None:
+    ) -> asyncio.Task | None:
         """Connect self and all child Devices.
 
         Contains a timeout that gets propagated to child.connect methods.
@@ -148,7 +148,7 @@ async def connect(
             elif not self._mock:
                 # Make one
                 self._mock = LazyMock()
-            await self._connector.connect_mock(self, self._mock)
+            self._connector.connect_mock(self, self._mock)
         else:
             # Try to cache the connect in real mode
             can_use_previous_connect = (
@@ -162,7 +162,7 @@ async def connect(
                 self._connect_task = asyncio.create_task(coro)
             assert self._connect_task, "Connect task not created, this shouldn't happen"
             # Wait for it to complete
-            await self._connect_task
+            return self._connect_task
 
 
 _not_device_attrs = {
diff --git a/src/ophyd_async/core/_signal.py b/src/ophyd_async/core/_signal.py
index d4d4d7ffe..d2f7985bb 100644
--- a/src/ophyd_async/core/_signal.py
+++ b/src/ophyd_async/core/_signal.py
@@ -58,7 +58,7 @@ class SignalConnector(DeviceConnector):
     def __init__(self, backend: SignalBackend):
         self.backend = self._init_backend = backend
 
-    async def connect_mock(self, device: Device, mock: LazyMock):
+    def connect_mock(self, device: Device, mock: LazyMock):
         self.backend = MockSignalBackend(self._init_backend, mock)
 
     async def connect_real(self, device: Device, timeout: float, force_reconnect: bool):
diff --git a/tests/core/test_device.py b/tests/core/test_device.py
index 70620a0bf..f091e59f2 100644
--- a/tests/core/test_device.py
+++ b/tests/core/test_device.py
@@ -177,9 +177,9 @@ def __init__(self, name: str) -> None:
 
 async def test_many_individual_device_connects_not_slow():
     start = time.time()
-    for i in range(100):
+    for i in range(5000):
         bundle = MotorBundle(f"bundle{i}")
-        await bundle.connect(mock=True)
+        bundle.connect(mock=True)
     duration = time.time() - start
     assert duration < 1