From cbe04343ff06d32d7a1ccdb0075bc61e862367d0 Mon Sep 17 00:00:00 2001
From: Zoheb Shaikh <zoheb.shaikh@diamond.ac.uk>
Date: Wed, 14 Aug 2024 15:21:35 +0100
Subject: [PATCH] updated to collect_while_completing

---
 pyproject.toml                     |  2 +-
 src/ophyd_async/plan_stubs/_fly.py | 25 +++----------------------
 tests/plan_stubs/test_fly.py       |  5 ++---
 3 files changed, 6 insertions(+), 26 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7f105afc1b..1f93ad1ccc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@ dependencies = [
     "numpy<2.0.0",
     "packaging",
     "pint",
-    "bluesky>=1.13.0a3",
+    "bluesky>=1.13.0a4",
     "event_model",
     "p4p",
     "pyyaml",
diff --git a/src/ophyd_async/plan_stubs/_fly.py b/src/ophyd_async/plan_stubs/_fly.py
index 087ec62dd1..014faa5ff6 100644
--- a/src/ophyd_async/plan_stubs/_fly.py
+++ b/src/ophyd_async/plan_stubs/_fly.py
@@ -1,7 +1,6 @@
 from typing import List, Optional
 
 import bluesky.plan_stubs as bps
-from bluesky.utils import short_uid
 
 from ophyd_async.core import (
     DetectorTrigger,
@@ -119,27 +118,9 @@ def fly_and_collect(
     for detector in detectors:
         yield from bps.kickoff(detector)
 
-    # collect_while_completing
-    group = short_uid(label="complete")
-
-    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)
-        except TimeoutError:
-            pass
-        else:
-            done = True
-        yield from bps.collect(
-            *detectors,
-            return_payload=False,
-            name=stream_name,
-        )
-    yield from bps.wait(group=group)
+    yield from bps.collect_while_completing(
+        flyers=[flyer], dets=detectors, flush_period=0.5, stream_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 791935c111..441e71eebb 100644
--- a/tests/plan_stubs/test_fly.py
+++ b/tests/plan_stubs/test_fly.py
@@ -114,9 +114,8 @@ def __init__(
         super().__init__(controller, writer, config_sigs, name)
 
     @WatchableAsyncStatus.wrap
-    async def complete(self):
-        assert self._arm_status, "Prepare not run"
-        assert self._trigger_info
+    async def kickoff(self):
+        super().kickoff()
         self.writer.increment_index()
         async for index in self.writer.observe_indices_written(
             self._trigger_info.frame_timeout