test(controller): Fix test for controller wait_for_connection

tests/tests_devices/test_galil.py

@@ -5,16 +5,16 @@ import pytest
 from bec_server.device_server.tests.utils import DMMock
 from ophyd_devices.tests.utils import SocketMock
 
-from csaxs_bec.devices.npoint.npoint import NPointAxis
-from csaxs_bec.devices.omny.galil.fgalil_ophyd import FlomniGalilMotor
-from csaxs_bec.devices.omny.galil.fupr_ophyd import FuprGalilMotor
+from csaxs_bec.devices.npoint.npoint import NPointAxis, NPointController
+from csaxs_bec.devices.omny.galil.fgalil_ophyd import FlomniGalilController, FlomniGalilMotor
+from csaxs_bec.devices.omny.galil.fupr_ophyd import FuprGalilController, FuprGalilMotor
 from csaxs_bec.devices.omny.galil.lgalil_ophyd import LamniGalilController, LamniGalilMotor
-from csaxs_bec.devices.omny.galil.ogalil_ophyd import OMNYGalilMotor
-from csaxs_bec.devices.omny.galil.sgalil_ophyd import SGalilMotor
-from csaxs_bec.devices.omny.rt.rt_flomni_ophyd import RtFlomniMotor
-from csaxs_bec.devices.omny.rt.rt_lamni_ophyd import RtLamniMotor
-from csaxs_bec.devices.omny.rt.rt_omny_ophyd import RtOMNYMotor
-from csaxs_bec.devices.smaract.smaract_ophyd import SmaractMotor
+from csaxs_bec.devices.omny.galil.ogalil_ophyd import OMNYGalilController, OMNYGalilMotor
+from csaxs_bec.devices.omny.galil.sgalil_ophyd import GalilController, SGalilMotor
+from csaxs_bec.devices.omny.rt.rt_flomni_ophyd import RtFlomniController, RtFlomniMotor
+from csaxs_bec.devices.omny.rt.rt_lamni_ophyd import RtLamniController, RtLamniMotor
+from csaxs_bec.devices.omny.rt.rt_omny_ophyd import RtOMNYController, RtOMNYMotor
+from csaxs_bec.devices.smaract.smaract_ophyd import SmaractController, SmaractMotor
 
 
 @pytest.fixture(scope="function")
@@ -176,23 +176,25 @@ def test_find_reference(leyex, axis_nr, socket_put_messages, socket_get_messages
 def test_wait_for_connection_called():
     """Test that wait_for_connection is called on all motors that have a socket controller."""
     dm = DMMock()
-    with (
-        mock.patch("ophyd_devices.utils.controller.Controller.on") as mock_on,
-        mock.patch("ophyd_devices.utils.controller.Controller.set_axis") as mock_set_axis,
-    ):
-        motors = [
-            FlomniGalilMotor,
-            FuprGalilMotor,
-            LamniGalilMotor,
-            OMNYGalilMotor,
-            SGalilMotor,
-            RtFlomniMotor,
-            RtLamniMotor,
-            RtOMNYMotor,
-            SmaractMotor,
-            NPointAxis,
-        ]
-        for motor_cls in motors:
+    testable_connections = [
+        (NPointAxis, NPointController),
+        (FlomniGalilMotor, FlomniGalilController),
+        (FuprGalilMotor, FuprGalilController),
+        (LamniGalilMotor, LamniGalilController),
+        (OMNYGalilMotor, OMNYGalilController),
+        (SGalilMotor, GalilController),
+        (RtFlomniMotor, RtFlomniController),
+        (RtLamniMotor, RtLamniController),
+        (RtOMNYMotor, RtOMNYController),
+        (SmaractMotor, SmaractController),
+    ]
+    for motor_cls, controller_cls in testable_connections:
+        # Store values to restore later
+        ctrl_axis_backup = controller_cls._axes_per_controller
+        try:
+            controller_cls._reset_controller()
+            controller_cls._axes_per_controller = 3
+
             motor = motor_cls(
                 "C",
                 name="test_motor",
@@ -201,5 +203,10 @@ def test_wait_for_connection_called():
                 socket_cls=SocketMock,
                 device_manager=dm,
             )
-            motor.wait_for_connection(timeout=5.0)
-            assert mock_on.call_args_list[-1] == mock.call(timeout=5.0)
+            with mock.patch.object(motor.controller, "on") as mock_on:
+
+                motor.wait_for_connection(timeout=5.0)
+                assert mock_on.call_args_list[-1] == mock.call(timeout=5.0)
+        finally:
+            controller_cls._reset_controller()
+            controller_cls._axes_per_controller = ctrl_axis_backup