diff --git a/src/dodal/devices/beamlines/i06_1/magnet/__init__.py b/src/dodal/devices/beamlines/i06_1/magnet/__init__.py index f81cc14696b..34a91107dc7 100644 --- a/src/dodal/devices/beamlines/i06_1/magnet/__init__.py +++ b/src/dodal/devices/beamlines/i06_1/magnet/__init__.py @@ -13,6 +13,7 @@ MagnetLimitStatus, MagnetMode, MagnetRampStatus, + MockSuperConductingMagnetController, SuperConductingMagnetController, ) @@ -30,5 +31,6 @@ "MagnetLimitStatus", "MagnetMode", "MagnetRampStatus", + "MockSuperConductingMagnetController", "SuperConductingMagnetController", ] diff --git a/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py b/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py index da154c500df..c46538fd311 100644 --- a/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py +++ b/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py @@ -243,8 +243,27 @@ async def _set_phi(self, phi: float): class MockSuperConductingMagnetController( DeviceMock["SuperConductingMagnetController"] ): - """Add additional callback logic to our device to get the mock behaviour to simulate - the hardware as best we can. + """Mock controller that simulates the behaviour of the + SuperConductingMagnetController hardware. + + The mock reproduces additional IOC behaviour that is not provided by the + standard device mock, including: + + - Updating readback positions when a ramp is triggered. + - Resetting demand positions back to zero, configuring PSU limits, and triggering a + ramp when the mode changes. + - Simulating the movement of readback positions over time. + + Movements are simulated over multiple steps by default so that beamline + operation in mock mode behaves similarly to the real hardware. This also + allows fly scans to be exercised in mock mode, with detector events + occurring while the magnet is moving. + + Unit tests that do not require simulated movement can disable it by + setting ``steps`` to zero:: + + scmc = SuperConductingMagnetController(..., name="scmc") + await scmc.connect(mock=MockSuperConductingMagnetController(steps=0)) """ # Pulled directly from live IOC so can replicate behaviour in mock mode. @@ -258,20 +277,50 @@ class MockSuperConductingMagnetController( MagnetMode.SPHERICAL: (2, 2, 2), } - async def connect(self, device: "SuperConductingMagnetController"): + def __init__( + self, + name: str = "", + parent: DeviceMock | None = None, + steps: int = 10, + ramp_time: float = 1.0, + ): + super().__init__(name, parent) + self.steps = steps + self.ramp_time = ramp_time + async def connect(self, device: "SuperConductingMagnetController"): async def _trigger_start_ramp(): - # Whenever ramp is triggered for the ioc, readback values move to the + # Whenever ramp is triggered for the IOC, readback values move to the # demand values. Simulate this behaviour here. - x_d, y_d, z_d = await asyncio.gather( + x_d, y_d, z_d, x_r, y_r, z_r = await asyncio.gather( device.cart.x.demand.get_value(), device.cart.y.demand.get_value(), device.cart.z.demand.get_value(), + device.cart.x.readback.get_value(), + device.cart.y.readback.get_value(), + device.cart.z.readback.get_value(), + ) + axes = ( + (device.cart.x.readback, x_r, x_d), + (device.cart.y.readback, y_r, y_d), + (device.cart.z.readback, z_r, z_d), ) + # Only move the axis that has changed + axes_to_move = [ + (rb, rb_val, demand) for rb, rb_val, demand in axes if rb_val != demand + ] + # Use configured number of steps or use a single step, whichever is larger + steps = max(self.steps, 1) + step_time = self.ramp_time / steps if steps > 1 else 0 + set_mock_value(device.ramp_status, MagnetRampStatus.RAMPING) - set_mock_value(device.cart.x.readback, x_d) - set_mock_value(device.cart.y.readback, y_d) - set_mock_value(device.cart.z.readback, z_d) + for step in range(1, steps + 1): + fraction = step / steps + for readback, readback_value, demand in axes_to_move: + set_mock_value( + readback, readback_value + (demand - readback_value) * fraction + ) + await asyncio.sleep(step_time) set_mock_value(device.ramp_status, MagnetRampStatus.RAMP_MADE) callback_on_mock_execute(device._start_ramp, _trigger_start_ramp) # noqa: SLF001 diff --git a/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py b/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py index 6310d585918..118dea79495 100644 --- a/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py +++ b/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py @@ -6,7 +6,7 @@ from bluesky import FailedStatus, RunEngine from bluesky.plan_stubs import mv from bluesky.protocols import Reading -from ophyd_async.core import DEFAULT_TIMEOUT, init_devices, set_mock_value +from ophyd_async.core import DEFAULT_TIMEOUT, SignalR, init_devices, set_mock_value from ophyd_async.testing import assert_configuration, assert_reading, partial_reading from dodal.devices.beamlines.i06_1.magnet import ( @@ -20,6 +20,7 @@ MagnetRampStatus, MagnetRequest, MagnetSphericalPosition, + MockSuperConductingMagnetController, SuperConductingMagnetController, ThreeMagnetAxisPowerSupply, movement, @@ -33,16 +34,17 @@ @pytest.fixture def scmc_psu() -> ThreeMagnetAxisPowerSupply: with init_devices(mock=True): - ramp_rate = ThreeMagnetAxisPowerSupply("TEST:") - return ramp_rate + scmc_psu = ThreeMagnetAxisPowerSupply("TEST:") + return scmc_psu @pytest.fixture -def scmc( +async def scmc( scmc_psu: ThreeMagnetAxisPowerSupply, ) -> SuperConductingMagnetController: - with init_devices(mock=True): - scmc = SuperConductingMagnetController("TEST:", scmc_psu) + # Optimise tests by making movement of readback to setpoint instant. + scmc = SuperConductingMagnetController("TEST:", scmc_psu, name="scmc") + await scmc.connect(mock=MockSuperConductingMagnetController(steps=0)) return scmc @@ -351,27 +353,30 @@ async def test_scmc_executes_movement_strategy_and_ramp_at_each_step( call(movement.MagnetRequest(x=0.5), timeout=DEFAULT_TIMEOUT), call(movement.MagnetRequest(x=1.2), timeout=DEFAULT_TIMEOUT), ] - scmc._MODE_MOVEMENT_STRATEGY[MagnetMode.UNIAXIAL_X] = movement_strategy - scmc._trigger_ramp = AsyncMock() - - with patch.object( - scmc, - "_apply_step", - wraps=scmc._apply_step, - ) as mock_apply_step: - # Configures PSU limits to X=2, Y=0, Z=0 - await scmc.mode.set(MagnetMode.UNIAXIAL_X) - - # Target is within the X axis limit - await scmc.cart.x.set(1.2) - - movement_strategy.move_steps.assert_called_once_with( - ANY, - movement.MagnetRequest(x=1.2), - ) + with patch.dict( + scmc._MODE_MOVEMENT_STRATEGY, + {MagnetMode.UNIAXIAL_X: movement_strategy}, + ): + scmc._trigger_ramp = AsyncMock() + + with patch.object( + scmc, + "_apply_step", + wraps=scmc._apply_step, + ) as mock_apply_step: + # Configures PSU limits to X=2, Y=0, Z=0 + await scmc.mode.set(MagnetMode.UNIAXIAL_X) + + # Target is within the X axis limit + await scmc.cart.x.set(1.2) + + movement_strategy.move_steps.assert_called_once_with( + ANY, + movement.MagnetRequest(x=1.2), + ) - assert mock_apply_step.call_args_list == expected_apply_step_calls - assert scmc._trigger_ramp.call_count == len(move_steps) + assert mock_apply_step.call_args_list == expected_apply_step_calls + assert scmc._trigger_ramp.call_count == len(move_steps) async def test_external_parallel_moves_for_scmc_raise_error( @@ -541,3 +546,61 @@ async def test_scmc_set_within_boundary_timeout_set_correctly( target_request, timeout=expected_timeout, ) + + +@pytest.mark.parametrize( + "steps, ramp_time", + [ + pytest.param(0, 0.0, id="instant"), + pytest.param(4, 0.04, id="stepped"), + ], +) +@pytest.mark.parametrize( + "axis, mode, value", + [ + pytest.param("x", MagnetMode.UNIAXIAL_X, 1.0, id="x"), + pytest.param("y", MagnetMode.UNIAXIAL_Y, 1.0, id="y"), + pytest.param("z", MagnetMode.UNIAXIAL_Z, 1.0, id="z"), + ], +) +async def test_mock_scmc_only_ramps_target_axis( + scmc_psu: ThreeMagnetAxisPowerSupply, + steps: int, + ramp_time: float, + axis: str, + mode: MagnetMode, + value: float, +): + scmc = SuperConductingMagnetController("PV:", scmc_psu, name="scmc") + await scmc.connect( + mock=MockSuperConductingMagnetController( + steps=steps, + ramp_time=ramp_time, + ) + ) + await scmc.mode.set(mode) + + readbacks: dict[str, SignalR[float]] = { + axis: getattr(scmc.cart, axis).readback for axis in ("x", "y", "z") + } + values: dict[str, list[float]] = {axis: [] for axis in readbacks} + + for axis_name, readback in readbacks.items(): + readback_name = readback.name + readback.subscribe( + lambda value, axis_name=axis_name, readback_name=readback_name: values[ + axis_name + ].append(value[readback_name]["value"]) + ) + await getattr(scmc.cart, axis).set(value) + # The initial 0.0 is emitted when the readback subscription is created, + # followed by each value produced during the ramp. + assert values[axis] == [ + 0.0, + *(value * step / max(steps, 1) for step in range(1, max(steps, 1) + 1)), + ] + # Non-target axes should only emit their initial readback value and should not + # be updated by the ramp as value not changed. + for other_axis in readbacks: + if other_axis != axis: + assert values[other_axis] == [0.0]