Wait for smargon enable in robot load and unload (#1791)

* Wait for smargon enable rather than pin mounted in robot load

* Fix tests

* Small fixes

* PR comments

* Fix lint

---------

Co-authored-by: Dominic Oram <dominic.oram@diamond.ac.uk>

Author: jacob720

src/dodal/devices/robot.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import FIRST_COMPLETED, CancelledError, Task, wait_for
 from dataclasses import dataclass
+from enum import IntEnum
 
 from bluesky.protocols import Movable
 from ophyd_async.core import (
@@ -21,8 +22,8 @@
 
 from dodal.log import LOGGER
 
-WAIT_FOR_OLD_PIN_MSG = "Waiting on old pin unloaded"
-WAIT_FOR_NEW_PIN_MSG = "Waiting on new pin loaded"
+WAIT_FOR_BEAMLINE_DISABLE_MSG = "Waiting on beamline disable"
+WAIT_FOR_BEAMLINE_ENABLE_MSG = "Waiting on beamline enable"
 
 
 class RobotLoadError(Exception):
@@ -51,6 +52,11 @@ class PinMounted(StrictEnum):
     PIN_MOUNTED = "Pin Mounted"
 
 
+class BeamlineStatus(IntEnum):
+    ENABLED = 0
+    DISABLED = 1
+
+
 class ErrorStatus(Device):
     def __init__(self, prefix: str) -> None:
         self.str = epics_signal_r(str, prefix + "_ERR_MSG")
@@ -89,6 +95,8 @@ def __init__(self, prefix: str, name: str = "") -> None:
             self.current_puck = epics_signal_r(float, prefix + "CURRENT_PUCK_RBV")
             self.current_pin = epics_signal_r(float, prefix + "CURRENT_PIN_RBV")
 
+        self.beamline_disabled = epics_signal_r(int, prefix + "ROBOT_OP_16_BITS.B8")
+
         self.next_pin = epics_signal_rw_rbv(float, prefix + "NEXT_PIN")
         self.next_puck = epics_signal_rw_rbv(float, prefix + "NEXT_PUCK")
 
@@ -119,8 +127,8 @@ def __init__(self, prefix: str, name: str = "") -> None:
         )
         super().__init__(name=name)
 
-    async def pin_state_or_error(self, expected_state=PinMounted.PIN_MOUNTED):
-        """This co-routine will finish when either the pin sensor reaches the specified
+    async def beamline_status_or_error(self, expected_state: BeamlineStatus):
+        """This co-routine will finish when either the beamline reaches the specified
         state or the robot gives an error (whichever happens first). In the case where
         there is an error a RobotLoadError error is raised.
         """
@@ -133,12 +141,12 @@ async def raise_if_error():
             error_msg = await self.prog_error.str.get_value()
             raise RobotLoadError(error_code, error_msg)
 
-        async def wfv():
-            await wait_for_value(self.gonio_pin_sensor, expected_state, None)
+        async def wait_for_expected_state():
+            await wait_for_value(self.beamline_disabled, expected_state.value, None)
 
         tasks = [
             (Task(raise_if_error())),
-            (Task(wfv())),
+            (Task(wait_for_expected_state())),
         ]
         try:
             finished, unfinished = await asyncio.wait(
@@ -174,12 +182,15 @@ async def _load_pin_and_puck(self, sample_location: SampleLocation):
             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(WAIT_FOR_OLD_PIN_MSG)
-            await self.pin_state_or_error(PinMounted.NO_PIN_MOUNTED)
-        LOGGER.info(WAIT_FOR_NEW_PIN_MSG)
+        await self._wait_for_beamline_enabled_after_load_or_unload()
+
+    async def _wait_for_beamline_enabled_after_load_or_unload(self):
+        if await self.beamline_disabled.get_value() == BeamlineStatus.ENABLED.value:
+            LOGGER.info(WAIT_FOR_BEAMLINE_DISABLE_MSG)
+            await self.beamline_status_or_error(BeamlineStatus.DISABLED)
 
-        await self.pin_state_or_error()
+        LOGGER.info(WAIT_FOR_BEAMLINE_ENABLE_MSG)
+        await self.beamline_status_or_error(BeamlineStatus.ENABLED)
 
     @AsyncStatus.wrap
     async def set(self, value: SampleLocation):
@@ -198,7 +209,10 @@ async def set(self, value: SampleLocation):
                 )
             else:
                 await self.unload.trigger(timeout=self.LOAD_TIMEOUT)
-                await wait_for_value(self.program_running, False, self.NOT_BUSY_TIMEOUT)
+                await wait_for(
+                    self._wait_for_beamline_enabled_after_load_or_unload(),
+                    timeout=self.LOAD_TIMEOUT + 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

@@ -15,9 +15,10 @@
 
 from dodal.devices.robot import (
     SAMPLE_LOCATION_EMPTY,
-    WAIT_FOR_NEW_PIN_MSG,
-    WAIT_FOR_OLD_PIN_MSG,
+    WAIT_FOR_BEAMLINE_DISABLE_MSG,
+    WAIT_FOR_BEAMLINE_ENABLE_MSG,
     BartRobot,
+    BeamlineStatus,
     PinMounted,
     RobotLoadError,
     SampleLocation,
@@ -36,14 +37,18 @@ async def robot_for_unload():
 
     async def finish_later():
         await drying_complete.wait()
+        await asyncio.sleep(0.1)
         set_mock_value(device.program_running, False)
+        set_mock_value(device.beamline_disabled, BeamlineStatus.ENABLED.value)
 
     async def fake_unload(*args, **kwargs):
         set_mock_value(device.program_running, True)
+        set_mock_value(device.beamline_disabled, BeamlineStatus.DISABLED.value)
         await trigger_complete.wait()
         asyncio.create_task(finish_later())
 
     get_mock_put(device.unload).side_effect = fake_unload
+    # device.unload.trigger = AsyncMock(side_effect=fake_unload)
     return device, trigger_complete, drying_complete
 
 
@@ -52,6 +57,7 @@ async def _get_bart_robot() -> BartRobot:
     device.LOAD_TIMEOUT = 1  # type: ignore
     device.NOT_BUSY_TIMEOUT = 1  # type: ignore
     await device.connect(mock=True)
+    # set_mock_value(device.beamline_disabled, BeamlineStatus.DISABLED.value)
     return device
 
 
@@ -109,7 +115,7 @@ async def test_given_program_not_running_but_pin_not_unmounting_when_load_pin_th
         await device.set(SampleLocation(15, 10))
     device.load.trigger.assert_called_once()  # type:ignore
     last_log = patch_logger.mock_calls[1].args[0]
-    assert "Waiting on old pin unloaded" in last_log
+    assert WAIT_FOR_BEAMLINE_DISABLE_MSG in last_log
 
 
 @patch("dodal.devices.robot.LOGGER")
@@ -127,34 +133,34 @@ async def test_given_program_not_running_and_pin_unmounting_but_new_pin_not_moun
     try:
         device.load.trigger.assert_called_once()  # type:ignore
         last_log = patch_logger.mock_calls[1].args[0]
-        assert "Waiting on new pin loaded" in last_log
+        assert WAIT_FOR_BEAMLINE_DISABLE_MSG in last_log
     except AssertionError:
         traceback.print_exception(exc_info.value)
         raise
 
 
-def _set_pin_sensor_on_log_messages(device: BartRobot, msg: str):
-    if msg == WAIT_FOR_OLD_PIN_MSG:
-        set_mock_value(device.gonio_pin_sensor, PinMounted.NO_PIN_MOUNTED)
-    elif msg == WAIT_FOR_NEW_PIN_MSG:
-        set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
+def _set_beamline_enabled_on_log_messages(device: BartRobot, msg: str):
+    if msg == WAIT_FOR_BEAMLINE_DISABLE_MSG:
+        set_mock_value(device.beamline_disabled, BeamlineStatus.DISABLED.value)
+    elif msg == WAIT_FOR_BEAMLINE_ENABLE_MSG:
+        set_mock_value(device.beamline_disabled, BeamlineStatus.ENABLED.value)
 
 
 def _error_on_unload_log_messages(device: BartRobot, msg: str):
-    if msg == WAIT_FOR_OLD_PIN_MSG:
+    if msg == WAIT_FOR_BEAMLINE_DISABLE_MSG:
         set_mock_value(device.prog_error.code, 40)
         set_mock_value(device.prog_error.str, "Test error")
 
 
-# Use log info messages to determine when to set the gonio_pin_sensor, so we don't have to use any sleeps during testing
+# Use log info messages to determine when to set the beamline enable, so we don't have to use any sleeps during testing
 async def set_with_happy_path(
     device: BartRobot, mock_log_info: MagicMock
 ) -> AsyncStatus:
     """Mocks the logic that the robot would do on a successful load"""
 
-    mock_log_info.side_effect = partial(_set_pin_sensor_on_log_messages, device)
+    mock_log_info.side_effect = partial(_set_beamline_enabled_on_log_messages, device)
     set_mock_value(device.program_running, False)
-    set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
+    set_mock_value(device.beamline_disabled, BeamlineStatus.ENABLED.value)
     status = device.set(SampleLocation(15, 10))
     return status
 
@@ -166,7 +172,7 @@ async def set_with_unhappy_path(
 
     mock_log_info.side_effect = partial(_error_on_unload_log_messages, device)
     set_mock_value(device.program_running, False)
-    set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
+    set_mock_value(device.beamline_disabled, BeamlineStatus.ENABLED.value)
     status = device.set(SampleLocation(15, 10))
     return status
 
@@ -187,15 +193,16 @@ async def test_given_program_not_running_and_pin_unmounts_then_mounts_when_load_
 async def test_given_waiting_for_pin_to_mount_when_no_pin_mounted_then_error_raised():
     device = await _get_bart_robot()
     set_mock_value(device.prog_error.code, 25)
-    status = device.pin_state_or_error()
+    set_mock_value(device.beamline_disabled, BeamlineStatus.DISABLED.value)
+    status = device.beamline_status_or_error(BeamlineStatus.ENABLED)
     with pytest.raises(RobotLoadError):
         await status
 
 
-async def test_given_waiting_for_pin_to_mount_when_pin_mounted_then_no_error_raised():
+async def test_given_waiting_for_beamline_to_enable_when_beamline_enabled_then_no_error_raised():
     device = await _get_bart_robot()
-    status = create_task(device.pin_state_or_error())
-    set_mock_value(device.gonio_pin_sensor, PinMounted.PIN_MOUNTED)
+    status = create_task(device.beamline_status_or_error(BeamlineStatus.ENABLED))
+    set_mock_value(device.beamline_disabled, BeamlineStatus.ENABLED.value)
     await status