diff --git a/src/dodal/devices/robot.py b/src/dodal/devices/robot.py index eff6171709..016573731e 100644 --- a/src/dodal/devices/robot.py +++ b/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) diff --git a/tests/devices/unit_tests/test_bart_robot.py b/tests/devices/unit_tests/test_bart_robot.py index ede43c5a7a..adbdd418c8 100644 --- a/tests/devices/unit_tests/test_bart_robot.py +++ b/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