Implement robot unload for udc default state (#1607)

* Implement robot unload for udc default state

Author: rtuck99

src/mx_bluesky/common/device_setup_plans/robot_load_unload.py

@@ -73,7 +73,7 @@ def robot_unload(
     smargon: Smargon,
     aperture_scatterguard: ApertureScatterguard,
     lower_gonio: XYZStage,
-    visit: str,
+    visit: str | None,
 ):
     """Unloads the currently mounted pin into the location that it was loaded from. The
     loaded location is stored on the robot and so need not be provided.

src/mx_bluesky/hyperion/experiment_plans/udc_default_state.py

@@ -4,7 +4,7 @@
 from dodal.common.beamlines.beamline_parameters import (
     get_beamline_parameters,
 )
-from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
 from dodal.devices.collimation_table import CollimationTable
 from dodal.devices.cryostream import (
     CryoStreamGantry,
@@ -16,6 +16,7 @@
 from dodal.devices.fluorescence_detector_motion import FluorescenceDetector
 from dodal.devices.fluorescence_detector_motion import InOut as FlouInOut
 from dodal.devices.hutch_shutter import HutchShutter, ShutterDemand
+from dodal.devices.motors import XYZStage
 from dodal.devices.mx_phase1.beamstop import BeamstopPositions
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.robot import BartRobot, PinMounted
@@ -24,6 +25,7 @@
 from dodal.devices.smargon import Smargon
 from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutterState
 
+from mx_bluesky.common.device_setup_plans.robot_load_unload import robot_unload
 from mx_bluesky.common.experiment_plans.beamstop_check import (
     BeamstopCheckDevices,
     move_beamstop_in_and_verify_using_diode,
@@ -47,6 +49,7 @@ class UDCDefaultDevices(BeamstopCheckDevices):
     cryostream_gantry: CryoStreamGantry
     fluorescence_det_motion: FluorescenceDetector
     hutch_shutter: HutchShutter
+    lower_gonio: XYZStage
     robot: BartRobot
     scintillator: Scintillator
     gonio: Smargon
@@ -65,8 +68,6 @@ def move_to_udc_default_state(devices: UDCDefaultDevices):
 
     yield from _check_cryostream(devices)
 
-    yield from _verify_no_sample_present(devices.robot)
-
     # Close fast shutter before opening hutch shutter
     yield from bps.abs_set(devices.sample_shutter, ZebraShutterState.CLOSE, wait=True)
 
@@ -108,6 +109,13 @@ def move_to_udc_default_state(devices: UDCDefaultDevices):
     # Wait for all of the above to complete
     yield from bps.wait(group=_GROUP_PRE_BEAMSTOP_CHECK, timeout=10)
 
+    yield from _unload_sample_if_present(
+        devices.robot,
+        devices.gonio,
+        devices.aperture_scatterguard,
+        devices.lower_gonio,
+    )
+
     feature_flags: HyperionFeatureSettings = (
         get_hyperion_config_client().get_feature_flags()
     )
@@ -175,11 +183,16 @@ def _verify_correct_cryostream_selected(
         )
 
 
-def _verify_no_sample_present(robot: BartRobot):
+def _unload_sample_if_present(
+    robot: BartRobot,
+    smargon: Smargon,
+    aperture_scatterguard: ApertureScatterguard,
+    lower_gonio: XYZStage,
+):
     pin_mounted = yield from bps.rd(robot.gonio_pin_sensor)
 
     if pin_mounted != PinMounted.NO_PIN_MOUNTED:
-        # Cannot unload this sample because we do not know the correct visit for it
-        raise UnexpectedSampleError(
-            "An unexpected sample was found, please unload the sample manually."
+        LOGGER.info("Pin detected, unloading sample...")
+        yield from robot_unload(
+            robot, smargon, aperture_scatterguard, lower_gonio, None
         )

src/mx_bluesky/hyperion/external_interaction/callbacks/robot_actions/ispyb_callback.py

@@ -51,20 +51,24 @@ def activity_gated_start(self, doc: RunStart):
             ISPYB_ZOCALO_CALLBACK_LOGGER.debug(
                 f"ISPyB robot load callback received: {doc}"
             )
-            self.run_uid = doc.get("uid")
             metadata = doc.get("metadata")
             assert isinstance(metadata, dict)
-            self._sample_id = metadata["sample_id"]
-            assert isinstance(self._sample_id, int)
-            proposal, session = get_proposal_and_session_from_visit_string(
-                metadata["visit"]
-            )
-            self.action_id = self.expeye.start_robot_action(
-                "LOAD" if subplan == CONST.PLAN.ROBOT_LOAD else "UNLOAD",
-                proposal,
-                session,
-                self._sample_id,
-            )
+            self._sample_id = metadata.get("sample_id")
+            if not isinstance(self._sample_id, int):
+                ISPYB_ZOCALO_CALLBACK_LOGGER.info(
+                    "No sample ID provided, not recording robot action in ispyb."
+                )
+            else:
+                self.run_uid = doc.get("uid")
+                proposal, session = get_proposal_and_session_from_visit_string(
+                    metadata["visit"]
+                )
+                self.action_id = self.expeye.start_robot_action(
+                    "LOAD" if subplan == CONST.PLAN.ROBOT_LOAD else "UNLOAD",
+                    proposal,
+                    session,
+                    self._sample_id,
+                )
         return super().activity_gated_start(doc)
 
     def activity_gated_descriptor(self, doc: EventDescriptor) -> EventDescriptor | None:
@@ -73,7 +77,11 @@ def activity_gated_descriptor(self, doc: EventDescriptor) -> EventDescriptor | N
 
     def activity_gated_event(self, doc: Event) -> Event | None:
         event_descriptor = self.descriptors.get(doc["descriptor"])
-        if (
+        if self._sample_id is None:
+            ISPYB_ZOCALO_CALLBACK_LOGGER.info(
+                "Ignoring robot update event as no sample ID provided."
+            )
+        elif (
             event_descriptor
             and event_descriptor.get("name") == CONST.DESCRIPTORS.ROBOT_UPDATE
         ):

tests/unit_tests/hyperion/experiment_plans/test_udc_default_state.py

@@ -17,6 +17,7 @@
 from dodal.devices.fluorescence_detector_motion import FluorescenceDetector
 from dodal.devices.fluorescence_detector_motion import InOut as FlouInOut
 from dodal.devices.hutch_shutter import HutchShutter, ShutterDemand
+from dodal.devices.motors import XYZStage
 from dodal.devices.mx_phase1.beamstop import BeamstopPositions
 from dodal.devices.robot import BartRobot, PinMounted
 from dodal.devices.scintillator import InOut, Scintillator
@@ -27,7 +28,6 @@
 from mx_bluesky.hyperion.experiment_plans.udc_default_state import (
     CryoStreamError,
     UDCDefaultDevices,
-    UnexpectedSampleError,
     move_to_udc_default_state,
 )
 from mx_bluesky.hyperion.parameters.constants import CONST, HyperionFeatureSettings
@@ -64,6 +64,7 @@ async def default_devices(
         hutch_shutter = HutchShutter("")
         scintillator = Scintillator("", MagicMock(), MagicMock(), name="scin")
         collimation_table = CollimationTable("")
+        lower_gonio = XYZStage("")
 
     with patch("dodal.devices.hutch_shutter.TEST_MODE", True):
         devices = UDCDefaultDevices(
@@ -73,6 +74,7 @@ async def default_devices(
             cryostream_gantry=cryostream_gantry,
             fluorescence_det_motion=fluo,
             hutch_shutter=hutch_shutter,
+            lower_gonio=lower_gonio,
             robot=robot,
             scintillator=scintillator,
             gonio=smargon,
@@ -156,6 +158,10 @@ def test_udc_default_state_runs_in_real_run_engine(
     run_engine(move_to_udc_default_state(default_devices))
 
 
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.udc_default_state._unload_sample_if_present",
+    MagicMock(return_value=iter([Msg("robot_unload")])),
+)
 def test_beamstop_moved_to_data_collection_if_diode_check_not_enabled(
     sim_run_engine: RunEngineSimulator,
     default_devices: UDCDefaultDevices,
@@ -166,13 +172,14 @@ def test_beamstop_moved_to_data_collection_if_diode_check_not_enabled(
         msgs,
         lambda msg: msg.command == "wait" and msg.kwargs["group"] == pre_beamstop_group,
     )
+    assert msgs[1].command == "robot_unload"
     assert (
-        msgs[1].command == "set"
-        and msgs[1].obj is default_devices.beamstop.selected_pos
-        and msgs[1].args[0] == BeamstopPositions.DATA_COLLECTION
+        msgs[2].command == "set"
+        and msgs[2].obj is default_devices.beamstop.selected_pos
+        and msgs[2].args[0] == BeamstopPositions.DATA_COLLECTION
     )
     assert (
-        msgs[2].command == "wait" and msgs[2].kwargs["group"] == msgs[1].kwargs["group"]
+        msgs[3].command == "wait" and msgs[3].kwargs["group"] == msgs[2].kwargs["group"]
     )
 
 
@@ -306,7 +313,19 @@ def test_udc_default_state_checks_cryostream_selection(
         run_engine(move_to_udc_default_state(default_devices))
 
 
-def test_udc_default_state_checks_that_pin_not_mounted(
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.udc_default_state._verify_correct_cryostream_selected",
+    MagicMock(return_value=iter([Msg("cryostream_gantry_check")])),
+)
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.udc_default_state._check_cryostream",
+    MagicMock(return_value=iter([Msg("cryostream_check")])),
+)
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.udc_default_state.robot_unload",
+    MagicMock(return_value=iter([Msg("unload")])),
+)
+def test_udc_default_state_unloads_if_sample_present(
     default_devices, sim_run_engine, beamline_parameters
 ):
     sim_run_engine.add_read_handler_for(
@@ -316,8 +335,38 @@ def test_udc_default_state_checks_that_pin_not_mounted(
         "mx_bluesky.hyperion.experiment_plans.udc_default_state.get_beamline_parameters",
         return_value=beamline_parameters,
     ):
-        with pytest.raises(UnexpectedSampleError):
-            sim_run_engine.simulate_plan(move_to_udc_default_state(default_devices))
+        msgs = sim_run_engine.simulate_plan(move_to_udc_default_state(default_devices))
+
+    msgs = assert_message_and_return_remaining(
+        msgs, lambda msg: msg.command == "cryostream_gantry_check"
+    )
+    msgs = assert_message_and_return_remaining(
+        msgs, lambda msg: msg.command == "cryostream_check"
+    )
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "set"
+        and msg.obj == default_devices.scintillator.selected_pos
+        and msg.args[0] == InOut.OUT,
+    )
+    msgs = assert_message_and_return_remaining(
+        msgs, lambda msg: msg.command == "unload"
+    )
+
+
+def test_udc_default_state_no_unload_if_no_sample_present(
+    default_devices, sim_run_engine, beamline_parameters
+):
+    sim_run_engine.add_read_handler_for(
+        default_devices.robot.gonio_pin_sensor, PinMounted.NO_PIN_MOUNTED
+    )
+    with patch(
+        "mx_bluesky.hyperion.experiment_plans.udc_default_state.get_beamline_parameters",
+        return_value=beamline_parameters,
+    ):
+        msgs = sim_run_engine.simulate_plan(move_to_udc_default_state(default_devices))
+
+    assert not [m for m in msgs if m.command == "unload"]
 
 
 def test_default_state_closes_sample_shutter_before_open_hutch_shutter(

tests/unit_tests/hyperion/external_interaction/callbacks/robot_load/test_robot_load_ispyb_callback.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, Mock, patch
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
@@ -35,6 +35,17 @@
     ],
 }
 
+metadata_unload_no_visit = {
+    "subplan_name": CONST.PLAN.ROBOT_UNLOAD,
+    "metadata": {
+        "sample_puck": SAMPLE_PUCK,
+        "sample_pin": SAMPLE_PIN,
+    },
+    "activate_callbacks": [
+        "RobotLoadISPyBCallback",
+    ],
+}
+
 update_doc_data = {
     "sampleBarcode": "BARCODE",
     "xtalSnapshotBefore": "test_webcam_snapshot",
@@ -127,6 +138,15 @@ def unsuccessful_robot_load_plan():
     raise AssertionError("Test failure")
 
 
+@bpp.run_decorator(md=metadata_unload_no_visit)
+def robot_unload_plan(robot: BartRobot, oav: OAV, webcam: Webcam):
+    yield from bps.create(name=CONST.DESCRIPTORS.ROBOT_UPDATE)
+    yield from bps.read(robot)
+    yield from bps.read(oav.snapshot)
+    yield from bps.read(webcam)
+    yield from bps.save()
+
+
 @patch(
     "mx_bluesky.hyperion.external_interaction.callbacks.robot_actions.ispyb_callback.ExpeyeInteraction",
     autospec=True,
@@ -195,3 +215,25 @@ def test_robot_load_fails_triggers_bl_sample_status_error(
     mock_sample_handling.return_value.update_sample_status.assert_called_with(
         SAMPLE_ID, BLSampleStatus.ERROR_BEAMLINE
     )
+
+
+@patch(
+    "mx_bluesky.hyperion.external_interaction.callbacks.robot_actions.ispyb_callback.ExpeyeInteraction",
+    autospec=True,
+)
+def test_robot_unload_event_without_sample_id_and_visit_is_ignored(
+    mock_expeye: MagicMock,
+    run_engine: RunEngine,
+    robot: BartRobot,
+    oav: OAV,
+    webcam: Webcam,
+):
+    unwrapped = RobotLoadISPyBCallback()
+    callback = Mock(wraps=unwrapped)
+    run_engine.subscribe(callback)
+
+    run_engine(robot_unload_plan(robot, oav, webcam))
+    mock_expeye.return_value.start_robot_action.assert_not_called()
+    mock_expeye.return_value.update_robot_action.assert_not_called()
+    mock_expeye.return_value.end_robot_action.assert_not_called()
+    mock_expeye.return_value.update_sample_status.assert_not_called()