From 713fa1a1b796c3e800980c33ffbf8deaac95446a Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Wed, 10 Dec 2025 14:40:15 +0100
Subject: [PATCH] fix(controller): Add method to remove axis from controller

---
 ophyd_devices/utils/controller.py | 12 ++++++++++++
 tests/test_controller.py          | 26 ++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/ophyd_devices/utils/controller.py b/ophyd_devices/utils/controller.py
index bca08be..7a15c45 100644
--- a/ophyd_devices/utils/controller.py
+++ b/ophyd_devices/utils/controller.py
@@ -248,6 +248,18 @@ class Controller(OphydObject):
         """
         self._axis[axis_nr] = axis
 
+    @axis_checked
+    def remove_axis(self, *, axis_nr: int) -> None:
+        """Remove the device instance assigned to a controller axis.
+
+        Args:
+            axis_nr (int): Controller axis number
+
+        """
+        self._axis[axis_nr] = None
+        if not any(self._axis):
+            self.off(update_config=False)
+
     @axis_checked
     def get_axis(self, axis_nr: int) -> Device:
         """Get device instance for a specified controller axis.
diff --git a/tests/test_controller.py b/tests/test_controller.py
index b93e2d9..ce7f218 100644
--- a/tests/test_controller.py
+++ b/tests/test_controller.py
@@ -74,3 +74,29 @@ def test_controller_with_multiple_axes(dm_with_devices):
         assert controller.connected is False
         controller.set_device_enabled("samx", True)
         assert controller.connected is True
+
+
+def test_remove_axis_from_controller(dm_with_devices):
+    """Test removing an axis from a controller."""
+    socket_cls = mock.MagicMock()
+    Controller._controller_instances = {}
+    Controller._axes_per_controller = 2
+    controller = Controller(
+        socket_cls=socket_cls, socket_host="dummy", socket_port=123, device_manager=dm_with_devices
+    )
+    # Set axes on the controller
+    controller.set_axis(axis=dm_with_devices.devices["samx"], axis_nr=0)
+    controller.set_axis(axis=dm_with_devices.devices["samy"], axis_nr=1)
+    assert len(controller._axis) == 2
+    with mock.patch.object(controller, "off") as mock_off:
+        # Remove one axis
+        controller.remove_axis(axis_nr=0)
+        assert len([a for a in controller._axis if a is not None]) == 1
+        assert "samy" in [a.name for a in controller._axis if a is not None]
+        mock_off.assert_not_called()
+
+        # Remove the last axis, should disable the controller
+        controller.remove_axis(axis_nr=1)
+        assert len([a for a in controller._axis if a is not None]) == 0
+        assert controller.connected is False
+        mock_off.assert_called_once_with(update_config=False)