432 create sample stage device for p99 (#458)

* added p99 stages

* fix lint test

* Apply suggestions from CR

* make sample_stage.py.SampleAngleStage signal only

* combined xyz-motors with xyzstage

* changed set_sim to set mock

* remove stages and test

* Fix some docs

---------

Co-authored-by: Dominic Oram <[email protected]>

src/dodal/beamlines/p99.py

@@ -0,0 +1,61 @@
+from dodal.common.beamlines.beamline_utils import device_instantiation, set_beamline
+from dodal.devices.motors import XYZPositioner
+from dodal.devices.p99.sample_stage import FilterMotor, SampleAngleStage
+from dodal.log import set_beamline as set_log_beamline
+from dodal.utils import get_beamline_name
+
+BL = get_beamline_name("BL99P")
+set_log_beamline(BL)
+set_beamline(BL)
+
+
+def sample_angle_stage(
+    wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
+) -> SampleAngleStage:
+    """Sample stage for p99"""
+
+    return device_instantiation(
+        SampleAngleStage,
+        prefix="-MO-STAGE-01:",
+        name="sample_angle_stage",
+        wait=wait_for_connection,
+        fake=fake_with_ophyd_sim,
+    )
+
+
+def sample_stage_filer(
+    wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
+) -> FilterMotor:
+    """Sample stage for p99"""
+
+    return device_instantiation(
+        FilterMotor,
+        prefix="-MO-STAGE-02:MP:SELECT",
+        name="sample_stage_filer",
+        wait=wait_for_connection,
+        fake=fake_with_ophyd_sim,
+    )
+
+
+def sample_xyz_stage(
+    wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
+) -> XYZPositioner:
+    return device_instantiation(
+        FilterMotor,
+        prefix="-MO-STAGE-02:",
+        name="sample_xyz_stage",
+        wait=wait_for_connection,
+        fake=fake_with_ophyd_sim,
+    )
+
+
+def sample_lab_xyz_stage(
+    wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False
+) -> XYZPositioner:
+    return device_instantiation(
+        FilterMotor,
+        prefix="-MO-STAGE-02:LAB:",
+        name="sample_lab_xyz_stage",
+        wait=wait_for_connection,
+        fake=fake_with_ophyd_sim,
+    )

src/dodal/devices/motors.py

@@ -3,8 +3,34 @@
 
 
 class XYZPositioner(Device):
-    def __init__(self, prefix: str, name: str):
-        self.x = Motor(prefix + "X")
-        self.y = Motor(prefix + "Y")
-        self.z = Motor(prefix + "Z")
-        super().__init__(name)
+    """
+
+    Standard ophyd_async xyz motor stage, by combining 3 Motors,
+    with added infix for extra flexibliy to allow different axes other than x,y,z.
+
+    Parameters
+    ----------
+    prefix:
+        EPICS PV (Common part up to and including :).
+    name:
+        name for the stage.
+    infix:
+        EPICS PV, default is the ["X", "Y", "Z"].
+    Notes
+    -----
+    Example usage::
+        async with DeviceCollector():
+            xyz_stage = XYZPositioner("BLXX-MO-STAGE-XX:")
+    Or::
+        with DeviceCollector():
+            xyz_stage = XYZPositioner("BLXX-MO-STAGE-XX:", suffix = ["A", "B", "C"])
+
+    """
+
+    def __init__(self, prefix: str, name: str, infix: list[str] | None = None):
+        if infix is None:
+            infix = ["X", "Y", "Z"]
+        self.x = Motor(prefix + infix[0])
+        self.y = Motor(prefix + infix[1])
+        self.z = Motor(prefix + infix[2])
+        super().__init__(name=name)

src/dodal/devices/p99/sample_stage.py

@@ -0,0 +1,43 @@
+from enum import Enum
+
+from ophyd_async.core import Device
+from ophyd_async.epics.signal import epics_signal_rw
+
+
+class SampleAngleStage(Device):
+    def __init__(self, prefix: str, name: str):
+        self.theta = epics_signal_rw(
+            float, prefix + "WRITETHETA:RBV", prefix + "WRITETHETA"
+        )
+        self.roll = epics_signal_rw(
+            float, prefix + "WRITEROLL:RBV", prefix + "WRITEROLL"
+        )
+        self.pitch = epics_signal_rw(
+            float, prefix + "WRITEPITCH:RBV", prefix + "WRITEPITCH"
+        )
+        super().__init__(name=name)
+
+
+class p99StageSelections(str, Enum):
+    Empty = "Empty"
+    Mn5um = "Mn 5um"
+    Fe = "Fe (empty)"
+    Co5um = "Co 5um"
+    Ni5um = "Ni 5um"
+    Cu5um = "Cu 5um"
+    Zn5um = "Zn 5um"
+    Zr = "Zr (empty)"
+    Mo = "Mo (empty)"
+    Rh = "Rh (empty)"
+    Pd = "Pd (empty)"
+    Ag = "Ag (empty)"
+    Cd25um = "Cd 25um"
+    W = "W (empty)"
+    Pt = "Pt (empty)"
+    User = "User"
+
+
+class FilterMotor(Device):
+    def __init__(self, prefix: str, name: str):
+        self.user_setpoint = epics_signal_rw(p99StageSelections, prefix)
+        super().__init__(name=name)

tests/devices/unit_tests/p99/test_p99_stage.py

@@ -0,0 +1,42 @@
+import pytest
+from ophyd_async.core import DeviceCollector, set_mock_value
+
+from dodal.devices.p99.sample_stage import (
+    FilterMotor,
+    SampleAngleStage,
+    p99StageSelections,
+)
+
+# Long enough for multiple asyncio event loop cycles to run so
+# all the tasks have a chance to run
+A_BIT = 0.001
+
+
+@pytest.fixture
+async def sim_sampleAngleStage():
+    async with DeviceCollector(mock=True):
+        sim_sampleAngleStage = SampleAngleStage(
+            "p99-MO-TABLE-01:", name="sim_sampleAngleStage"
+        )
+        # Signals connected here
+    yield sim_sampleAngleStage
+
+
+@pytest.fixture
+async def sim_filter_wheel():
+    async with DeviceCollector(mock=True):
+        sim_filter_wheel = FilterMotor("p99-MO-TABLE-01:", name="sim_filter_wheel")
+    yield sim_filter_wheel
+
+
+async def test_sampleAngleStage(sim_sampleAngleStage: SampleAngleStage) -> None:
+    assert sim_sampleAngleStage.name == "sim_sampleAngleStage"
+    assert sim_sampleAngleStage.theta.name == "sim_sampleAngleStage-theta"
+    assert sim_sampleAngleStage.roll.name == "sim_sampleAngleStage-roll"
+    assert sim_sampleAngleStage.pitch.name == "sim_sampleAngleStage-pitch"
+
+
+async def test_filter_wheel(sim_filter_wheel: FilterMotor) -> None:
+    assert sim_filter_wheel.name == "sim_filter_wheel"
+    set_mock_value(sim_filter_wheel.user_setpoint, p99StageSelections.Cd25um)
+    assert await sim_filter_wheel.user_setpoint.get_value() == p99StageSelections.Cd25um