From b0b211404ee9b2d6840e072727a777e359bc0178 Mon Sep 17 00:00:00 2001
From: Dominic Oram <dominic.oram@diamond.ac.uk>
Date: Tue, 16 Jul 2024 13:02:32 +0100
Subject: [PATCH] Use set and wait in ophyd_async

---
 src/dodal/devices/tetramm.py           |  7 ++--
 src/dodal/devices/xspress3/xspress3.py | 55 ++++++++++++++++++++------
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/src/dodal/devices/tetramm.py b/src/dodal/devices/tetramm.py
index 7ced99c28e..3d76b5198d 100644
--- a/src/dodal/devices/tetramm.py
+++ b/src/dodal/devices/tetramm.py
@@ -11,8 +11,8 @@
     DirectoryProvider,
     ShapeProvider,
     StandardDetector,
-    set_and_wait_for_value,
     soft_signal_r_and_setter,
+    wait_for_value,
 )
 from ophyd_async.epics.areadetector.utils import stop_busy_record
 from ophyd_async.epics.areadetector.writers import HDFWriter, NDFileHDF
@@ -132,9 +132,10 @@ async def arm(
             self._drv.averaging_time.set(exposure), self.set_exposure(exposure)
         )
 
-        status = await set_and_wait_for_value(self._drv.acquire, 1)
+        finished_status = self._drv.acquire.set(True)
+        await wait_for_value(self._drv.acquire, True, 10)
 
-        return status
+        return finished_status
 
     def _validate_trigger(self, trigger: DetectorTrigger) -> None:
         supported_trigger_types = {
diff --git a/src/dodal/devices/xspress3/xspress3.py b/src/dodal/devices/xspress3/xspress3.py
index f8ae574f0a..207f96fb62 100644
--- a/src/dodal/devices/xspress3/xspress3.py
+++ b/src/dodal/devices/xspress3/xspress3.py
@@ -7,7 +7,7 @@
     AsyncStatus,
     Device,
     DeviceVector,
-    wait_for_value,
+    set_and_wait_for_other_value,
 )
 from ophyd_async.epics.signal.signal import (
     epics_signal_r,
@@ -45,6 +45,16 @@ class AcquireRBVState(str, Enum):
     ACQUIRE = "Acquiring"
 
 
+class WritingRBVState(str, Enum):
+    DONE = "Done"
+    CAPTURE = "Capturing"
+
+
+class WritingState(str, Enum):
+    DONE = "Done"
+    CAPTURE = "Capture"
+
+
 class DetectorState(str, Enum):
     IDLE = "Idle"
     ACQUIRE = "Acquire"
@@ -115,6 +125,13 @@ def __init__(
         self.acquire_rbv = epics_signal_r(AcquireRBVState, prefix + "Acquire_RBV")
         self.trigger_mode = epics_signal_rw_rbv(TriggerMode, prefix + "TriggerMode")
 
+        self.file_path = epics_signal_rw_rbv(str, prefix + "HDF5:FilePath")
+        self.file_name = epics_signal_rw_rbv(str, prefix + "HDF5:FileName")
+        self.do_file_writing = epics_signal_rw(WritingState, prefix + "HDF5:Capture")
+        self.file_writing_rbv = epics_signal_r(
+            WritingRBVState, prefix + "HDF5:Capture_RBV"
+        )
+
         self.detector_state = epics_signal_r(
             DetectorState, prefix + "DetectorState_RBV"
         )
@@ -130,21 +147,37 @@ def __init__(
     @AsyncStatus.wrap
     async def stage(self) -> None:
         LOGGER.info("Arming Xspress3 detector...")
-        await self.trigger_mode.set(TriggerMode.BURST)
-        await wait_for_value(
-            self.detector_state,
-            lambda v: v in self.detector_busy_states,
+        await set_and_wait_for_other_value(
+            self.do_file_writing,
+            WritingState.CAPTURE,
+            self.file_writing_rbv,
+            WritingRBVState.CAPTURE,
             timeout=self.timeout,
         )
-        await self.acquire.set(AcquireState.ACQUIRE)
-        await wait_for_value(
-            self.acquire_rbv, AcquireRBVState.ACQUIRE, timeout=self.timeout
+        await self.trigger_mode.set(TriggerMode.BURST)
+
+        await set_and_wait_for_other_value(
+            self.acquire,
+            AcquireState.ACQUIRE,
+            self.acquire_rbv,
+            AcquireRBVState.ACQUIRE,
+            timeout=self.timeout,
         )
 
     @AsyncStatus.wrap
     async def unstage(self) -> None:
-        await self.acquire.set(AcquireState.DONE)
         LOGGER.info("unstaging Xspress3 detector...")
-        await wait_for_value(
-            self.acquire_rbv, AcquireRBVState.DONE, timeout=self.timeout
+        await set_and_wait_for_other_value(
+            self.do_file_writing,
+            WritingState.DONE,
+            self.file_writing_rbv,
+            WritingRBVState.DONE,
+            timeout=self.timeout,
+        )
+        await set_and_wait_for_other_value(
+            self.acquire,
+            AcquireState.DONE,
+            self.acquire_rbv,
+            AcquireRBVState.DONE,
+            timeout=self.timeout,
         )