Make beamstop selected_pos attribute settable (#1227)

* Add beamstop position setter to device

* Raise ValueError if set to UNKNOWN

* Move beamstop tests to mx_phase1 directory

Author: jacob720

src/dodal/devices/mx_phase1/beamstop.py

@@ -1,6 +1,11 @@
+import asyncio
 from math import isclose
 
-from ophyd_async.core import StandardReadable, StrictEnum, derived_signal_r
+from ophyd_async.core import (
+    StandardReadable,
+    StrictEnum,
+    derived_signal_rw,
+)
 from ophyd_async.epics.motor import Motor
 
 from dodal.common.beamlines.beamline_parameters import GDABeamlineParameters
@@ -37,8 +42,7 @@ class Beamstop(StandardReadable):
         x: beamstop x position in mm
         y: beamstop y position in mm
         z: beamstop z position in mm
-        selected_pos: Get the current position of the beamstop as an enum. Currently this
-            is read-only.
+        selected_pos: Get or set the current position of the beamstop as an enum.
     """
 
     def __init__(
@@ -51,10 +55,13 @@ def __init__(
             self.x_mm = Motor(prefix + "X")
             self.y_mm = Motor(prefix + "Y")
             self.z_mm = Motor(prefix + "Z")
-            self.selected_pos = derived_signal_r(
-                self._get_selected_position, x=self.x_mm, y=self.y_mm, z=self.z_mm
+            self.selected_pos = derived_signal_rw(
+                self._get_selected_position,
+                self._set_selected_position,
+                x=self.x_mm,
+                y=self.y_mm,
+                z=self.z_mm,
             )
-
         self._in_beam_xyz_mm = [
             float(beamline_parameters[f"in_beam_{axis}_STANDARD"])
             for axis in ("x", "y", "z")
@@ -77,3 +84,13 @@ def _get_selected_position(self, x: float, y: float, z: float) -> BeamstopPositi
             return BeamstopPositions.DATA_COLLECTION
         else:
             return BeamstopPositions.UNKNOWN
+
+    async def _set_selected_position(self, position: BeamstopPositions) -> None:
+        if position == BeamstopPositions.DATA_COLLECTION:
+            await asyncio.gather(
+                self.x_mm.set(self._in_beam_xyz_mm[0]),
+                self.y_mm.set(self._in_beam_xyz_mm[1]),
+                self.z_mm.set(self._in_beam_xyz_mm[2]),
+            )
+        elif position == BeamstopPositions.UNKNOWN:
+            raise ValueError(f"Cannot set beamstop to position {position}")

tests/devices/mx_phase1/__init__.py

@@ -1,7 +1,8 @@
 from itertools import dropwhile
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
+from bluesky import FailedStatus
 from bluesky import plan_stubs as bps
 from bluesky.preprocessors import run_decorator
 from bluesky.run_engine import RunEngine
@@ -64,3 +65,38 @@ def check_in_beam():
     assert data["beamstop-y_mm"] == y
     assert data["beamstop-z_mm"] == z
     assert data["beamstop-selected_pos"] == expected_pos
+
+
+async def test_set_beamstop_position_to_data_collection_moves_beamstop_into_beam(
+    beamline_parameters: GDABeamlineParameters, RE: RunEngine
+):
+    beamstop = Beamstop("-MO-BS-01:", beamline_parameters, name="beamstop")
+    await beamstop.connect(mock=True)
+
+    beamstop.x_mm.set = AsyncMock()
+    beamstop.y_mm.set = AsyncMock()
+    beamstop.z_mm.set = AsyncMock()
+    set_mock_value(beamstop.x_mm.user_readback, 0)
+    set_mock_value(beamstop.y_mm.user_readback, 0)
+    set_mock_value(beamstop.z_mm.user_readback, 0)
+
+    RE(bps.abs_set(beamstop.selected_pos, BeamstopPositions.DATA_COLLECTION))
+
+    assert beamstop.x_mm.set.call_count == 1
+    assert beamstop.x_mm.set.call_args[0][0] == 1.52
+
+    assert beamstop.y_mm.set.call_count == 1
+    assert beamstop.y_mm.set.call_args[0][0] == 44.78
+
+    assert beamstop.z_mm.set.call_count == 1
+    assert beamstop.z_mm.set.call_args[0][0] == 30.0
+
+
+async def test_set_beamstop_position_to_unknown_raises_error(
+    beamline_parameters: GDABeamlineParameters, RE: RunEngine
+):
+    beamstop = Beamstop("-MO-BS-01:", beamline_parameters, name="beamstop")
+    await beamstop.connect(mock=True)
+    with pytest.raises(FailedStatus) as e:
+        RE(bps.abs_set(beamstop.selected_pos, BeamstopPositions.UNKNOWN, wait=True))
+        assert isinstance(e.value.args[0].exception(), ValueError)