From ef9568cb462c0faac422c7ebf85991117f734473 Mon Sep 17 00:00:00 2001
From: appel_c <christian.appel@psi.ch>
Date: Tue, 9 Dec 2025 16:32:44 +0100
Subject: [PATCH] fix(controller): cleanup and fix logic for enable/off
 dependency.

---
 ophyd_devices/utils/controller.py |  9 +++++----
 ophyd_devices/utils/socket.py     | 11 +++++++----
 tests/test_socket.py              |  2 +-
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/ophyd_devices/utils/controller.py b/ophyd_devices/utils/controller.py
index a8ffc6b..d08edca 100644
--- a/ophyd_devices/utils/controller.py
+++ b/ophyd_devices/utils/controller.py
@@ -200,7 +200,7 @@ class Controller(OphydObject):
                 if axis is not None
             )
             if all_disabled:
-                self.off()
+                self.off(update_config=False)
 
     def set_all_devices_enable(self, enabled: bool) -> None:
         """
@@ -281,14 +281,15 @@ class Controller(OphydObject):
         else:
             logger.info("The connection has already been established.")
 
-    def off(self) -> None:
+    def off(self, update_config: bool = True) -> None:
         """Close the socket connection to the controller"""
         if self.connected and self.sock is not None:
             self.sock.close()
             self.connected = False
             self.sock = None
-            # Disable all axes associated with this controller
-            self.set_all_devices_enable(False)
+            if update_config:
+                # Disable all axes associated with this controller
+                self.set_all_devices_enable(False)
         else:
             logger.info("The connection is already closed.")
 
diff --git a/ophyd_devices/utils/socket.py b/ophyd_devices/utils/socket.py
index 5764cd6..156729e 100644
--- a/ophyd_devices/utils/socket.py
+++ b/ophyd_devices/utils/socket.py
@@ -187,7 +187,7 @@ class SocketIO:
         Args:
             timeout (int): Time in seconds to wait for connection
         """
-        logger.info(f"Connecting to {self.host}:{self.port}")
+        logger.info(f"Connecting to {self.host}:{self.port}.")
         start_time = time.time()
         while time.time() - start_time < timeout:
             try:
@@ -197,11 +197,14 @@ class SocketIO:
                 break
             except Exception as exc:
                 self.sock = None
-                logger.warning(f"Connection failed, retrying after 0.2 seconds... {exc}")
+                logger.warning(
+                    f"Connection to {self.host}:{self.port} failed after {time.time()-start_time:.2f} seconds"
+                    f" with exception: {exc}. Retrying after 1 second..."
+                )
                 time.sleep(1)
         else:
             raise ConnectionError(
-                f"Could not connect to {self.host}:{self.port} within {time.time()-start_time} seconds"
+                f"Could not connect to {self.host}:{self.port} within {time.time()-start_time:.2f} seconds"
             )
 
     def _put(self, msg_bytes):
@@ -224,7 +227,7 @@ class SocketIO:
         return self._recv(buffer_length=buffer_length)
 
     def open(self, timeout: int = 10):
-        """ "
+        """
         Open the socket connection to the host:port
 
         Args:
diff --git a/tests/test_socket.py b/tests/test_socket.py
index 17f7bcc..be40a91 100644
--- a/tests/test_socket.py
+++ b/tests/test_socket.py
@@ -77,7 +77,7 @@ def test_socket_open_with_timeout():
         mock_connect.side_effect = Exception("Connection failed")
         with pytest.raises(ConnectionError):
             socketio.open(timeout=0.4)
-            mock_connect.assert_called_once()
+        mock_connect.assert_called_once()
 
 
 def test_close():