1495 make apple2 and pgm start flying in sync

src/dodal/devices/apple2_undulator.py

@@ -11,6 +11,7 @@
     AsyncStatus,
     Device,
     FlyMotorInfo,
+    Reference,
     SignalR,
     SignalW,
     StandardReadable,
@@ -19,15 +20,18 @@
     WatchableAsyncStatus,
     WatcherUpdate,
     derived_signal_rw,
+    error_if_none,
     observe_value,
     soft_signal_r_and_setter,
+    soft_signal_rw,
     wait_for_value,
 )
 from ophyd_async.core import StandardReadableFormat as Format
 from ophyd_async.epics.core import epics_signal_r, epics_signal_rw
 from ophyd_async.epics.motor import Motor
 from pydantic import BaseModel, ConfigDict, RootModel
 
+from dodal.devices.pgm import PGM
 from dodal.log import LOGGER
 
 T = TypeVar("T")
@@ -794,3 +798,77 @@ def determine_phase_from_hardware(
 
         LOGGER.warning("Unable to determine polarisation. Defaulting to NONE.")
         return Pol.NONE, 0.0
+
+
+Apple2ID = TypeVar("Apple2ID", bound=Apple2)
+
+
+class EnergySetter(
+    StandardReadable, Movable[float], Preparable, Flyable, Generic[Apple2ID]
+):
+    """
+    Compound device to set both ID and PGM energy at the same time.
+
+    """
+
+    def __init__(self, id: Apple2ID, pgm: PGM, name: str = "") -> None:
+        """
+        Parameters
+        ----------
+        id:
+            An Apple2 device (should be I10Apple2 or subclass).
+        pgm:
+            A PGM/mono device.
+        name:
+            New device name.
+        """
+
+        self.id = id
+        self.pgm_ref = Reference(pgm)
+        self.add_readables(
+            [self.id.energy, self.pgm_ref().energy.user_readback],
+            StandardReadableFormat.HINTED_SIGNAL,
+        )
+
+        with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL):
+            self.energy_offset = soft_signal_rw(float, initial_value=0)
+        super().__init__(name=name)
+        self._fly_status: AsyncStatus | None = None
+
+    @AsyncStatus.wrap
+    async def set(self, value: float) -> None:
+        LOGGER.info(f"Moving f{self.name} energy to {value}.")
+        await asyncio.gather(
+            self.id.set(value=value + await self.energy_offset.get_value()),
+            self.pgm_ref().energy.set(value),
+        )
+
+    @AsyncStatus.wrap
+    async def prepare(self, value: FlyMotorInfo) -> None:
+        await asyncio.gather(
+            self.id.prepare(value), self.pgm_ref().energy.prepare(value)
+        )
+
+    @AsyncStatus.wrap
+    async def kickoff(self):
+        pgm_acceleration_time, gap_acceleration_time = await asyncio.gather(
+            self.pgm_ref().energy.acceleration_time.get_value(),
+            self.id.gap.acceleration_time.get_value(),
+        )
+        start_offset_time = pgm_acceleration_time - gap_acceleration_time
+
+        await self.pgm_ref().energy.kickoff()
+        await asyncio.sleep(start_offset_time)
+        await self.id.kickoff()
+        self._fly_status = self._combined_fly_status()
+
+    def complete(self) -> AsyncStatus:
+        """Stop when both pgm and id is done moving."""
+        fly_status = error_if_none(self._fly_status, "kickoff not called")
+        return fly_status
+
+    @AsyncStatus.wrap
+    async def _combined_fly_status(self):
+        status_pgm = self.pgm_ref().energy.complete()
+        status_id = self.id.complete()
+        await asyncio.gather(status_pgm, status_id)

src/dodal/devices/i10/i10_apple2.py

@@ -1,4 +1,3 @@
-import asyncio
 import csv
 from dataclasses import dataclass
 from pathlib import Path
@@ -13,9 +12,9 @@
     StandardReadable,
     StandardReadableFormat,
     soft_signal_r_and_setter,
-    soft_signal_rw,
 )
 
+from dodal.devices.apple2_undulator import EnergySetter
 from dodal.log import LOGGER
 
 from ..apple2_undulator import (
@@ -191,44 +190,6 @@ def update_lookuptable(self):
         self._available_pol = list(self.lookup_tables["Gap"].keys())
 
 
-class EnergySetter(StandardReadable, Movable[float]):
-    """
-    Compound device to set both ID and PGM energy at the same time.
-
-    """
-
-    def __init__(self, id: I10Apple2, pgm: PGM, name: str = "") -> None:
-        """
-        Parameters
-        ----------
-        id:
-            An Apple2 device.
-        pgm:
-            A PGM/mono device.
-        name:
-            New device name.
-        """
-        super().__init__(name=name)
-        self.id = id
-        self.pgm_ref = Reference(pgm)
-
-        self.add_readables(
-            [self.id.energy, self.pgm_ref().energy.user_readback],
-            StandardReadableFormat.HINTED_SIGNAL,
-        )
-
-        with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL):
-            self.energy_offset = soft_signal_rw(float, initial_value=0)
-
-    @AsyncStatus.wrap
-    async def set(self, value: float) -> None:
-        LOGGER.info(f"Moving f{self.name} energy to {value}.")
-        await asyncio.gather(
-            self.id.set(value=value + await self.energy_offset.get_value()),
-            self.pgm_ref().energy.set(value),
-        )
-
-
 class I10Apple2Pol(StandardReadable, Movable[Pol]):
     """
     Compound device to set polorisation of ID.

tests/devices/i10/test_i10Apple2.py

@@ -1,16 +1,23 @@
+import asyncio
 import os
 import pickle
 from collections import defaultdict
 from unittest import mock
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock, patch
 
 import numpy as np
 import pytest
-from bluesky.plan_stubs import prepare
+from bluesky import plan_stubs as bps
 from bluesky.plans import scan
 from bluesky.run_engine import RunEngine
 from numpy import poly1d
-from ophyd_async.core import FlyMotorInfo, init_devices
+from ophyd_async.core import (
+    AsyncStatus,
+    FlyMotorInfo,
+    init_devices,
+    soft_signal_rw,
+    wait_for_value,
+)
 from ophyd_async.testing import (
     assert_emitted,
     callback_on_mock_put,
@@ -28,6 +35,7 @@
 )
 
 from dodal.devices.apple2_undulator import (
+    EnergySetter,
     Pol,
     UndulatorGap,
     UndulatorGateStatus,
@@ -36,7 +44,6 @@
 )
 from dodal.devices.i10.i10_apple2 import (
     DEFAULT_JAW_PHASE_POLY_PARAMS,
-    EnergySetter,
     I10Apple2,
     I10Apple2Pol,
     LinearArbitraryAngle,
@@ -598,10 +605,69 @@ def test_convert_csv_to_lookup_failed():
         )
 
 
-async def test_i10apple2_prepare_success(RE: RunEngine, mock_id: I10Apple2):
-    fly_motor_info = FlyMotorInfo(
-        start_position=600,
-        end_position=700,
-        time_for_move=60,
+async def test_energysetter_prepare_success(RE: RunEngine, mock_id_pgm: EnergySetter):
+    fly_info = FlyMotorInfo(start_position=700, end_position=800, time_for_move=10)
+    mock_id_pgm.id.prepare = AsyncMock()
+    mock_id_pgm.pgm_ref().energy.prepare = AsyncMock()
+    RE(bps.prepare(mock_id_pgm, fly_info))
+    mock_id_pgm.id.prepare.assert_awaited_once_with(fly_info)
+    mock_id_pgm.pgm_ref().energy.prepare.assert_awaited_once_with(fly_info)  # type: ignore
+
+
+@patch("asyncio.sleep", new_callable=AsyncMock)
+async def test_energysetter_kickoff_set_correct_delay(
+    mock_sleep: AsyncMock, RE: RunEngine, mock_id_pgm: EnergySetter
+):
+    fly_info = FlyMotorInfo(start_position=700, end_position=800, time_for_move=10)
+    id_acc_time = 3
+    pgm_acc_time = 1
+    set_mock_value(mock_id_pgm.id.gap.max_velocity, 30)
+    set_mock_value(mock_id_pgm.id.gap.min_velocity, 0.1)
+    set_mock_value(mock_id_pgm.id.gap.acceleration_time, id_acc_time)
+    set_mock_value(mock_id_pgm.pgm_ref().energy.max_velocity, 30)
+    set_mock_value(mock_id_pgm.id.gap.low_limit_travel, 0)
+    set_mock_value(mock_id_pgm.id.gap.high_limit_travel, 200)
+    set_mock_value(mock_id_pgm.pgm_ref().energy.low_limit_travel, 0)
+    set_mock_value(mock_id_pgm.pgm_ref().energy.high_limit_travel, 1000)
+    set_mock_value(mock_id_pgm.pgm_ref().energy.acceleration_time, pgm_acc_time)
+    set_mock_value(mock_id_pgm.id.gap.gate, UndulatorGateStatus.CLOSE)
+    mock_id_pgm.id.kickoff = AsyncMock()
+    mock_id_pgm.pgm_ref().energy.kickoff = AsyncMock()
+    await mock_id_pgm.prepare(fly_info)
+    await mock_id_pgm.kickoff()
+    mock_sleep.assert_awaited_once_with(pgm_acc_time - id_acc_time)
+    mock_id_pgm.id.kickoff.assert_awaited_once()
+    mock_id_pgm.pgm_ref().energy.kickoff.assert_awaited_once()  # type: ignore
+
+
+@patch("asyncio.sleep", new_callable=AsyncMock)
+async def test_energysetter_complete(mock_sleep, mock_id_pgm: EnergySetter) -> None:
+    fake_id_fly_status = soft_signal_rw(bool, initial_value=False)
+    fake_pgm_fly_status = soft_signal_rw(bool, initial_value=False)
+    await asyncio.gather(
+        fake_pgm_fly_status.connect(mock=True), fake_id_fly_status.connect(mock=True)
     )
-    RE(prepare(mock_id, fly_motor_info))
+
+    @AsyncStatus.wrap
+    async def get_status(signal):
+        await wait_for_value(signal, True, timeout=1)
+
+    mock_id_pgm.id.gap._fly_status = get_status(fake_id_fly_status)
+
+    mock_id_pgm.pgm_ref().energy._fly_status = get_status(fake_pgm_fly_status)  # type: ignore
+    mock_id_pgm.id.kickoff = AsyncMock()
+    mock_id_pgm.pgm_ref().energy.kickoff = AsyncMock()
+    pgm_status = mock_id_pgm.pgm_ref().energy.complete()
+    id_status = mock_id_pgm.id.complete()
+    assert not id_status.done
+    assert not pgm_status.done
+    await mock_id_pgm.kickoff()
+    energy_setter_fly_status = mock_id_pgm.complete()
+    assert not energy_setter_fly_status.done
+    await fake_id_fly_status.set(True)
+    assert mock_id_pgm.id.gap._fly_status.done
+    assert not energy_setter_fly_status.done
+    await fake_pgm_fly_status.set(True)
+    assert mock_id_pgm.pgm_ref().energy._fly_status.done  # type: ignore
+    await energy_setter_fly_status
+    assert energy_setter_fly_status.done

tests/devices/test_apple2_undulator.py

@@ -5,7 +5,7 @@
 import pytest
 from bluesky.plans import scan
 from bluesky.run_engine import RunEngine
-from ophyd_async.core import AsyncStatus, FlyMotorInfo, init_devices
+from ophyd_async.core import AsyncStatus, FlyMotorInfo, StrictEnum, init_devices
 from ophyd_async.testing import (
     assert_configuration,
     assert_emitted,
@@ -20,13 +20,15 @@
     DEFAULT_TIMEOUT,
     Apple2,
     Apple2PhasesVal,
+    EnergySetter,
     MotorWithoutStop,
     Pol,
     UndulatorGap,
     UndulatorGateStatus,
     UndulatorJawPhase,
     UndulatorPhaseAxes,
 )
+from dodal.devices.pgm import PGM
 
 
 @pytest.fixture
@@ -101,7 +103,6 @@ def __init__(
         id_gap: UndulatorGap,
         id_phase: UndulatorPhaseAxes,
         prefix: str = "",
-        name: str = "",
     ) -> None:
         super().__init__(id_gap, id_phase, prefix)
 
@@ -120,6 +121,24 @@ async def mock_apple2(
     return mock_apple2
 
 
+class MockGrating(StrictEnum):
+    LINES = "lines"
+
+
+@pytest.fixture
+async def mock_pgm() -> PGM:
+    async with init_devices(mock=True):
+        mock_pgm = PGM(prefix="", grating=MockGrating)
+    return mock_pgm
+
+
+@pytest.fixture
+async def mock_energy_setter(mock_apple2: test_apple2, mock_pgm: PGM) -> EnergySetter:
+    async with init_devices(mock=True):
+        mock_energy_setter = EnergySetter(id=mock_apple2, pgm=mock_pgm)
+    return mock_energy_setter
+
+
 async def test_in_motion_error(
     mock_id_gap: UndulatorGap,
     mock_phaseAxes: UndulatorPhaseAxes,
@@ -495,17 +514,30 @@ async def test_apple2_prepare_success(
     assert await mock_apple2.gap.velocity.get_value() == abs(velocity)
 
 
-async def test_apple2_kickoff__call_gap_kickoff(
+async def test_apple2_kickoff_call_gap_kickoff(
     mock_apple2: test_apple2,
 ):
     mock_apple2.gap.kickoff = AsyncMock()
     await mock_apple2.kickoff()
     mock_apple2.gap.kickoff.assert_awaited_once()
 
 
-def test_apple2_complete__call_gap_complete(
+def test_apple2_complete_call_gap_complete(
     mock_apple2: test_apple2,
 ):
     mock_apple2.gap.complete = MagicMock()
     mock_apple2.complete()
     mock_apple2.gap.complete.assert_called_once()
+
+
+async def test_energy_setter_prepare_success(
+    mock_energy_setter: EnergySetter,
+    RE: RunEngine,
+):
+    mock_energy_setter.id.prepare = AsyncMock()
+    mock_energy_setter.pgm_ref().energy.prepare = AsyncMock()
+
+    fly_info = FlyMotorInfo(start_position=700, end_position=800, time_for_move=10)
+    RE(bps.prepare(mock_energy_setter, fly_info, wait=True))
+    mock_energy_setter.id.prepare.assert_awaited_once_with(fly_info)
+    mock_energy_setter.pgm_ref().energy.prepare.assert_awaited_once_with(fly_info)  # type: ignore