Make more use of AsyncStatus.wrap

src/dodal/devices/aperturescatterguard.py

@@ -195,12 +195,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, 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}")
 
-        return AsyncStatus(self._safe_move_within_datacollection_range(pos.location))
+        await self._safe_move_within_datacollection_range(value.location)
 
     def _get_motor_list(self):
         return [

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
@@ -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, **kwargs):
         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, **kwargs):
         await self.thaw_for_time_s.stop()
         await self.control.set(ThawerStates.OFF)

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()

src/dodal/devices/zebra.py

@@ -109,10 +109,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):

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):

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):
@@ -122,11 +123,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)
         )
 

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
@@ -91,7 +91,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(