From 9c16e6a6dd136b16e5ddafba46f671f2d2d79c9c Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Mon, 22 Jul 2024 12:38:57 +0100 Subject: [PATCH 1/4] Make more use of AsyncStatus.wrap --- src/dodal/devices/aperturescatterguard.py | 5 +++-- src/dodal/devices/undulator_dcm.py | 14 ++++++-------- src/dodal/devices/zebra.py | 7 +++---- .../unit_tests/test_aperture_scatterguard.py | 11 ++++++----- 4 files changed, 18 insertions(+), 19 deletions(-) diff --git a/src/dodal/devices/aperturescatterguard.py b/src/dodal/devices/aperturescatterguard.py index 96d49265fd..54e67452b6 100644 --- a/src/dodal/devices/aperturescatterguard.py +++ b/src/dodal/devices/aperturescatterguard.py @@ -199,12 +199,13 @@ def load_aperture_positions(self, positions: AperturePositions): LOGGER.info(f"{self.name} loaded in {positions}") self.aperture_positions = positions - def set(self, pos: SingleAperturePosition) -> AsyncStatus: + @AsyncStatus.wrap + async def set(self, pos: SingleAperturePosition): assert isinstance(self.aperture_positions, AperturePositions) if pos not in self.aperture_positions.as_list(): raise InvalidApertureMove(f"Unknown aperture: {pos}") - return AsyncStatus(self._safe_move_within_datacollection_range(pos.location)) + await self._safe_move_within_datacollection_range(pos.location) def _get_motor_list(self): return [ diff --git a/src/dodal/devices/undulator_dcm.py b/src/dodal/devices/undulator_dcm.py index c8efd1ac1a..4fed3d524a 100644 --- a/src/dodal/devices/undulator_dcm.py +++ b/src/dodal/devices/undulator_dcm.py @@ -74,14 +74,12 @@ def __init__( daq_configuration_path + "/domain/beamlineParameters" )["DCM_Perp_Offset_FIXED"] - def set(self, value: float) -> AsyncStatus: - async def _set(): - await asyncio.gather( - self._set_dcm_energy(value), - self._set_undulator_gap_if_required(value), - ) - - return AsyncStatus(_set()) + @AsyncStatus.wrap + async def set(self, value: float): + await asyncio.gather( + self._set_dcm_energy(value), + self._set_undulator_gap_if_required(value), + ) async def _set_dcm_energy(self, energy_kev: float) -> None: access_level = await self.undulator.gap_access.get_value() diff --git a/src/dodal/devices/zebra.py b/src/dodal/devices/zebra.py index 2abaa1f7c1..c1d904b380 100644 --- a/src/dodal/devices/zebra.py +++ b/src/dodal/devices/zebra.py @@ -110,10 +110,9 @@ async def _set_armed(self, demand: ArmDemand): if reading == demand.value: return - def set(self, demand: ArmDemand) -> AsyncStatus: - return AsyncStatus( - asyncio.wait_for(self._set_armed(demand), timeout=self.TIMEOUT) - ) + @AsyncStatus.wrap + async def set(self, demand: ArmDemand): + await asyncio.wait_for(self._set_armed(demand), timeout=self.TIMEOUT) class PositionCompare(StandardReadable): diff --git a/tests/devices/unit_tests/test_aperture_scatterguard.py b/tests/devices/unit_tests/test_aperture_scatterguard.py index 6c40e01e39..3c6e067648 100644 --- a/tests/devices/unit_tests/test_aperture_scatterguard.py +++ b/tests/devices/unit_tests/test_aperture_scatterguard.py @@ -7,6 +7,7 @@ import pytest from bluesky.run_engine import RunEngine from ophyd_async.core import ( + DeviceCollector, get_mock_put, set_mock_value, ) @@ -36,10 +37,10 @@ def get_all_motors(ap_sg: ApertureScatterguard): @pytest.fixture -async def ap_sg_and_call_log(aperture_positions: AperturePositions): +async def ap_sg_and_call_log(RE: RunEngine, aperture_positions: AperturePositions): call_log = MagicMock() - ap_sg = ApertureScatterguard(name="test_ap_sg") - await ap_sg.connect(mock=True) + async with DeviceCollector(mock=True): + ap_sg = ApertureScatterguard(name="test_ap_sg") ap_sg.load_aperture_positions(aperture_positions) with ExitStack() as motor_patch_stack: for motor in get_all_motors(ap_sg): @@ -121,11 +122,11 @@ def _assert_patched_ap_sg_has_call( ) -def test_aperture_scatterguard_rejects_unknown_position(aperture_in_medium_pos): +async def test_aperture_scatterguard_rejects_unknown_position(aperture_in_medium_pos): position_to_reject = ApertureFiveDimensionalLocation(0, 0, 0, 0, 0) with pytest.raises(InvalidApertureMove): - aperture_in_medium_pos.set( + await aperture_in_medium_pos.set( SingleAperturePosition("test", "GDA_NAME", 10, position_to_reject) ) From 85a5a8d2ea04b49421c0b7fec9b2afc38c6b0947 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Mon, 22 Jul 2024 14:52:11 +0100 Subject: [PATCH 2/4] Use RE in smargon so it has an event loop --- tests/devices/unit_tests/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/devices/unit_tests/conftest.py b/tests/devices/unit_tests/conftest.py index 3f6f1d82d4..247ad2967e 100644 --- a/tests/devices/unit_tests/conftest.py +++ b/tests/devices/unit_tests/conftest.py @@ -2,6 +2,7 @@ from pathlib import Path import pytest +from bluesky.run_engine import RunEngine from ophyd_async.core import ( DirectoryInfo, DirectoryProvider, @@ -39,7 +40,7 @@ def static_directory_provider(tmp_path: Path) -> DirectoryProvider: @pytest.fixture -def smargon(): +def smargon(RE: RunEngine): smargon = i03.smargon(fake_with_ophyd_sim=True) def mock_set(motor, value, *args, **kwargs): From 7deec9432f03390b6457628d3106bc5795dbb2c1 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Mon, 22 Jul 2024 16:36:40 +0100 Subject: [PATCH 3/4] Fix signature of stop --- src/dodal/devices/thawer.py | 6 ++++-- tests/devices/unit_tests/test_thawer.py | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/dodal/devices/thawer.py b/src/dodal/devices/thawer.py index abd9c748c2..32c849bc3b 100644 --- a/src/dodal/devices/thawer.py +++ b/src/dodal/devices/thawer.py @@ -32,7 +32,8 @@ async def set(self, time_to_thaw_for: float): finally: await self._control_signal.set(ThawerStates.OFF) - async def stop(self): + @AsyncStatus.wrap + async def stop(self, *args): if self._thawing_task: self._thawing_task.cancel() @@ -43,6 +44,7 @@ def __init__(self, prefix: str, name: str = "") -> None: self.thaw_for_time_s = ThawingTimer(self.control) super().__init__(name) - async def stop(self): + @AsyncStatus.wrap + async def stop(self, *args): await self.thaw_for_time_s.stop() await self.control.set(ThawerStates.OFF) diff --git a/tests/devices/unit_tests/test_thawer.py b/tests/devices/unit_tests/test_thawer.py index da145fe7c2..2ef257a0b9 100644 --- a/tests/devices/unit_tests/test_thawer.py +++ b/tests/devices/unit_tests/test_thawer.py @@ -1,5 +1,5 @@ import asyncio -from unittest.mock import ANY, AsyncMock, MagicMock, call, patch +from unittest.mock import ANY, AsyncMock, call, patch import pytest from ophyd_async.core import DeviceCollector, get_mock_put @@ -87,7 +87,7 @@ async def test_given_thawing_already_triggered_when_stop_called_then_stop_thawin async def test_calling_stop_on_thawer_stops_thawing_timer_and_turns_thawer_off( thawer: Thawer, ): - thawer.thaw_for_time_s = MagicMock(spec=ThawingTimer) + thawer.thaw_for_time_s.stop = AsyncMock(spec=ThawingTimer) await thawer.stop() thawer.thaw_for_time_s.stop.assert_called_once() get_mock_put(thawer.control).assert_called_once_with( From 4e52a00be2dd80b7c7ba6dcf9476ea572d5e48a6 Mon Sep 17 00:00:00 2001 From: Dominic Oram Date: Tue, 23 Jul 2024 11:27:36 +0100 Subject: [PATCH 4/4] Fix some typing issues --- src/dodal/devices/aperturescatterguard.py | 8 ++++---- src/dodal/devices/thawer.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/dodal/devices/aperturescatterguard.py b/src/dodal/devices/aperturescatterguard.py index 5537b59755..755543fd93 100644 --- a/src/dodal/devices/aperturescatterguard.py +++ b/src/dodal/devices/aperturescatterguard.py @@ -196,12 +196,12 @@ def load_aperture_positions(self, positions: AperturePositions): self.aperture_positions = positions @AsyncStatus.wrap - async def set(self, pos: SingleAperturePosition): + async def set(self, value: SingleAperturePosition): assert isinstance(self.aperture_positions, AperturePositions) - if pos not in self.aperture_positions.as_list(): - raise InvalidApertureMove(f"Unknown aperture: {pos}") + if value not in self.aperture_positions.as_list(): + raise InvalidApertureMove(f"Unknown aperture: {value}") - await self._safe_move_within_datacollection_range(pos.location) + await self._safe_move_within_datacollection_range(value.location) def _get_motor_list(self): return [ diff --git a/src/dodal/devices/thawer.py b/src/dodal/devices/thawer.py index 32c849bc3b..f329024b60 100644 --- a/src/dodal/devices/thawer.py +++ b/src/dodal/devices/thawer.py @@ -15,7 +15,7 @@ class ThawerStates(str, Enum): ON = "On" -class ThawingTimer(Device): +class ThawingTimer(Device, Stoppable): def __init__(self, control_signal: SignalRW[ThawerStates]) -> None: self._control_signal = control_signal self._thawing_task: Task | None = None @@ -33,7 +33,7 @@ async def set(self, time_to_thaw_for: float): await self._control_signal.set(ThawerStates.OFF) @AsyncStatus.wrap - async def stop(self, *args): + async def stop(self, *args, **kwargs): if self._thawing_task: self._thawing_task.cancel() @@ -45,6 +45,6 @@ def __init__(self, prefix: str, name: str = "") -> None: super().__init__(name) @AsyncStatus.wrap - async def stop(self, *args): + async def stop(self, *args, **kwargs): await self.thaw_for_time_s.stop() await self.control.set(ThawerStates.OFF)