From db72a81d916a0c8146fe30a4bcc0394b36948224 Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Thu, 7 Aug 2025 12:52:35 +0200
Subject: [PATCH] refactor(undulator): raise for setpoint if undulator is
 operator controlled

---
 ophyd_devices/devices/undulator.py | 44 ++++++++++++++++++++++++++----
 tests/test_undulator.py            | 10 +++++--
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/ophyd_devices/devices/undulator.py b/ophyd_devices/devices/undulator.py
index 1df9a5f..f60eabd 100644
--- a/ophyd_devices/devices/undulator.py
+++ b/ophyd_devices/devices/undulator.py
@@ -24,7 +24,7 @@ class UNDULATORCONTROL(int, enum.Enum):
     BEAMLINE = 1
 
 
-class UndulatorEpicsSignal(EpicsSignal):
+class UndulatorSetointSignal(EpicsSignal):
     """
     SLS Undulator setpoint control
     """
@@ -45,9 +45,43 @@ class UndulatorEpicsSignal(EpicsSignal):
         If the undulator is operator controlled, it will not move.
         """
         if self.parent.select_control.get() == UNDULATORCONTROL.OPERATOR.value:
-            raise PermissionError(
+            raise PermissionError("Undulator is operator controlled!")
+        return super().put(
+            value,
+            force=force,
+            connection_timeout=connection_timeout,
+            callback=callback,
+            use_complete=use_complete,
+            timeout=timeout,
+            **kwargs,
+        )
+
+
+class UndulatorStopSignal(EpicsSignal):
+    """
+    SLS Undulator setpoint control
+    """
+
+    def put(
+        self,
+        value,
+        force=False,
+        connection_timeout=DEFAULT_CONNECTION_TIMEOUT,
+        callback=None,
+        use_complete=None,
+        timeout=DEFAULT_WRITE_TIMEOUT,
+        **kwargs,
+    ):
+        """
+        Put a value to the setpoint PV.
+
+        If the undulator is operator controlled, it will not move.
+        """
+        if self.parent.select_control.get() == UNDULATORCONTROL.OPERATOR.value:
+            logger.error(
                 f"Cannot use put for signal {self.name}; Undulator is operator controlled!"
             )
+            return None
         return super().put(
             value,
             force=force,
@@ -64,10 +98,10 @@ class UndulatorGap(PVPositioner):
     SLS Undulator gap control
     """
 
-    setpoint = Cpt(UndulatorEpicsSignal, suffix="GAP-SP")
-    readback = Cpt(EpicsSignal, suffix="GAP-RBV", kind="hinted", auto_monitor=True)
+    setpoint = Cpt(UndulatorSetointSignal, suffix="GAP-SP")
+    readback = Cpt(EpicsSignalRO, suffix="GAP-RBV", kind="hinted", auto_monitor=True)
 
-    stop_signal = Cpt(UndulatorEpicsSignal, suffix="STOP")
+    stop_signal = Cpt(UndulatorStopSignal, suffix="STOP")
     done = Cpt(EpicsSignalRO, suffix="DONE", auto_monitor=True)
 
     select_control = Cpt(EpicsSignalRO, suffix="SCTRL", auto_monitor=True)
diff --git a/tests/test_undulator.py b/tests/test_undulator.py
index 967ae19..98e481e 100644
--- a/tests/test_undulator.py
+++ b/tests/test_undulator.py
@@ -52,11 +52,17 @@ def test_undulator_raises_when_disabled(mock_undulator):
     with pytest.raises(PermissionError) as e:
         mock_undulator.move(5)
     assert e.match("Undulator is operator controlled!")
+    with pytest.raises(PermissionError) as e:
+        mock_undulator.setpoint.put(5)
+    assert e.match("Undulator is operator controlled!")
 
 
 def test_undulator_stop_call(mock_undulator):
     mock_undulator.select_control._read_pv.mock_data = 1
+    mock_undulator.stop_signal._read_pv.mock_data = 0
     mock_undulator.stop()
+    assert mock_undulator.stop_signal.get() == 1
+    mock_undulator.stop_signal._read_pv.mock_data = 0
     mock_undulator.select_control._read_pv.mock_data = 0
-    with pytest.raises(PermissionError) as e:
-        mock_undulator.stop()
+    mock_undulator.stop()
+    assert mock_undulator.stop_signal.get() == 0