Port UndulatorDCM to ophyd-async and remove energy signal

src/dodal/devices/undulator_dcm.py

@@ -1,14 +1,16 @@
+import asyncio
+
 import numpy as np
+from bluesky.protocols import Movable
 from numpy import argmin, loadtxt, ndarray
-from ophyd.status import Status
-from ophyd_async.core import SignalW, StandardReadable
+from ophyd_async.core import AsyncStatus, StandardReadable
 
 from dodal.devices.dcm import DCM
 from dodal.devices.undulator import Undulator, UndulatorGapAccess
 from dodal.log import LOGGER
 
-ENERGY_TIMEOUT_S = 30
-STATUS_TIMEOUT_S = 10
+ENERGY_TIMEOUT_S: float = 30.0
+STATUS_TIMEOUT_S: float = 10.0
 
 # Enable to allow testing when the beamline is down, do not change in production!
 TEST_MODE = False
@@ -19,6 +21,7 @@ class AccessError(Exception):
 
 
 def _get_energy_distance_table(lookup_table_path: str) -> ndarray:
+    # TODO: Make IO async
     return loadtxt(lookup_table_path, comments=["#", "Units"])
 
 
@@ -30,60 +33,66 @@ def _get_closest_gap_for_energy(
     return table[1][idx]
 
 
-class UndulatorDCM(StandardReadable):
+class UndulatorDCM(StandardReadable, Movable):
     """
     Composite device to handle changing beamline energies
     """
 
-    class EnergySignal(Signal):
-        parent: "UndulatorDCM"
+    _setpoint: float
+
+    def __init__(self, undulator: Undulator, dcm: DCM, name: str = ""):
+        self.undulator = undulator
+        self.dcm = dcm
+
+        super().__init__(name)
 
-        def set(self, value, *, timeout=None, settle_time=None, **kwargs) -> Status:
+    def set(self, value: float) -> AsyncStatus:
+        async def _set(value: float):
+            # TODO: Break up into smaller methods
             energy_kev = value
-            access_level = self.parent.undulator.gap_access.get(as_string=True)
-            if access_level == UndulatorGapAccess.DISABLED.value and not TEST_MODE:
+            access_level = await self.undulator.gap_access.get_value()
+            if access_level is UndulatorGapAccess.DISABLED and not TEST_MODE:
                 raise AccessError(
                     "Undulator gap access is disabled. Contact Control Room"
                 )
 
             # Get 2d np.array converting energies to undulator gap distance, from lookup table
             energy_to_distance_table = _get_energy_distance_table(
-                self.parent.undulator.lookup_table_path
+                self.undulator.lookup_table_path
             )
             LOGGER.info(f"Setting DCM energy to {energy_kev:.2f} kev")
 
-            status = self.parent.dcm.energy_in_kev.move(
-                energy_kev, timeout=ENERGY_TIMEOUT_S
-            )
+            statuses = [
+                self.dcm.energy_in_kev.set(
+                    energy_kev,
+                    timeout=ENERGY_TIMEOUT_S,
+                )
+            ]
 
             # Use the lookup table to get the undulator gap associated with this dcm energy
             gap_to_match_dcm_energy = _get_closest_gap_for_energy(
-                energy_kev * 1000, energy_to_distance_table
+                energy_kev * 1000,
+                energy_to_distance_table,
             )
 
             # Check if undulator gap is close enough to the value from the DCM
-            current_gap = self.parent.undulator.current_gap.get()
+            current_gap = await self.undulator.current_gap.get_value()
 
             if (
                 abs(gap_to_match_dcm_energy - current_gap)
-                > self.parent.undulator.gap_discrepancy_tolerance_mm
+                > self.undulator.gap_discrepancy_tolerance_mm
             ):
                 LOGGER.info(
                     f"Undulator gap mismatch. {abs(gap_to_match_dcm_energy-current_gap):.3f}mm is outside tolerance.\
                     Moving gap to nominal value, {gap_to_match_dcm_energy:.3f}mm"
                 )
                 if not TEST_MODE:
-                    status &= self.parent.undulator.gap_motor.move(
-                        gap_to_match_dcm_energy, timeout=STATUS_TIMEOUT_S
+                    statuses.append(
+                        self.undulator.gap_motor.set(
+                            gap_to_match_dcm_energy,
+                            timeout=STATUS_TIMEOUT_S,
+                        )
                     )
+            await asyncio.gather(*statuses)
 
-            return status
-
-    energy_kev = Component(EnergySignal)
-
-    def __init__(self, undulator: Undulator, dcm: DCM, *args, **kwargs):
-        self.undulator = undulator
-        self.dcm = dcm
-        self.energy_kev = SignalW()
-
-        super().__init__(*args, **kwargs)
+        return AsyncStatus(_set(value))

tests/devices/unit_tests/test_focusing_mirror.py

@@ -63,9 +63,7 @@ def not_ok(_):
 def test_mirror_set_voltage_sets_and_waits_happy_path(
     vfm_mirror_voltages_with_set: VFMMirrorVoltages,
 ):
-    vfm_mirror_voltages_with_set._channel14_voltage_device._setpoint_v.set.return_value = (
-        NullStatus()
-    )
+    vfm_mirror_voltages_with_set._channel14_voltage_device._setpoint_v.set.return_value = NullStatus()
     vfm_mirror_voltages_with_set._channel14_voltage_device._demand_accepted.sim_put(
         MirrorVoltageDemand.OK
     )
@@ -103,9 +101,7 @@ def test_mirror_set_voltage_sets_and_waits_set_fail(
 def test_mirror_set_voltage_sets_and_waits_settle_timeout_expires(
     vfm_mirror_voltages_with_set_timing_out: VFMMirrorVoltages,
 ):
-    vfm_mirror_voltages_with_set_timing_out._channel14_voltage_device._setpoint_v.set.return_value = (
-        NullStatus()
-    )
+    vfm_mirror_voltages_with_set_timing_out._channel14_voltage_device._setpoint_v.set.return_value = NullStatus()
 
     status: StatusBase = vfm_mirror_voltages_with_set_timing_out.voltage_channels[
         0

tests/devices/unit_tests/test_undulator_dcm.py

@@ -1,9 +1,13 @@
+import asyncio
 from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
-from ophyd.sim import make_fake_device
-from ophyd.status import Status
+from ophyd_async.core import (
+    AsyncStatus,
+    DeviceCollector,
+    set_sim_value,
+)
 
 from dodal.devices.dcm import DCM
 from dodal.devices.undulator import Undulator, UndulatorGapAccess
@@ -18,26 +22,24 @@
 
 
 @pytest.fixture
-def fake_undulator_dcm() -> UndulatorDCM:
-    undulator: Undulator = make_fake_device(Undulator)(
-        name="undulator",
-        lookup_table_path="./tests/devices/unit_tests/test_beamline_undulator_to_gap_lookup_table.txt",
-    )
-    dcm: DCM = make_fake_device(DCM)(
-        name="dcm", daq_configuration_path=MOCK_DAQ_CONFIG_PATH
-    )
-    undulator_dcm: UndulatorDCM = make_fake_device(UndulatorDCM)(
-        undulator, dcm, name="undulator_dcm"
-    )
+async def fake_undulator_dcm() -> UndulatorDCM:
+    async with DeviceCollector(sim=True):
+        undulator = Undulator(
+            "UND-01",
+            name="undulator",
+            lookup_table_path="./tests/devices/unit_tests/test_beamline_undulator_to_gap_lookup_table.txt",
+        )
+        dcm = DCM("DCM-01", name="dcm", daq_configuration_path=MOCK_DAQ_CONFIG_PATH)
+        undulator_dcm = UndulatorDCM(undulator, dcm, name="undulator_dcm")
     return undulator_dcm
 
 
-def test_when_gap_access_is_disabled_set_energy_then_error_is_raised(
+async def test_when_gap_access_is_disabled_set_energy_then_error_is_raised(
     fake_undulator_dcm: UndulatorDCM,
 ):
-    fake_undulator_dcm.undulator.gap_access.sim_put(UndulatorGapAccess.DISABLED.value)  # type: ignore
+    set_sim_value(fake_undulator_dcm.undulator.gap_access, UndulatorGapAccess.DISABLED)
     with pytest.raises(AccessError):
-        fake_undulator_dcm.energy_kev.set(5)
+        await fake_undulator_dcm.set(5)
 
 
 def test_energy_to_distance_table_correct_format(fake_undulator_dcm: UndulatorDCM):
@@ -60,80 +62,88 @@ def test_correct_closest_distance_to_energy_from_table(dcm_energy, expected_outp
 
 @patch("dodal.devices.undulator_dcm.loadtxt")
 @patch("dodal.devices.undulator_dcm.LOGGER")
-def test_if_gap_is_wrong_then_logger_info_is_called_and_gap_is_set_correctly(
+async def test_if_gap_is_wrong_then_logger_info_is_called_and_gap_is_set_correctly(
     mock_logger: MagicMock, mock_load: MagicMock, fake_undulator_dcm: UndulatorDCM
 ):
-    fake_undulator_dcm.undulator.current_gap.sim_put(5.3)  # type: ignore
-    fake_undulator_dcm.undulator.gap_motor.move = MagicMock()
-    fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock()
+    set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.3)
+    set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.readback, 5.7)
+
+    # fake_undulator_dcm.undulator.gap_motor.move = MagicMock()
+    # fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock()
     mock_load.return_value = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]])
-    fake_undulator_dcm.dcm.energy_in_kev.user_readback.sim_put(5.7)  # type: ignore
 
-    fake_undulator_dcm.energy_kev.set(6.9)
+    await fake_undulator_dcm.set(6.9)
 
-    fake_undulator_dcm.dcm.energy_in_kev.move.assert_called_once_with(6.9, timeout=30)
-    fake_undulator_dcm.undulator.gap_motor.move.assert_called_once_with(
-        6.045, timeout=10
-    )
+    assert (await fake_undulator_dcm.dcm.energy_in_kev.setpoint.get_value()) == 6.9
+    assert (await fake_undulator_dcm.undulator.gap_motor.setpoint.get_value()) == 6.045
     mock_logger.info.assert_called()
 
 
 @patch("dodal.devices.undulator_dcm.loadtxt")
 @patch("dodal.devices.undulator_dcm.LOGGER")
 @patch("dodal.devices.undulator_dcm.TEST_MODE", True)
-def test_when_gap_access_is_not_checked_if_test_mode_enabled(
+async def test_when_gap_access_is_not_checked_if_test_mode_enabled(
     mock_logger: MagicMock, mock_load: MagicMock, fake_undulator_dcm: UndulatorDCM
 ):
-    fake_undulator_dcm.undulator.gap_access.sim_put(UndulatorGapAccess.DISABLED.value)  # type: ignore
-    fake_undulator_dcm.undulator.current_gap.sim_put(5.3)  # type: ignore
-    fake_undulator_dcm.undulator.gap_motor.move = MagicMock()
-    fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock()
+    set_sim_value(fake_undulator_dcm.undulator.gap_access, UndulatorGapAccess.DISABLED)
+    set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.3)
+    set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.readback, 5.7)
+
+    set_sim_value(fake_undulator_dcm.undulator.gap_motor.setpoint, 0.0)
+    set_sim_value(fake_undulator_dcm.undulator.gap_motor.readback, 0.0)
+
+    # fake_undulator_dcm.undulator.gap_motor.move = MagicMock()
+    # fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock()
     mock_load.return_value = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]])
-    fake_undulator_dcm.dcm.energy_in_kev.user_readback.sim_put(5.7)  # type: ignore
 
-    fake_undulator_dcm.energy_kev.set(6.9)
+    await fake_undulator_dcm.set(6.9)
+
+    assert (await fake_undulator_dcm.dcm.energy_in_kev.setpoint.get_value()) == 6.9
+    # Verify undulator has not been asked to move
+    assert (await fake_undulator_dcm.undulator.gap_motor.setpoint.get_value()) == 0.0
 
-    fake_undulator_dcm.dcm.energy_in_kev.move.assert_called_once_with(6.9, timeout=30)
-    fake_undulator_dcm.undulator.gap_motor.move.assert_not_called()
     mock_logger.info.assert_called()
 
 
 @patch("dodal.devices.undulator_dcm.loadtxt")
 @patch("dodal.devices.undulator_dcm.LOGGER")
-def test_if_gap_is_already_correct_then_dont_move_gap(
+async def test_if_gap_is_already_correct_then_dont_move_gap(
     mock_logger: MagicMock, mock_load: MagicMock, fake_undulator_dcm: UndulatorDCM
 ):
-    fake_undulator_dcm.undulator.gap_motor.move = MagicMock()
-    fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock()
+    set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.setpoint, 0.0)
+    set_sim_value(fake_undulator_dcm.dcm.energy_in_kev.readback, 0.0)
+    # fake_undulator_dcm.undulator.gap_motor.move = MagicMock()
+    # fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock()
     mock_load.return_value = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]])
-    fake_undulator_dcm.undulator.current_gap.sim_put(5.4605)  # type: ignore
+    set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.4605)
 
-    fake_undulator_dcm.energy_kev.set(5.8).wait(timeout=0.01)
+    status = fake_undulator_dcm.set(5.8)
+    await asyncio.wait_for(status, timeout=0.01)
 
-    fake_undulator_dcm.undulator.gap_motor.move.assert_not_called()
+    # Verify undulator has not been asked to move
+    assert (await fake_undulator_dcm.undulator.gap_motor.setpoint.get_value()) == 0.0
     mock_logger.info.assert_called_once()
 
 
-def test_energy_set_only_complete_when_all_statuses_are_finished(
+async def test_energy_set_only_complete_when_all_statuses_are_finished(
     fake_undulator_dcm: UndulatorDCM,
 ):
-    dcm_energy_move_status = Status()
-    undulator_gap_move_status = Status()
+    set_sim_value(fake_undulator_dcm.undulator.current_gap, 5.0)
+
+    release_dcm = asyncio.Event()
+    release_undulator = asyncio.Event()
 
-    fake_undulator_dcm.dcm.energy_in_kev.move = MagicMock(
-        return_value=dcm_energy_move_status
+    fake_undulator_dcm.dcm.energy_in_kev.set = MagicMock(
+        return_value=AsyncStatus(release_dcm.wait())
     )
-    fake_undulator_dcm.undulator.gap_motor.move = MagicMock(
-        return_value=undulator_gap_move_status
+    fake_undulator_dcm.undulator.gap_motor.set = MagicMock(
+        return_value=AsyncStatus(release_undulator.wait())
     )
-    _get_energy_distance_table = MagicMock()
-    _get_closest_gap_for_energy = MagicMock(return_value=10)
-    fake_undulator_dcm.undulator.current_gap.sim_put(5)  # type: ignore
 
-    status: Status = fake_undulator_dcm.energy_kev.set(5.8)
+    status = fake_undulator_dcm.set(5.0)
 
-    assert not status.success
-    dcm_energy_move_status.set_finished()
-    assert not status.success
-    undulator_gap_move_status.set_finished()
-    status.wait(timeout=0.01)
+    assert not status.done
+    release_dcm.set()
+    assert not status.done
+    release_undulator.set()
+    await asyncio.wait_for(status, timeout=0.01)