From 5ddf953c34bfa7fbb49067fbef71c5912ae05cea Mon Sep 17 00:00:00 2001
From: Zoheb Shaikh <zoheb.shaikh@diamond.ac.uk>
Date: Wed, 28 Aug 2024 11:44:47 +0000
Subject: [PATCH] moved code from bluesky

---
 src/ophyd_async/plan_stubs/_fly.py | 18 +++++++++++++++---
 tests/plan_stubs/test_fly.py       |  5 +++--
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/ophyd_async/plan_stubs/_fly.py b/src/ophyd_async/plan_stubs/_fly.py
index 014faa5ff6..e0fa2e4743 100644
--- a/src/ophyd_async/plan_stubs/_fly.py
+++ b/src/ophyd_async/plan_stubs/_fly.py
@@ -1,6 +1,7 @@
 from typing import List, Optional
 
 import bluesky.plan_stubs as bps
+from bluesky.utils import short_uid
 
 from ophyd_async.core import (
     DetectorTrigger,
@@ -117,10 +118,21 @@ def fly_and_collect(
     yield from bps.kickoff(flyer, wait=True)
     for detector in detectors:
         yield from bps.kickoff(detector)
+    group = short_uid(label="complete")
 
-    yield from bps.collect_while_completing(
-        flyers=[flyer], dets=detectors, flush_period=0.5, stream_name=stream_name
-    )
+    yield from bps.complete(flyer, wait=False, group=group)
+    for detector in detectors:
+        yield from bps.complete(detector, wait=False, group=group)
+
+    done = False
+    while not done:
+        try:
+            yield from bps.wait(group=group, timeout=0.5, move_on=True)
+        except TimeoutError:
+            pass
+        else:
+            done = True
+        yield from bps.collect(*detectors, name=stream_name)
 
 
 def fly_and_collect_with_static_pcomp(
diff --git a/tests/plan_stubs/test_fly.py b/tests/plan_stubs/test_fly.py
index 441e71eebb..791935c111 100644
--- a/tests/plan_stubs/test_fly.py
+++ b/tests/plan_stubs/test_fly.py
@@ -114,8 +114,9 @@ def __init__(
         super().__init__(controller, writer, config_sigs, name)
 
     @WatchableAsyncStatus.wrap
-    async def kickoff(self):
-        super().kickoff()
+    async def complete(self):
+        assert self._arm_status, "Prepare not run"
+        assert self._trigger_info
         self.writer.increment_index()
         async for index in self.writer.observe_indices_written(
             self._trigger_info.frame_timeout