Undulator order (harmonics) class (#1661)

* all in one shot

* improve error message, rename internal signal

---------

Co-authored-by: eir17846 <[email protected]>
Co-authored-by: oliwenmandiamond <[email protected]>

Author: 3 people

src/dodal/beamlines/i09_1.py

@@ -11,6 +11,7 @@
 from dodal.devices.electron_analyser.specs import SpecsDetector
 from dodal.devices.i09_1 import LensMode, PsuMode
 from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.undulator import UndulatorOrder
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
@@ -47,3 +48,8 @@ def analyser() -> SpecsDetector[LensMode, PsuMode]:
         psu_mode_type=PsuMode,
         energy_source=energy_source(),
     )
+
+
+@device_factory()
+def harmonics() -> UndulatorOrder:
+    return UndulatorOrder(name="harmonics")

src/dodal/devices/undulator.py

@@ -1,14 +1,15 @@
 import os
 
 import numpy as np
-from bluesky.protocols import Movable
+from bluesky.protocols import Locatable, Location, Movable
 from numpy import ndarray
 from ophyd_async.core import (
     AsyncStatus,
     Reference,
     StandardReadable,
     StandardReadableFormat,
     soft_signal_r_and_setter,
+    soft_signal_rw,
 )
 from ophyd_async.epics.core import epics_signal_r
 from ophyd_async.epics.motor import Motor
@@ -157,3 +158,30 @@ async def _get_gap_to_match_energy(self, energy_kev: float) -> float:
             energy_kev * 1000,
             energy_to_distance_table,
         )
+
+
+class UndulatorOrder(StandardReadable, Locatable[int]):
+    """
+    Represents the order of an undulator device. Allows setting and locating the order.
+    """
+
+    def __init__(self, name: str = "") -> None:
+        """
+        Args:
+            name: Name for device. Defaults to ""
+        """
+        with self.add_children_as_readables():
+            self._value = soft_signal_rw(int, initial_value=3)
+        super().__init__(name=name)
+
+    @AsyncStatus.wrap
+    async def set(self, value: int) -> None:
+        if (value >= 0) and isinstance(value, int):
+            await self._value.set(value)
+        else:
+            raise ValueError(
+                f"Undulator order must be a positive integer. Requested value: {value}"
+            )
+
+    async def locate(self) -> Location[int]:
+        return await self._value.locate()

tests/devices/test_undulator.py

@@ -3,6 +3,8 @@
 
 import numpy as np
 import pytest
+from bluesky import RunEngine
+from bluesky.plan_stubs import mv
 from ophyd_async.core import init_devices
 from ophyd_async.testing import (
     assert_configuration,
@@ -17,6 +19,7 @@
 from dodal.devices.undulator import (
     AccessError,
     Undulator,
+    UndulatorOrder,
     _get_gap_for_energy,
 )
 from dodal.testing import patch_all_motors
@@ -48,6 +51,13 @@ def undulator_in_commissioning_mode(
     yield undulator
 
 
+@pytest.fixture
+async def undulator_order() -> UndulatorOrder:
+    async with init_devices(mock=True):
+        order = UndulatorOrder(name="undulator_order")
+    return order
+
+
 async def test_reading_includes_read_fields(undulator: Undulator):
     await assert_reading(
         undulator,
@@ -147,3 +157,36 @@ async def test_gap_access_check_move_not_inhibited_when_commissioning_mode_disab
         get_mock_put(undulator.gap_motor.user_setpoint).assert_called_once_with(
             15.0, wait=True
         )
+
+
+async def test_order_read(
+    undulator_order: UndulatorOrder,
+):
+    await assert_reading(
+        undulator_order,
+        {"undulator_order-_value": partial_reading(3)},
+    )
+
+
+async def test_move_order(
+    undulator_order: UndulatorOrder,
+    run_engine: RunEngine,
+):
+    assert (await undulator_order.locate())["readback"] == 3  # default order
+    run_engine(mv(undulator_order, 1))
+    assert (await undulator_order.locate())["readback"] == 1  # no error
+
+
+@pytest.mark.parametrize(
+    "order_value",
+    [-1, 1.56],
+)
+async def test_move_order_fails(
+    undulator_order: UndulatorOrder,
+    order_value: float | int,
+):
+    with pytest.raises(
+        ValueError,
+        match=f"Undulator order must be a positive integer. Requested value: {order_value}",
+    ):
+        await undulator_order.set(order_value)  # type: ignore