Device connect lazily

src/ophyd_async/core/device.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import sys
 from functools import cached_property
 from logging import LoggerAdapter, getLogger
@@ -33,6 +34,10 @@ class Device(HasName):
     #: The parent Device if it exists
     parent: Optional[Device] = None
 
+    # Previous connect was successful, False on initialization
+    # since there was no previous connect
+    _previous_connect_success: bool = False
+
     def __init__(self, name: str = "") -> None:
         self.set_name(name)
 
@@ -71,7 +76,12 @@ def set_name(self, name: str):
             child.set_name(child_name)
             child.parent = self
 
-    async def connect(self, mock: bool = False, timeout: float = DEFAULT_TIMEOUT):
+    async def connect(
+        self,
+        mock: bool = False,
+        timeout: float = DEFAULT_TIMEOUT,
+        force_reconnect=False,
+    ):
         """Connect self and all child Devices.
 
         Contains a timeout that gets propagated to child.connect methods.
@@ -83,12 +93,18 @@ async def connect(self, mock: bool = False, timeout: float = DEFAULT_TIMEOUT):
         timeout:
             Time to wait before failing with a TimeoutError.
         """
-        coros = {
-            name: child_device.connect(mock=mock, timeout=timeout)
-            for name, child_device in self.children()
-        }
-        if coros:
-            await wait_for_connection(**coros)
+
+        if force_reconnect or not self._previous_connect_success:
+            # Kick off a connection
+            coros = {
+                name: child_device.connect(mock, timeout=timeout)
+                for name, child_device in self.children()
+            }
+            connect_task = asyncio.create_task(wait_for_connection(**coros))
+
+            # Wait for it to complete
+            await connect_task
+            self._previous_connect_success = not connect_task.exception()
 
 
 VT = TypeVar("VT", bound=Device)

src/ophyd_async/epics/pvi/pvi.py

@@ -91,7 +91,7 @@ def _verify_common_blocks(entry: PVIEntry, common_device: Type[Device]):
         return
     common_sub_devices = get_type_hints(common_device)
     for sub_name, sub_device in common_sub_devices.items():
-        if sub_name in ("_name", "parent"):
+        if sub_name.startswith("_") or sub_name == "parent":
             continue
         assert entry.sub_entries
         device_t, is_optional = _strip_union(sub_device)
@@ -161,7 +161,7 @@ def _mock_common_blocks(device: Device, stripped_type: Optional[Type] = None):
     sub_devices = (
         (field, field_type)
         for field, field_type in get_type_hints(device_t).items()
-        if field not in ("_name", "parent")
+        if not field.startswith("_") and field != "parent"
     )
 
     for device_name, device_cls in sub_devices:

src/ophyd_async/planstubs/__init__.py

@@ -1,5 +1,9 @@
+from .ensure_connected import ensure_connected
 from .prepare_trigger_and_dets import (
     prepare_static_seq_table_flyer_and_detectors_with_same_trigger,
 )
 
-__all__ = ["prepare_static_seq_table_flyer_and_detectors_with_same_trigger"]
+__all__ = [
+    "prepare_static_seq_table_flyer_and_detectors_with_same_trigger",
+    "ensure_connected",
+]

src/ophyd_async/planstubs/ensure_connected.py

@@ -0,0 +1,22 @@
+import bluesky.plan_stubs as bps
+
+from ophyd_async.core.device import Device
+from ophyd_async.core.utils import DEFAULT_TIMEOUT, wait_for_connection
+
+
+def ensure_connected(
+    *devices: Device,
+    mock: bool = False,
+    timeout: float = DEFAULT_TIMEOUT,
+    force_reconnect=False,
+):
+    yield from bps.wait_for(
+        [
+            lambda: wait_for_connection(
+                **{
+                    device.name: device.connect(mock, timeout, force_reconnect)
+                    for device in devices
+                }
+            )
+        ]
+    )

tests/core/test_device.py

@@ -1,5 +1,6 @@
 import asyncio
 import traceback
+from unittest.mock import Mock
 
 import pytest
 
@@ -11,6 +12,10 @@
     NotConnected,
     wait_for_connection,
 )
+from ophyd_async.epics.motion import motor
+from ophyd_async.planstubs.ensure_connected import (
+    ensure_connected,
+)
 
 
 class DummyBaseDevice(Device):
@@ -117,3 +122,79 @@ async def test_device_log_has_correct_name():
     assert device.log.extra["ophyd_async_device_name"] == ""
     device.set_name("device")
     assert device.log.extra["ophyd_async_device_name"] == "device"
+
+
+async def test_device_lazily_connects(RE):
+    async with DeviceCollector(mock=True, connect=False):
+        mock_motor = motor.Motor("BLxxI-MO-TABLE-01:X")
+
+    assert not mock_motor._previous_connect_success
+
+    # When ready to connect
+    RE(ensure_connected(mock_motor, mock=True))
+
+    assert mock_motor._previous_connect_success
+
+
+class MotorBundle(Device):
+    def __init__(self, name: str) -> None:
+        self.X = motor.Motor("BLxxI-MO-TABLE-01:X")
+        self.Y = motor.Motor("BLxxI-MO-TABLE-01:Y")
+        self.V: DeviceVector[motor.Motor] = DeviceVector(
+            {
+                0: motor.Motor("BLxxI-MO-TABLE-21:X"),
+                1: motor.Motor("BLxxI-MO-TABLE-21:Y"),
+                2: motor.Motor("BLxxI-MO-TABLE-21:Z"),
+            }
+        )
+
+
+async def test_device_with_children_lazily_connects(RE):
+    parentMotor = MotorBundle("parentMotor")
+
+    for device in [parentMotor, parentMotor.X, parentMotor.Y] + list(
+        parentMotor.V.values()
+    ):
+        assert not device._previous_connect_success
+    RE(ensure_connected(parentMotor, mock=True))
+
+    for device in [parentMotor, parentMotor.X, parentMotor.Y] + list(
+        parentMotor.V.values()
+    ):
+        assert device._previous_connect_success
+
+
+async def test_device_with_device_collector_lazily_connects():
+    mock_motor = motor.Motor("NONE_EXISTENT")
+    with pytest.raises(NotConnected):
+        await mock_motor.connect(mock=False, timeout=0.01)
+    assert not mock_motor._previous_connect_success
+    await mock_motor.connect(mock=True, timeout=0.01)
+    assert mock_motor._previous_connect_success
+
+
+async def test_no_reconnect_signals_if_not_forced():
+    parent = DummyDeviceGroup("parent")
+
+    async def inner_connect(mock, timeout):
+        parent.child1.connected = True
+
+    parent.child1.connect = Mock(side_effect=inner_connect)
+    await parent.connect(mock=True, timeout=0.01)
+    assert parent.child1.connected
+    assert parent.child1.connect.call_count == 1
+    await parent.connect(mock=True, timeout=0.01)
+    assert parent.child1.connected
+    assert parent.child1.connect.call_count == 1
+
+    for count in range(2, 10):
+        await parent.connect(mock=True, timeout=0.01, force_reconnect=True)
+        assert parent.child1.connected
+        assert parent.child1.connect.call_count == count
+
+
+async def test_parent_not_connected():
+    parent = DummyDeviceGroup("parent")
+    await parent.child1.connect(mock=True, timeout=0.01)
+    assert not parent._previous_connect_success
+    assert not parent.child2._previous_connect_success

tests/epics/motion/test_motor.py

@@ -59,7 +59,7 @@ async def test_motor_moving_well(sim_motor: motor.Motor) -> None:
         target=0.55,
         unit="mm",
         precision=3,
-        time_elapsed=pytest.approx(0.1, abs=0.05),
+        time_elapsed=pytest.approx(0.1, abs=0.2),
     )
     set_mock_put_proceeds(sim_motor.user_setpoint, True)
     await asyncio.sleep(A_BIT)