From ec082ae616d004fe2b155ee402eb136290246f0a Mon Sep 17 00:00:00 2001
From: Tom Willemsen <tom.willemsen@stfc.ac.uk>
Date: Fri, 1 Nov 2024 09:06:46 +0000
Subject: [PATCH] Try wait_for_wakeups

---
 tests/epics/test_motor.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/tests/epics/test_motor.py b/tests/epics/test_motor.py
index cba0fde7e..a87122e78 100644
--- a/tests/epics/test_motor.py
+++ b/tests/epics/test_motor.py
@@ -18,9 +18,17 @@
 )
 from ophyd_async.epics import motor
 
-# Long enough for multiple asyncio event loop cycles to run so
-# all the tasks have a chance to run
-A_BIT = 0.001
+
+async def wait_for_wakeups(max_yields=10):
+    loop = asyncio.get_event_loop()
+    # If anything has called loop.call_soon or is scheduled a wakeup
+    # then let it run
+    for _ in range(max_yields):
+        if loop._scheduled:
+            await asyncio.sleep(0)
+        else:
+            return
+    raise RuntimeError(f"Tasks still scheduling wakeups after {max_yields} yields")
 
 
 @pytest.fixture
@@ -37,7 +45,7 @@ async def sim_motor():
 async def wait_for_eq(item, attribute, comparison, timeout):
     timeout_time = time.monotonic() + timeout
     while getattr(item, attribute) != comparison:
-        await asyncio.sleep(A_BIT)
+        await wait_for_wakeups()
         if time.monotonic() > timeout_time:
             raise TimeoutError
 
@@ -49,6 +57,7 @@ async def test_motor_moving_well(sim_motor: motor.Motor) -> None:
     s.watch(watcher)
     done = Mock()
     s.add_callback(done)
+    await wait_for_wakeups()
     await wait_for_eq(watcher, "call_count", 1, 1)
     assert watcher.call_args == call(
         name="sim_motor",
@@ -78,7 +87,7 @@ async def test_motor_moving_well(sim_motor: motor.Motor) -> None:
     set_mock_value(sim_motor.motor_done_move, True)
     set_mock_value(sim_motor.user_readback, 0.55)
     set_mock_put_proceeds(sim_motor.user_setpoint, True)
-    await asyncio.sleep(A_BIT)
+    await wait_for_wakeups()
     await wait_for_eq(s, "done", True, 1)
     done.assert_called_once_with(s)
 
@@ -90,7 +99,7 @@ async def test_motor_moving_well_2(sim_motor: motor.Motor) -> None:
     s.watch(watcher)
     done = Mock()
     s.add_callback(done)
-    await asyncio.sleep(0.1)
+    await wait_for_wakeups()
     assert watcher.call_count == 1
     assert watcher.call_args == call(
         name="sim_motor",
@@ -118,7 +127,7 @@ async def test_motor_moving_well_2(sim_motor: motor.Motor) -> None:
         time_elapsed=pytest.approx(0.1, abs=0.2),
     )
     set_mock_put_proceeds(sim_motor.user_setpoint, True)
-    await asyncio.sleep(0.1)
+    await wait_for_wakeups()
     assert s.done
     done.assert_called_once_with(s)
 
@@ -157,7 +166,7 @@ async def test_motor_moving_stopped(sim_motor: motor.Motor):
     assert not s.done
     await sim_motor.stop()
     set_mock_put_proceeds(sim_motor.user_setpoint, True)
-    await asyncio.sleep(0.1)
+    await wait_for_wakeups()
     assert s.done
     assert s.success is False