From d0c81d20e71c482662e507fe84306de24f25896f Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Mon, 29 Jul 2024 09:49:25 +0200
Subject: [PATCH] fix: improve error handling in wait_for_status and
 _move_and_finish

---
 debye_bec/devices/mo1_bragg.py             | 25 ++++++++++++++++------
 tests/tests_scans/test_mono_bragg_scans.py |  1 +
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/debye_bec/devices/mo1_bragg.py b/debye_bec/devices/mo1_bragg.py
index c07e4b7..d09d8ec 100644
--- a/debye_bec/devices/mo1_bragg.py
+++ b/debye_bec/devices/mo1_bragg.py
@@ -30,6 +30,7 @@ from ophyd import (
 )
 from ophyd.utils import LimitError
 from ophyd_devices.utils import bec_scaninfo_mixin, bec_utils
+from ophyd_devices.utils.errors import DeviceStopError, DeviceTimeoutError
 
 logger = bec_logger.logger
 
@@ -410,7 +411,6 @@ class Mo1Bragg(Device, PositionerBase):
             status (DeviceStatus)   : status object to set the status of the motion
             update_frequency (float): Optional, frequency to update the current position of the motion, defaults to 0.1s
         """
-        success = True
         try:
             # Set the target position on IOC
             move_cpt.put(target_pos)
@@ -419,11 +419,10 @@ class Mo1Bragg(Device, PositionerBase):
             time.sleep(0.5)
             while self.motor_is_moving.get() == 0:
                 if self.stopped:
-                    success = False
-                    break
+                    raise DeviceStopError(f"Device {self.name} was stopped")
                 time.sleep(update_frequency)
             # pylint: disable=protected-access
-            status._finished(success=success)
+            status.set_finished()
         # pylint: disable=broad-except
         except Exception as exc:
             content = traceback.format_exc()
@@ -799,7 +798,7 @@ class Mo1Bragg(Device, PositionerBase):
         check_stopped: bool = False,
         interval: float = 0.05,
         all_signals: bool = False,
-        exception_on_timeout: Exception = TimeoutError("Timeout while waiting for signals"),
+        exception_on_timeout: Exception = None,
     ) -> DeviceStatus:
         """Wrapper around wait_for_signals to be started in thread and attach a DeviceStatus object.
         This allows BEC to perform actinos in parallel and not be blocked by method calls on a device.
@@ -816,6 +815,10 @@ class Mo1Bragg(Device, PositionerBase):
         Returns:
             DeviceStatus: DeviceStatus object that resolves either to set_finished or set_exception
         """
+        if exception_on_timeout is None:
+            exception_on_timeout = DeviceTimeoutError(
+                f"Timeout error for {self.name} while waiting for signals {signal_conditions}"
+            )
 
         status = DeviceStatus(device=self)
 
@@ -826,7 +829,7 @@ class Mo1Bragg(Device, PositionerBase):
             check_stopped: bool,
             interval: float,
             all_signals: bool,
-            exception_on_timeout: Exception = TimeoutError("Timeout while waiting for signals"),
+            exception_on_timeout: Exception = None,
         ):
             """Convenient wrapper around wait_for_signals to set status based on the result.
 
@@ -847,9 +850,17 @@ class Mo1Bragg(Device, PositionerBase):
                     # pylint: disable=protected-access
                     status.set_finished()
                 else:
-                    status.set_exception(exception_on_timeout)
+                    if self.parent.stopped:
+                        # INFO This will execute a callback to the parent device.stop() method
+                        status.set_exception(exc=DeviceStopError(f"{self.name} was stopped"))
+                    else:
+                        # INFO This will execute a callback to the parent device.stop() method
+                        status.set_exception(exc=exception_on_timeout)
             # pylint: disable=broad-except
             except Exception as exc:
+                content = traceback.format_exc()
+                logger.warning(f"Error in wait_for_signals in {self.name}; Traceback: {content}")
+                # INFO This will execute a callback to the parent device.stop() method
                 status.set_exception(exc=exc)
 
         thread = threading.Thread(
diff --git a/tests/tests_scans/test_mono_bragg_scans.py b/tests/tests_scans/test_mono_bragg_scans.py
index 235f3dc..0d6a567 100644
--- a/tests/tests_scans/test_mono_bragg_scans.py
+++ b/tests/tests_scans/test_mono_bragg_scans.py
@@ -1,3 +1,4 @@
+# pylint: skip-file
 from unittest import mock
 
 from bec_lib.messages import DeviceInstructionMessage