From 358fecadd42cbcf60e26ca8ff8f3109f0df61b91 Mon Sep 17 00:00:00 2001
From: smathis <stefan.mathis@psi.ch>
Date: Thu, 24 Jul 2025 09:16:34 +0200
Subject: [PATCH] Simplified enabling / disabling checking

---
 common.py                           | 45 ++++++++++++++++++++++++-----
 tests/sinqMotor/turboPmac/common.py | 18 ++----------
 tests/sinqMotor/turboPmac/reset.py  |  7 +++--
 3 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/common.py b/common.py
index b344137..e0a7709 100755
--- a/common.py
+++ b/common.py
@@ -100,7 +100,7 @@ class Motor:
 
     def move_and_wait(self, target):
         """
-        Move the motor to the given target and return, once the motor has 
+        Move the motor to the given target and return, once the motor has
         finished moving. The motor status is polled regulary to see if an error
         occurred during movement. In case this happens, an error is raised.
         """
@@ -108,7 +108,7 @@ class Motor:
         # Is the motor already at its target?
         if self.at_target(target):
             return
-        
+
         # Is the motor already moving?
         already_moving = self.read_field('donemoving') == 0
 
@@ -117,7 +117,7 @@ class Motor:
         if already_moving:
             # Wait until the motor has stopped
             self.wait_for_done()
-            
+
         self.wait_for_start()
         self.wait_for_done()
 
@@ -132,9 +132,9 @@ class Motor:
         """
         if target is None:
             target = self.read_field('writepv')
-            
-        return (self.read_field('miss') == 0 and 
-                math.isclose(self.read_field('readpv'), target, 
+
+        return (self.read_field('miss') == 0 and
+                math.isclose(self.read_field('readpv'), target,
                              rel_tol=1e-9, abs_tol=self.read_field('precision')))
 
     def wait_for_start(self, maxwaittime=None):
@@ -144,7 +144,8 @@ class Motor:
         now = time.time()
         while self.read_field('donemoving') != 0:
             if now + maxwaittime < time.time():
-                raise TimeoutError(f'Motor did not start in {maxwaittime} seconds')
+                raise TimeoutError(
+                    f'Motor did not start in {maxwaittime} seconds')
             time.sleep(0.1)
 
     def wait_for_done(self):
@@ -184,6 +185,36 @@ class SinqMotor(Motor):
             return self.read(self.suffixes[fieldname], as_string)
         return self.read('.' + self.fields[fieldname], as_string)
 
+    def wait_enabled(self, timeout=10):
+        """
+        Wait until the motor is enabled or a timeout has been reached.
+        """
+        return self._wait_enabled_disabled(timeout, True)
+
+    def wait_disabled(self, timeout=10):
+        """
+        Wait until the motor is enabled or a timeout has been reached.
+        """
+        return self._wait_enabled_disabled(timeout, False)
+
+    def _wait_enabled_disabled(self, timeout, wait_for_enabling):
+
+        def check(enable_rbv):
+            if wait_for_enabling:
+                return enable_rbv != 0
+            return enable_rbv == 0
+
+        now = time.time()
+        while check(self.read_field('enable_rbv')):
+            if time.time() > now + timeout:
+                if wait_for_enabling:
+                    pytest.fail(
+                        f'Motor {self.pv} could not be enabled in {timeout} seconds')
+                else:
+                    pytest.fail(
+                        f'Motor {self.pv} could not be disabled in {timeout} seconds')
+            time.sleep(0.1)
+
 
 class TurboPMAC(SinqMotor):
     pass
diff --git a/tests/sinqMotor/turboPmac/common.py b/tests/sinqMotor/turboPmac/common.py
index 4695d29..33ffa6b 100644
--- a/tests/sinqMotor/turboPmac/common.py
+++ b/tests/sinqMotor/turboPmac/common.py
@@ -8,26 +8,12 @@ def stop_reset_enable_move_sequence(motor, target):
     motor.write_field('reseterrorpv', 1)
 
     # Wait until the motor is disabled.
-    now = time.time()
-    max_enable_time = 10
-    while motor.read_field('enable_rbv') == 1:
-        # Wait a maximum of max_enable_time seconds until enabling
-        if time.time() > now + max_enable_time:
-            pytest.fail(
-                f'Motor {motor.pv} could not be stopped and resetted in {max_enable_time} seconds')
-        time.sleep(0.1)
+    motor.wait_disabled()
 
     motor.write_field('enable', 1)
 
     # Wait until the motor is enabled.
-    now = time.time()
-    max_enable_time = 10
-    while motor.read_field('enable_rbv') == 0:
-        # Wait a maximum of max_enable_time seconds until enabling
-        if time.time() > now + max_enable_time:
-            pytest.fail(
-                f'Motor {motor.pv} could not be enabled in {max_enable_time} seconds')
-        time.sleep(0.1)
+    motor.wait_enabled()
 
     motor.move_and_wait(target)
     assert motor.at_target(target)
diff --git a/tests/sinqMotor/turboPmac/reset.py b/tests/sinqMotor/turboPmac/reset.py
index 7bfb9fb..2f679d2 100755
--- a/tests/sinqMotor/turboPmac/reset.py
+++ b/tests/sinqMotor/turboPmac/reset.py
@@ -1,4 +1,6 @@
 import time
+import pytest
+
 
 def reset(motor, initpos=0):
     """
@@ -25,7 +27,8 @@ def reset(motor, initpos=0):
         while motor.read_field('enable_rbv') == 0:
             # Wait a maximum of max_enable_time seconds until enabling
             if time.time() > now + max_enable_time:
-                pytest.fail(f'Motor {self.pv} could not be enabled in {max_enable_time} seconds')
+                pytest.fail(
+                    f'Motor {motor.pv} could not be enabled in {max_enable_time} seconds')
             time.sleep(0.1)
 
-    motor.move_and_wait(initpos)
\ No newline at end of file
+    motor.move_and_wait(initpos)