From ecc6111c7e75ae68afe571a5878a6033baf14941 Mon Sep 17 00:00:00 2001
From: Devin Burke <devin.burke@desy.de>
Date: Thu, 23 May 2024 16:23:46 +0200
Subject: [PATCH] Tuning sample device definitions

---
 .../tango/device_controllers/dgg2.py          | 31 ++++++++++++++++---
 .../tango/device_controllers/omsvme58.py      |  4 ++-
 .../tango/device_controllers/sis3820.py       |  7 +++--
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/ophyd_async/tango/device_controllers/dgg2.py b/src/ophyd_async/tango/device_controllers/dgg2.py
index 6321d1c0f0..8de6ce5b56 100644
--- a/src/ophyd_async/tango/device_controllers/dgg2.py
+++ b/src/ophyd_async/tango/device_controllers/dgg2.py
@@ -1,17 +1,23 @@
 from __future__ import annotations
 
+import asyncio
 import time
 from typing import Callable, List
 
 from bluesky.protocols import Preparable, Triggerable
 
 from ophyd_async.core import AsyncStatus
-from ophyd_async.tango import TangoReadableDevice, tango_signal_rw, tango_signal_x
+from ophyd_async.tango import (
+    TangoReadableDevice,
+    tango_signal_r,
+    tango_signal_rw,
+    tango_signal_x,
+)
+from tango import DevState
 
 
 # --------------------------------------------------------------------
 class DGG2Timer(TangoReadableDevice, Triggerable, Preparable):
-
     # --------------------------------------------------------------------
     def __init__(self, trl: str, name="") -> None:
         TangoReadableDevice.__init__(self, trl, name)
@@ -19,7 +25,6 @@ def __init__(self, trl: str, name="") -> None:
 
     # --------------------------------------------------------------------
     def register_signals(self):
-
         self.sampletime = tango_signal_rw(
             float, self.trl + "/sampletime", device_proxy=self.proxy
         )
@@ -35,6 +40,10 @@ def register_signals(self):
             self.trl + "/startandwaitfortimer", device_proxy=self.proxy
         )
 
+        self.start = tango_signal_x(self.trl + "/start", device_proxy=self.proxy)
+
+        self._state = tango_signal_r(DevState, self.trl + "/state", self.proxy)
+
         self.set_name(self.name)
 
     # --------------------------------------------------------------------
@@ -58,12 +67,24 @@ def update_watchers(remaining_time: float):
         else:
             update_watchers(total_time)
         try:
-            await self.startandwaitfortimer.trigger()
+            await self.start.trigger()
         finally:
             if self.remainingtime.is_cachable():
                 self.remainingtime.clear_sub(update_watchers)
             else:
-                update_watchers(await self.remainingtime.get_value())
+                counter = 0
+                state = await self._state.get_value()
+                while state == DevState.MOVING:
+                    # Update the watchers with the current position every 0.5 seconds
+                    if counter % 5 == 0:
+                        remaining_time = await self.remainingtime.get_value()
+                        update_watchers(remaining_time)
+                        counter = 0
+                    await asyncio.sleep(0.1)
+                    state = await self._state.get_value()
+                    counter += 1
+
+                # update_watchers(await self.remainingtime.get_value())
         if not self._set_success:
             raise RuntimeError("Timer was not triggered")
 
diff --git a/src/ophyd_async/tango/device_controllers/omsvme58.py b/src/ophyd_async/tango/device_controllers/omsvme58.py
index 8c51c0bd33..0f62f4a400 100644
--- a/src/ophyd_async/tango/device_controllers/omsvme58.py
+++ b/src/ophyd_async/tango/device_controllers/omsvme58.py
@@ -75,13 +75,15 @@ def update_watchers(current_position: float):
             await self.position.set(new_position)
             await asyncio.sleep(0.1)
             counter = 0
-            while await self._state.get_value() == DevState.MOVING:
+            state = await self._state.get_value()
+            while state == DevState.MOVING:
                 # Update the watchers with the current position every 0.5 seconds
                 if counter % 5 == 0:
                     current_position = await self.position.get_value()
                     update_watchers(current_position)
                     counter = 0
                 await asyncio.sleep(0.1)
+                state = await self._state.get_value()
                 counter += 1
         finally:
             if self.position.is_cachable():
diff --git a/src/ophyd_async/tango/device_controllers/sis3820.py b/src/ophyd_async/tango/device_controllers/sis3820.py
index 845e23d2de..fd56a62fe8 100644
--- a/src/ophyd_async/tango/device_controllers/sis3820.py
+++ b/src/ophyd_async/tango/device_controllers/sis3820.py
@@ -10,7 +10,6 @@
 
 # --------------------------------------------------------------------
 class SIS3820Counter(TangoReadableDevice):  # Triggerable
-
     # --------------------------------------------------------------------
     def __init__(self, trl: str, name="") -> None:
         TangoReadableDevice.__init__(self, trl, name)
@@ -18,7 +17,6 @@ def __init__(self, trl: str, name="") -> None:
 
     # --------------------------------------------------------------------
     def register_signals(self):
-
         self.counts = tango_signal_rw(
             float, self.trl + "/counts", device_proxy=self.proxy
         )
@@ -40,5 +38,8 @@ def register_signals(self):
     # --------------------------------------------------------------------
     async def read(self) -> Dict[str, Reading]:
         ret = await super().read()
-        await self.reset.trigger()
         return ret
+
+    async def stage(self):
+        await self.reset.trigger()
+        await super().stage()