Allow triggering an unload by calling robot.set(None) (#1678)

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

Author: rtuck99

src/dodal/devices/robot.py

@@ -61,7 +61,7 @@ async def raise_if_error(self, raise_from: Exception):
             raise RobotLoadError(int(error_code), error_string) from raise_from
 
 
-class BartRobot(StandardReadable, Movable[SampleLocation]):
+class BartRobot(StandardReadable, Movable[SampleLocation | None]):
     """The sample changing robot."""
 
     # How long to wait for the robot if it is busy soaking/drying
@@ -104,7 +104,6 @@ def __init__(self, prefix: str, name: str = "") -> None:
         self.init = epics_signal_x(prefix + "INIT.PROC")
         self.soak = epics_signal_x(prefix + "SOAK.PROC")
         self.home = epics_signal_x(prefix + "GOHM.PROC")
-        self.unload = epics_signal_x(prefix + "UNLD.PROC")
         self.dry = epics_signal_x(prefix + "DRY.PROC")
         self.open = epics_signal_x(prefix + "COLO.PROC")
         self.close = epics_signal_x(prefix + "COLC.PROC")
@@ -175,12 +174,23 @@ async def _load_pin_and_puck(self, sample_location: SampleLocation):
         await self.pin_mounted_or_no_pin_found()
 
     @AsyncStatus.wrap
-    async def set(self, value: SampleLocation):
+    async def set(self, value: SampleLocation | None):
+        """
+        Perform a sample load from the specified sample location
+        Args:
+            value: The pin and puck to load, or None to unload the sample.
+        Raises:
+            RobotLoadError if a timeout occurs, or if an error occurs loading the smaple.
+        """
         try:
-            await wait_for(
-                self._load_pin_and_puck(value),
-                timeout=self.LOAD_TIMEOUT + self.NOT_BUSY_TIMEOUT,
-            )
+            if value is not None:
+                await wait_for(
+                    self._load_pin_and_puck(value),
+                    timeout=self.LOAD_TIMEOUT + self.NOT_BUSY_TIMEOUT,
+                )
+            else:
+                await self.unload.trigger(timeout=self.LOAD_TIMEOUT)
+                await wait_for_value(self.program_running, False, self.NOT_BUSY_TIMEOUT)
         except TimeoutError as e:
             await self.prog_error.raise_if_error(e)
             await self.controller_error.raise_if_error(e)

tests/devices/test_bart_robot.py

@@ -1,5 +1,6 @@
+import asyncio
 import traceback
-from asyncio import create_task
+from asyncio import Event, create_task
 from functools import partial
 from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
 
@@ -22,6 +23,29 @@
 )
 
 
+@pytest.fixture
+async def robot_for_unload():
+    device = BartRobot("robot", "-MO-ROBOT-01:")
+    device.NOT_BUSY_TIMEOUT = 0.3  # type: ignore
+    device.LOAD_TIMEOUT = 0.3  # type: ignore
+    await device.connect(mock=True)
+
+    trigger_complete = Event()
+    drying_complete = Event()
+
+    async def finish_later():
+        await drying_complete.wait()
+        set_mock_value(device.program_running, False)
+
+    async def fake_unload(*args, **kwargs):
+        set_mock_value(device.program_running, True)
+        await trigger_complete.wait()
+        asyncio.create_task(finish_later())
+
+    get_mock_put(device.unload).side_effect = fake_unload
+    return device, trigger_complete, drying_complete
+
+
 async def _get_bart_robot() -> BartRobot:
     device = BartRobot("robot", "-MO-ROBOT-01:")
     device.LOAD_TIMEOUT = 1  # type: ignore
@@ -206,3 +230,41 @@ async def test_moving_the_robot_will_reset_error_if_light_curtain_is_tripped_and
             call.load.put(None, wait=True),
         ]
     )
+
+
+async def test_unloading_the_robot_waits_for_drying_to_complete(robot_for_unload):
+    robot, trigger_completed, drying_completed = robot_for_unload
+    drying_completed.set()
+    unload_status = robot.set(None)
+
+    await asyncio.sleep(0.1)
+    assert not unload_status.done
+    get_mock_put(robot.unload).assert_called_once()
+
+    trigger_completed.set()
+    await unload_status
+    assert unload_status.done
+
+
+async def test_unloading_the_robot_times_out_if_unloading_takes_too_long(
+    robot_for_unload,
+):
+    robot, trigger_completed, drying_completed = robot_for_unload
+    drying_completed.set()
+    unload_status = robot.set(None)
+
+    with pytest.raises(RobotLoadError) as exc_info:
+        await unload_status
+
+    assert isinstance(exc_info.value.__cause__, TimeoutError)
+
+
+async def test_unloading_the_robot_times_out_if_drying_takes_too_long(robot_for_unload):
+    robot, trigger_completed, drying_completed = robot_for_unload
+    trigger_completed.set()
+    unload_status = robot.set(None)
+
+    with pytest.raises(RobotLoadError) as exc_info:
+        await unload_status
+
+    assert isinstance(exc_info.value.__cause__, TimeoutError)