diff --git a/tests/tests_devices/test_galil.py b/tests/tests_devices/test_galil.py index 919e8d2..dc0ed68 100644 --- a/tests/tests_devices/test_galil.py +++ b/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