Enhance MockSuperConductingMagnetController to be configurable

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",
 ]

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

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]