debye_bec/devices/nidaq/nidaq.py

@@ -5,9 +5,9 @@ from typing import TYPE_CHECKING, Literal
 from bec_lib.logger import bec_logger
 from bec_server.scan_server.scans.scan_base import ScanInfo as ScanServerScanInfo
 from ophyd import Component as Cpt
-from ophyd import Device, DeviceStatus, EpicsSignal, EpicsSignalRO, Kind, StatusBase
+from ophyd import Device, DeviceStatus, EpicsSignal, EpicsSignalRO, Kind
 from ophyd.status import WaitTimeoutError
-from ophyd_devices import CompareStatus, ProgressSignal, TransitionStatus
+from ophyd_devices import CompareStatus, ProgressSignal, StatusBase, TransitionStatus
 from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase
 from ophyd_devices.sim.sim_signals import SetableSignal
 
@@ -432,13 +432,18 @@ class Nidaq(PSIDeviceBase, NidaqControl):
         status.wait(timeout=self.timeout_wait_for_signal)
         if self.scan_parameters.scan_name != "nidaq_continuous_scan":
             status = self.on_kickoff()
-            self.cancel_on_stop(status)
-            status.wait(timeout=self._timeout_wait_for_pv)
+            status.wait()
         logger.info(f"Device {self.name} was staged: {NidaqState(self.state.get())}")
 
     def on_kickoff(self) -> DeviceStatus | StatusBase:
         """Kickoff the Nidaq"""
-        status = self.kickoff_call.set(1)
+        status = CompareStatus(
+            self.state,
+            NidaqState.KICKOFF.value,
+            timeout=self._timeout_wait_for_pv,
+            description="Waiting for NIDAQ to enter KICKOFF state",
+        )
+        self.kickoff_call.put(1)
         self.cancel_on_stop(status)
         return status