diff --git a/src/ophyd_async/planstubs/ensure_connected.py b/src/ophyd_async/planstubs/ensure_connected.py index cb4b5caa5f..78607be36c 100644 --- a/src/ophyd_async/planstubs/ensure_connected.py +++ b/src/ophyd_async/planstubs/ensure_connected.py @@ -4,14 +4,19 @@ from ophyd_async.core.utils import DEFAULT_TIMEOUT, wait_for_connection -def ensure_connected(*devices: Device, - sim: bool = False, - timeout: float = DEFAULT_TIMEOUT, - force_reconnect=False): - yield from bps.wait_for([lambda : - wait_for_connection(**{ - device.name: - device.connect(sim, timeout, force_reconnect) - for device in devices - }) - ]) +def ensure_connected( + *devices: Device, + sim: bool = False, + timeout: float = DEFAULT_TIMEOUT, + force_reconnect=False, +): + yield from bps.wait_for( + [ + lambda: wait_for_connection( + **{ + device.name: device.connect(sim, timeout, force_reconnect) + for device in devices + } + ) + ] + ) diff --git a/tests/core/test_device.py b/tests/core/test_device.py index eebf1b265b..2cd81dae90 100644 --- a/tests/core/test_device.py +++ b/tests/core/test_device.py @@ -137,20 +137,29 @@ async def test_device_lazily_connects(RE): and sim_motor._connect_task.done() and not sim_motor._connect_task.exception() ) + + 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")} + 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") - assert parentMotor._connect_task is None and parentMotor.X._connect_task is None and parentMotor.Y._connect_task is None - + assert ( + parentMotor._connect_task is None + and parentMotor.X._connect_task is None + and parentMotor.Y._connect_task is None + ) RE(ensure_connected(parentMotor, sim=True)) @@ -169,9 +178,9 @@ async def test_device_with_children_lazily_connects(RE): and parentMotor._connect_task.done() and not parentMotor._connect_task.exception() ) - for motor in parentMotor.V.values(): + for vector in parentMotor.V.values(): assert ( - motor._connect_task - and motor._connect_task.done() - and not motor._connect_task.exception() - ) \ No newline at end of file + vector._connect_task + and vector._connect_task.done() + and not vector._connect_task.exception() + )