Add logic to stop waiting on pin if robot raises pin not found error (#…

…629)

* Add logic to stop waiting on pin if robot raises pin not found error

* Fix unit tests

src/dodal/devices/robot.py

@@ -1,15 +1,33 @@
 import asyncio
+from asyncio import FIRST_COMPLETED, Task
 from dataclasses import dataclass
 from enum import Enum
 
 from bluesky.protocols import Movable
-from ophyd_async.core import AsyncStatus, StandardReadable, wait_for_value
+from ophyd_async.core import (
+    AsyncStatus,
+    StandardReadable,
+    set_and_wait_for_value,
+    wait_for_value,
+)
 from ophyd_async.epics.signal import epics_signal_r, epics_signal_x
 from ophyd_async.epics.signal.signal import epics_signal_rw_rbv
 
 from dodal.log import LOGGER
 
 
+class RobotLoadFailed(Exception):
+    error_code: int
+    error_string: str
+
+    def __init__(self, error_code: int, error_string: str) -> None:
+        self.error_code, self.error_string = error_code, error_string
+        super().__init__(error_string)
+
+    def __str__(self) -> str:
+        return self.error_string
+
+
 @dataclass
 class SampleLocation:
     puck: int
@@ -25,6 +43,7 @@ class BartRobot(StandardReadable, Movable):
     """The sample changing robot."""
 
     LOAD_TIMEOUT = 60
+    NO_PIN_ERROR_CODE = 25
 
     def __init__(
         self,
@@ -40,8 +59,34 @@ def __init__(
         self.load = epics_signal_x(prefix + "LOAD.PROC")
         self.program_running = epics_signal_r(bool, prefix + "PROGRAM_RUNNING")
         self.program_name = epics_signal_r(str, prefix + "PROGRAM_NAME")
+        self.error_str = epics_signal_r(str, prefix + "PRG_ERR_MSG")
+        self.error_code = epics_signal_r(int, prefix + "PRG_ERR_CODE")
         super().__init__(name=name)
 
+    async def pin_mounted_or_no_pin_found(self):
+        """This co-routine will finish when either a pin is detected or the robot gives
+        an error saying no pin was found (whichever happens first). In the case where no
+        pin was found a RobotLoadFailed error is raised.
+        """
+
+        async def raise_if_no_pin():
+            await wait_for_value(self.error_code, self.NO_PIN_ERROR_CODE, None)
+            raise RobotLoadFailed(self.NO_PIN_ERROR_CODE, "Pin was not detected")
+
+        finished, unfinished = await asyncio.wait(
+            [
+                Task(raise_if_no_pin()),
+                Task(
+                    wait_for_value(self.gonio_pin_sensor, PinMounted.PIN_MOUNTED, None)
+                ),
+            ],
+            return_when=FIRST_COMPLETED,
+        )
+        for task in unfinished:
+            task.cancel()
+        for task in finished:
+            await task
+
     async def _load_pin_and_puck(self, sample_location: SampleLocation):
         LOGGER.info(f"Loading pin {sample_location}")
         if await self.program_running.get_value():
@@ -50,19 +95,24 @@ async def _load_pin_and_puck(self, sample_location: SampleLocation):
             )
             await wait_for_value(self.program_running, False, None)
         await asyncio.gather(
-            self.next_puck.set(sample_location.puck),
-            self.next_pin.set(sample_location.pin),
+            set_and_wait_for_value(self.next_puck, sample_location.puck),
+            set_and_wait_for_value(self.next_pin, sample_location.pin),
         )
         await self.load.trigger()
         if await self.gonio_pin_sensor.get_value() == PinMounted.PIN_MOUNTED:
             LOGGER.info("Waiting on old pin unloaded")
             await wait_for_value(self.gonio_pin_sensor, PinMounted.NO_PIN_MOUNTED, None)
         LOGGER.info("Waiting on new pin loaded")
-        await wait_for_value(self.gonio_pin_sensor, PinMounted.PIN_MOUNTED, None)
 
-    def set(self, sample_location: SampleLocation) -> AsyncStatus:
-        return AsyncStatus(
-            asyncio.wait_for(
+        await self.pin_mounted_or_no_pin_found()
+
+    @AsyncStatus.wrap
+    async def set(self, sample_location: SampleLocation):
+        try:
+            await asyncio.wait_for(
                 self._load_pin_and_puck(sample_location), timeout=self.LOAD_TIMEOUT
             )
-        )
+        except asyncio.TimeoutError:
+            error_code = await self.error_code.get_value()
+            error_string = await self.error_str.get_value()
+            raise RobotLoadFailed(error_code, error_string)

tests/devices/unit_tests/test_bart_robot.py

@@ -1,10 +1,10 @@
-from asyncio import TimeoutError, sleep
+from asyncio import create_task, sleep
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from ophyd_async.core import set_mock_value
 
-from dodal.devices.robot import BartRobot, PinMounted, SampleLocation
+from dodal.devices.robot import BartRobot, PinMounted, RobotLoadFailed, SampleLocation
 
 
 async def _get_bart_robot() -> BartRobot:
@@ -19,6 +19,24 @@ async def test_bart_robot_can_be_connected_in_sim_mode():
     await device.connect(mock=True)
 
 
+async def test_given_robot_load_times_out_when_load_called_then_exception_contains_error_info():
+    device = await _get_bart_robot()
+
+    async def _sleep(*args):
+        await sleep(1)
+
+    device._load_pin_and_puck = AsyncMock(side_effect=_sleep)
+
+    set_mock_value(device.error_code, (expected_error_code := 10))
+    set_mock_value(device.error_str, (expected_error_string := "BAD"))
+
+    with pytest.raises(RobotLoadFailed) as e:
+        await device.set(SampleLocation(0, 0))
+    assert e.value.error_code == expected_error_code
+    assert e.value.error_string == expected_error_string
+    assert str(e.value) == expected_error_string
+
+
 @patch("dodal.devices.robot.LOGGER")
 async def test_given_program_running_when_load_pin_then_logs_the_program_name_and_times_out(
     patch_logger: MagicMock,
@@ -27,7 +45,7 @@ async def test_given_program_running_when_load_pin_then_logs_the_program_name_an
     program_name = "BAD_PROGRAM"
     set_mock_value(device.program_running, True)
     set_mock_value(device.program_name, program_name)
-    with pytest.raises(TimeoutError):
+    with pytest.raises(RobotLoadFailed):
         await device.set(SampleLocation(0, 0))
     last_log = patch_logger.mock_calls[1].args[0]
     assert program_name in last_log
@@ -41,7 +59,7 @@ async def test_given_program_not_running_but_pin_not_unmounting_when_load_pin_th
     set_mock_value(device.program_running, False)
     set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
     device.load = AsyncMock(side_effect=device.load)
-    with pytest.raises(TimeoutError):
+    with pytest.raises(RobotLoadFailed):
         await device.set(SampleLocation(15, 10))
     device.load.trigger.assert_called_once()  # type:ignore
     last_log = patch_logger.mock_calls[1].args[0]
@@ -56,7 +74,7 @@ async def test_given_program_not_running_and_pin_unmounting_but_new_pin_not_moun
     set_mock_value(device.program_running, False)
     set_mock_value(device.gonio_pin_sensor, PinMounted.NO_PIN_MOUNTED)
     device.load = AsyncMock(side_effect=device.load)
-    with pytest.raises(TimeoutError):
+    with pytest.raises(RobotLoadFailed):
         await device.set(SampleLocation(15, 10))
     device.load.trigger.assert_called_once()  # type:ignore
     last_log = patch_logger.mock_calls[1].args[0]
@@ -72,6 +90,7 @@ async def test_given_program_not_running_and_pin_unmounts_then_mounts_when_load_
     device.load = AsyncMock(side_effect=device.load)
     status = device.set(SampleLocation(15, 10))
     await sleep(0.01)
+    device.load.trigger.assert_called_once()  # type:ignore
     set_mock_value(device.gonio_pin_sensor, PinMounted.NO_PIN_MOUNTED)
     await sleep(0.005)
     set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
@@ -80,3 +99,22 @@ async def test_given_program_not_running_and_pin_unmounts_then_mounts_when_load_
     assert (await device.next_puck.get_value()) == 15
     assert (await device.next_pin.get_value()) == 10
     device.load.trigger.assert_called_once()  # type:ignore
+
+
+async def test_given_waiting_for_pin_to_mount_when_no_pin_mounted_then_error_raised():
+    device = await _get_bart_robot()
+    status = create_task(device.pin_mounted_or_no_pin_found())
+    await sleep(0.2)
+    set_mock_value(device.error_code, 25)
+    await sleep(0.01)
+    with pytest.raises(RobotLoadFailed):
+        await status
+
+
+async def test_given_waiting_for_pin_to_mount_when_pin_mounted_then_no_error_raised():
+    device = await _get_bart_robot()
+    status = create_task(device.pin_mounted_or_no_pin_found())
+    await sleep(0.01)
+    set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
+    await sleep(0.01)
+    await status