Convert mirror voltage devices to use ophyd async (#636)

* (#604) Convert focusing_mirror mirror voltages to ophyd_async * (#604) Refactor VFMMirrorVoltages to use DeviceVector * (#604) response to PR comments --------- Co-authored-by: Dominic Oram <[email protected]>
DiamondLightSource · Jul 11, 2024 · 555f696 · 555f696
1 parent 36e289a
commit 555f696
Show file tree

Hide file tree

Showing 3 changed files with 255 additions and 135 deletions.
diff --git a/src/dodal/devices/focusing_mirror.py b/src/dodal/devices/focusing_mirror.py
@@ -1,12 +1,16 @@
-from enum import Enum, IntEnum
-from typing import Any
-
-from ophyd import Component, Device, EpicsSignal
-from ophyd.status import Status, StatusBase
-from ophyd_async.core import StandardReadable
+from enum import Enum
+
+from ophyd_async.core import (
+    AsyncStatus,
+    Device,
+    DeviceVector,
+    StandardReadable,
+    observe_value,
+)
 from ophyd_async.core.signal import soft_signal_r_and_setter
 from ophyd_async.epics.motion import Motor
 from ophyd_async.epics.signal import (
+    epics_signal_r,
     epics_signal_rw,
     epics_signal_x,
 )
@@ -32,93 +36,89 @@ class MirrorStripe(str, Enum):
     PLATINUM = "Platinum"
 
 
-class MirrorVoltageDemand(IntEnum):
-    N_A = 0
-    OK = 1
-    FAIL = 2
-    SLEW = 3
+class MirrorVoltageDemand(str, Enum):
+    N_A = "N/A"
+    OK = "OK"
+    FAIL = "FAIL"
+    SLEW = "SLEW"
 
 
 class MirrorVoltageDevice(Device):
     """Abstract the bimorph mirror voltage PVs into a single device that can be set asynchronously and returns when
     the demanded voltage setpoint is accepted, without blocking the caller as this process can take significant time.
     """
 
-    _actual_v: EpicsSignal = Component(EpicsSignal, "R")
-    _setpoint_v: EpicsSignal = Component(EpicsSignal, "D")
-    _demand_accepted: EpicsSignal = Component(EpicsSignal, "DSEV")
+    def __init__(self, name: str = "", prefix: str = ""):
+        self._actual_v = epics_signal_r(int, prefix + "R")
+        self._setpoint_v = epics_signal_rw(int, prefix + "D")
+        self._demand_accepted = epics_signal_r(MirrorVoltageDemand, prefix + "DSEV")
+        super().__init__(name=name)
 
-    def set(self, value, *args, **kwargs) -> StatusBase:
+    @AsyncStatus.wrap
+    async def set(self, value, *args, **kwargs):
         """Combine the following operations into a single set:
         1. apply the value to the setpoint PV
-        2. Return to the caller with a Status future
         3. Wait until demand is accepted
         4. When either demand is accepted or DEFAULT_SETTLE_TIME expires, signal the result on the Status
         """
 
         setpoint_v = self._setpoint_v
         demand_accepted = self._demand_accepted
 
-        if demand_accepted.get() != MirrorVoltageDemand.OK:
+        if await demand_accepted.get_value() != MirrorVoltageDemand.OK:
             raise AssertionError(
                 f"Attempted to set {setpoint_v.name} when demand is not accepted."
             )
 
-        if setpoint_v.get() == value:
+        if await setpoint_v.get_value() == value:
             LOGGER.debug(f"{setpoint_v.name} already at {value} - skipping set")
-            return Status(success=True, done=True)
+            return
 
         LOGGER.debug(f"setting {setpoint_v.name} to {value}")
-        demand_accepted_status = Status(self, DEFAULT_SETTLE_TIME_S)
-
-        subscription: dict[str, Any] = {"handle": None}
 
-        def demand_check_callback(old_value, value, **kwargs):
-            LOGGER.debug(f"Got event old={old_value} new={value} for {setpoint_v.name}")
-            if old_value != MirrorVoltageDemand.OK and value == MirrorVoltageDemand.OK:
-                LOGGER.debug(f"Demand accepted for {setpoint_v.name}")
-                subs_handle = subscription.pop("handle", None)
-                if subs_handle is None:
-                    raise AssertionError("Demand accepted before set attempted")
-                demand_accepted.unsubscribe(subs_handle)
-
-                demand_accepted_status.set_finished()
-            # else timeout handled by parent demand_accepted_status
+        # Register an observer up front to ensure we don't miss events after we
+        # perform the set
+        demand_accepted_iterator = observe_value(
+            demand_accepted, timeout=DEFAULT_SETTLE_TIME_S
+        )
+        # discard the current value (OK) so we can await a subsequent change
+        await anext(demand_accepted_iterator)
+        await setpoint_v.set(value)
+
+        # The set should always change to SLEW regardless of whether we are
+        # already at the set point, then change back to OK/FAIL depending on
+        # success
+        accepted_value = await anext(demand_accepted_iterator)
+        assert accepted_value == MirrorVoltageDemand.SLEW
+        LOGGER.debug(
+            f"Demand not accepted for {setpoint_v.name}, waiting for acceptance..."
+        )
+        while MirrorVoltageDemand.SLEW == (
+            accepted_value := await anext(demand_accepted_iterator)
+        ):
+            pass
 
-        subscription["handle"] = demand_accepted.subscribe(demand_check_callback)
-        setpoint_status = setpoint_v.set(value)
-        status = setpoint_status & demand_accepted_status
-        return status
+        if accepted_value != MirrorVoltageDemand.OK:
+            raise AssertionError(
+                f"Voltage slew failed for {setpoint_v.name}, new state={accepted_value}"
+            )
 
 
-class VFMMirrorVoltages(Device):
-    def __init__(self, *args, daq_configuration_path: str, **kwargs):
-        super().__init__(*args, **kwargs)
+class VFMMirrorVoltages(StandardReadable):
+    def __init__(
+        self, name: str, prefix: str, *args, daq_configuration_path: str, **kwargs
+    ):
         self.voltage_lookup_table_path = (
             daq_configuration_path + "/json/mirrorFocus.json"
         )
-
-    _channel14_voltage_device = Component(MirrorVoltageDevice, "BM:V14")
-    _channel15_voltage_device = Component(MirrorVoltageDevice, "BM:V15")
-    _channel16_voltage_device = Component(MirrorVoltageDevice, "BM:V16")
-    _channel17_voltage_device = Component(MirrorVoltageDevice, "BM:V17")
-    _channel18_voltage_device = Component(MirrorVoltageDevice, "BM:V18")
-    _channel19_voltage_device = Component(MirrorVoltageDevice, "BM:V19")
-    _channel20_voltage_device = Component(MirrorVoltageDevice, "BM:V20")
-    _channel21_voltage_device = Component(MirrorVoltageDevice, "BM:V21")
-
-    @property
-    def voltage_channels(self) -> list[MirrorVoltageDevice]:
-        return [
-            self._channel14_voltage_device,
-            self._channel15_voltage_device,
-            self._channel16_voltage_device,
-            self._channel17_voltage_device,
-            self._channel18_voltage_device,
-            self._channel19_voltage_device,
-            self._channel20_voltage_device,
-            self._channel21_voltage_device,
-        ]
+        with self.add_children_as_readables():
+            self.voltage_channels = DeviceVector(
+                {
+                    i - 14: MirrorVoltageDevice(prefix=f"{prefix}BM:V{i}")
+                    for i in range(14, 22)
+                }
+            )
+        super().__init__(*args, name=name, **kwargs)
 
 
 class FocusingMirror(StandardReadable):

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -6,12 +6,11 @@
 import time
 from os import environ, getenv
 from pathlib import Path
-from typing import Mapping, cast
+from typing import Mapping
 from unittest.mock import MagicMock, patch
 
 import pytest
 from bluesky.run_engine import RunEngine
-from ophyd.sim import make_fake_device
 
 from dodal.beamlines import i03
 from dodal.common.beamlines import beamline_utils
@@ -73,17 +72,11 @@ def pytest_runtest_teardown():
 
 
 @pytest.fixture
-def vfm_mirror_voltages() -> VFMMirrorVoltages:
-    voltages = cast(
-        VFMMirrorVoltages,
-        make_fake_device(VFMMirrorVoltages)(
-            name="vfm_mirror_voltages",
-            prefix="BL-I03-MO-PSU-01:",
-            daq_configuration_path=i03.DAQ_CONFIGURATION_PATH,
-        ),
-    )
+def vfm_mirror_voltages(RE: RunEngine) -> VFMMirrorVoltages:
+    voltages = i03.vfm_mirror_voltages(fake_with_ophyd_sim=True)
     voltages.voltage_lookup_table_path = "tests/test_data/test_mirror_focus.json"
-    return voltages
+    yield voltages
+    beamline_utils.clear_devices()
 
 
 s03_epics_server_port = getenv("S03_EPICS_CA_SERVER_PORT")