Merge branch 'main' into 488_warn_rather_than_error_on_webcam_failure

src/mx_bluesky/hyperion/experiment_plans/robot_load_then_centre_plan.py

@@ -238,53 +238,73 @@ def robot_load_and_snapshots(
     yield from bps.wait("reset-lower_gonio")
 
 
-def robot_load_then_centre_plan(
+def centring_plan_from_robot_load_params(
     composite: RobotLoadThenCentreComposite,
     params: RobotLoadThenCentre,
 ):
-    # TODO: get these from one source of truth #254
-    assert params.sample_puck is not None
-    assert params.sample_pin is not None
-
-    if not (
-        yield from _pin_already_loaded(
-            composite.robot, params.sample_pin, params.sample_puck
-        )
-    ):
-        yield from prepare_for_robot_load(composite)
-        yield from bpp.run_wrapper(
-            robot_load_and_snapshots(
-                composite, params, SampleLocation(params.sample_puck, params.sample_pin)
-            ),
-            md={
-                "subplan_name": CONST.PLAN.ROBOT_LOAD,
-                "metadata": {
-                    "visit_path": str(params.visit_directory),
-                    "sample_id": params.sample_id,
-                    "sample_puck": params.sample_puck,
-                    "sample_pin": params.sample_pin,
-                },
-                "activate_callbacks": [
-                    "RobotLoadISPyBCallback",
-                ],
-            },
-        )
-    else:
-        LOGGER.info(
-            f"Pin/puck {params.sample_pin}/{params.sample_puck} already loaded, will not reload."
-        )
     yield from pin_centre_then_xray_centre_plan(
         cast(GridDetectThenXRayCentreComposite, composite),
         params.pin_centre_then_xray_centre_params(),
     )
 
 
+def robot_load_then_centre_plan(
+    composite: RobotLoadThenCentreComposite,
+    params: RobotLoadThenCentre,
+    sample_location: SampleLocation,
+):
+    yield from prepare_for_robot_load(composite)
+    yield from bpp.run_wrapper(
+        robot_load_and_snapshots(composite, params, sample_location),
+        md={
+            "subplan_name": CONST.PLAN.ROBOT_LOAD,
+            "metadata": {
+                "visit_path": str(params.visit_directory),
+                "sample_id": params.sample_id,
+                "sample_puck": params.sample_puck,
+                "sample_pin": params.sample_pin,
+            },
+            "activate_callbacks": [
+                "RobotLoadISPyBCallback",
+            ],
+        },
+    )
+
+    yield from centring_plan_from_robot_load_params(composite, params)
+
+
 def robot_load_then_centre(
     composite: RobotLoadThenCentreComposite,
     parameters: RobotLoadThenCentre,
 ) -> MsgGenerator:
     eiger: EigerDetector = composite.eiger
 
+    # TODO: get these from one source of truth #254
+    assert parameters.sample_puck is not None
+    assert parameters.sample_pin is not None
+
+    doing_sample_load = not (
+        yield from _pin_already_loaded(
+            composite.robot, parameters.sample_pin, parameters.sample_puck
+        )
+    )
+
+    doing_chi_change = parameters.chi_start_deg is not None
+
+    if doing_sample_load:
+        plan = robot_load_then_centre_plan(
+            composite,
+            parameters,
+            SampleLocation(parameters.sample_puck, parameters.sample_pin),
+        )
+        LOGGER.info("Pin not loaded, loading and centring")
+    elif doing_chi_change:
+        plan = centring_plan_from_robot_load_params(composite, parameters)
+        LOGGER.info("Pin already loaded but chi changed so centring")
+    else:
+        LOGGER.info("Pin already loaded and chi not changed so doing nothing")
+        return
+
     detector_params = parameters.detector_params
     if not detector_params.expected_energy_ev:
         actual_energy_ev = 1000 * (
@@ -297,6 +317,6 @@ def robot_load_then_centre(
         eiger,
         composite.detector_motion,
         parameters.detector_distance_mm,
-        robot_load_then_centre_plan(composite, parameters),
+        plan,
         group=CONST.WAIT.GRID_READY_FOR_DC,
     )

tests/conftest.py

@@ -4,7 +4,7 @@
 import logging
 import sys
 import threading
-from collections.abc import Generator, Sequence
+from collections.abc import Callable, Generator, Sequence
 from functools import partial
 from typing import Any
 from unittest.mock import MagicMock, patch
@@ -896,3 +896,10 @@ def feature_flags():
     return FeatureFlags(
         **{field_name: False for field_name in FeatureFlags.model_fields.keys()}
     )
+
+
+def assert_none_matching(
+    messages: list[Msg],
+    predicate: Callable[[Msg], bool],
+):
+    assert not list(filter(predicate, messages))

tests/unit_tests/hyperion/experiment_plans/test_wait_for_robot_load_then_centre.py

@@ -8,6 +8,7 @@
 from bluesky.utils import Msg
 from dodal.devices.aperturescatterguard import ApertureValue
 from dodal.devices.oav.oav_detector import OAV
+from dodal.devices.robot import SampleLocation
 from dodal.devices.smargon import StubPosition
 from dodal.devices.webcam import Webcam
 from ophyd.sim import NullStatus
@@ -20,7 +21,6 @@
     RobotLoadThenCentreComposite,
     prepare_for_robot_load,
     robot_load_then_centre,
-    robot_load_then_centre_plan,
     take_robot_snapshots,
 )
 from mx_bluesky.hyperion.external_interaction.callbacks.robot_load.ispyb_callback import (
@@ -31,7 +31,7 @@
     RobotLoadThenCentre,
 )
 
-from ....conftest import raw_params_from_file
+from ....conftest import assert_none_matching, raw_params_from_file
 
 
 @pytest.fixture
@@ -546,33 +546,185 @@ def test_when_plan_run_then_thawing_turned_on_for_expected_time(
     )
 
 
+def mock_current_sample(sim_run_engine: RunEngineSimulator, sample: SampleLocation):
+    sim_run_engine.add_handler(
+        "read",
+        lambda msg: {"robot-current_puck": {"value": sample.puck}},
+        "robot-current_puck",
+    )
+    sim_run_engine.add_handler(
+        "read",
+        lambda msg: {"robot-current_pin": {"value": sample.pin}},
+        "robot-current_pin",
+    )
+
+
 @patch(
-    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan"
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan",
+    MagicMock(return_value=iter([Msg("centre_plan")])),
 )
-@patch("mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.do_robot_load")
 @patch(
-    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.prepare_for_robot_load"
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.set_energy_plan",
+    MagicMock(return_value=iter([])),
 )
-def test_given_sample_already_loaded_when_plan_run_then_sample_not_loaded(
-    mock_prepare_robot: MagicMock,
-    mock_do_robot_load: MagicMock,
-    mock_centring_plan: MagicMock,
+def test_given_sample_already_loaded_and_chi_not_changed_when_robot_load_called_then_eiger_not_staged_and_centring_not_run(
     robot_load_composite: RobotLoadThenCentreComposite,
-    robot_load_then_centre_params_no_energy: RobotLoadThenCentre,
-    RE: RunEngine,
+    robot_load_then_centre_params: RobotLoadThenCentre,
+    sim_run_engine: RunEngineSimulator,
 ):
-    set_mock_value(robot_load_composite.robot.current_pin, 1)
-    set_mock_value(robot_load_composite.robot.current_puck, 2)
+    sample_location = SampleLocation(2, 6)
+    robot_load_then_centre_params.sample_puck = sample_location.puck
+    robot_load_then_centre_params.sample_pin = sample_location.pin
+    robot_load_then_centre_params.chi_start_deg = None
 
-    robot_load_then_centre_params_no_energy.sample_pin = 1
-    robot_load_then_centre_params_no_energy.sample_puck = 2
+    mock_current_sample(sim_run_engine, sample_location)
 
-    RE(
-        robot_load_then_centre_plan(
+    messages = sim_run_engine.simulate_plan(
+        robot_load_then_centre(
             robot_load_composite,
-            robot_load_then_centre_params_no_energy,
+            robot_load_then_centre_params,
+        )
+    )
+
+    assert_none_matching(
+        messages, lambda msg: msg.command == "set" and msg.obj.name == "eiger_do_arm"
+    )
+
+    assert_none_matching(
+        messages, lambda msg: msg.command == "set" and msg.obj.name == "robot"
+    )
+
+    assert_none_matching(
+        messages,
+        lambda msg: msg.command == "centre_plan",
+    )
+
+
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan",
+    MagicMock(return_value=iter([Msg("centre_plan")])),
+)
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.set_energy_plan",
+    MagicMock(return_value=iter([])),
+)
+def test_given_sample_already_loaded_and_chi_is_changed_when_robot_load_called_then_eiger_staged_and_centring_run(
+    robot_load_composite: RobotLoadThenCentreComposite,
+    robot_load_then_centre_params: RobotLoadThenCentre,
+    sim_run_engine: RunEngineSimulator,
+):
+    sample_location = SampleLocation(2, 6)
+    robot_load_then_centre_params.sample_puck = sample_location.puck
+    robot_load_then_centre_params.sample_pin = sample_location.pin
+    robot_load_then_centre_params.chi_start_deg = 30
+
+    mock_current_sample(sim_run_engine, sample_location)
+
+    messages = sim_run_engine.simulate_plan(
+        robot_load_then_centre(
+            robot_load_composite,
+            robot_load_then_centre_params,
         )
     )
-    mock_prepare_robot.assert_not_called()
-    mock_do_robot_load.assert_not_called()
-    mock_centring_plan.assert_called()
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "set" and msg.obj.name == "eiger_do_arm",
+    )
+
+    assert_none_matching(
+        messages, lambda msg: msg.command == "set" and msg.obj.name == "robot"
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "centre_plan",
+    )
+
+
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan",
+    MagicMock(return_value=iter([Msg("centre_plan")])),
+)
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.set_energy_plan",
+    MagicMock(return_value=iter([])),
+)
+def test_given_sample_not_loaded_and_chi_not_changed_when_robot_load_called_then_eiger_staged_before_robot_and_centring_run_after(
+    robot_load_composite: RobotLoadThenCentreComposite,
+    robot_load_then_centre_params: RobotLoadThenCentre,
+    sim_run_engine: RunEngineSimulator,
+):
+    robot_load_then_centre_params.sample_puck = (puck := 2)
+    robot_load_then_centre_params.sample_pin = (pin := 6)
+    robot_load_then_centre_params.chi_start_deg = None
+
+    mock_current_sample(sim_run_engine, SampleLocation(1, 1))
+
+    messages = sim_run_engine.simulate_plan(
+        robot_load_then_centre(
+            robot_load_composite,
+            robot_load_then_centre_params,
+        )
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "set" and msg.obj.name == "eiger_do_arm",
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "set"
+        and msg.obj.name == "robot"
+        and msg.args[0] == SampleLocation(puck, pin),
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "centre_plan",
+    )
+
+
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan",
+    MagicMock(return_value=iter([Msg("centre_plan")])),
+)
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan.set_energy_plan",
+    MagicMock(return_value=iter([])),
+)
+def test_given_sample_not_loaded_and_chi_changed_when_robot_load_called_then_eiger_staged_before_robot_and_centring_run(
+    robot_load_composite: RobotLoadThenCentreComposite,
+    robot_load_then_centre_params: RobotLoadThenCentre,
+    sim_run_engine: RunEngineSimulator,
+):
+    robot_load_then_centre_params.sample_puck = (puck := 2)
+    robot_load_then_centre_params.sample_pin = (pin := 6)
+    robot_load_then_centre_params.chi_start_deg = 30
+
+    mock_current_sample(sim_run_engine, SampleLocation(1, 1))
+
+    messages = sim_run_engine.simulate_plan(
+        robot_load_then_centre(
+            robot_load_composite,
+            robot_load_then_centre_params,
+        )
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "set" and msg.obj.name == "eiger_do_arm",
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "set"
+        and msg.obj.name == "robot"
+        and msg.args[0] == SampleLocation(puck, pin),
+    )
+
+    messages = assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "centre_plan",
+    )