From 2d47e50304325313f5b8b44d73abde7843e77f1c Mon Sep 17 00:00:00 2001
From: Dominic Oram <dominic.oram@diamond.ac.uk>
Date: Thu, 15 Aug 2024 10:01:40 +0100
Subject: [PATCH] Do not move lower gonio until the robot is safely away
 (#1526)

* Do not move lower gonio until the robot is safely away

* Added group to thawing
---
 .../robot_load_then_centre_plan.py            | 14 +++--
 .../test_wait_for_robot_load_then_centre.py   | 56 +++++++++++++++++++
 2 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/src/hyperion/experiment_plans/robot_load_then_centre_plan.py b/src/hyperion/experiment_plans/robot_load_then_centre_plan.py
index 61404c9fe..c56b6dd77 100644
--- a/src/hyperion/experiment_plans/robot_load_then_centre_plan.py
+++ b/src/hyperion/experiment_plans/robot_load_then_centre_plan.py
@@ -149,10 +149,11 @@ def prepare_for_robot_load(composite: RobotLoadThenCentreComposite):
     yield from bps.wait("prepare_robot_load")
 
 
-def robot_load_and_energy_change(
+def do_robot_load(
     composite: RobotLoadThenCentreComposite,
     sample_location: SampleLocation,
     demand_energy_ev: float | None,
+    thawing_time: float,
 ):
     yield from bps.abs_set(
         composite.robot,
@@ -168,6 +169,11 @@ def robot_load_and_energy_change(
 
     yield from bps.wait("robot_load")
 
+    yield from bps.abs_set(
+        composite.thawer.thaw_for_time_s, thawing_time, group="thawing_finished"
+    )
+    yield from wait_for_smargon_not_disabled(composite.smargon)
+
 
 def raise_exception_if_moved_out_of_cryojet(exception):
     yield from bps.null()
@@ -204,10 +210,11 @@ def robot_load_and_snapshots():
         assert params.sample_puck is not None
         assert params.sample_pin is not None
 
-        robot_load_plan = robot_load_and_energy_change(
+        robot_load_plan = do_robot_load(
             composite,
             SampleLocation(params.sample_puck, params.sample_pin),
             params.demand_energy_ev,
+            params.thawing_time,
         )
 
         # The lower gonio must be in the correct position for the robot load and we
@@ -226,9 +233,6 @@ def robot_load_and_snapshots():
             except_plan=raise_exception_if_moved_out_of_cryojet,
         )
 
-        yield from bps.abs_set(composite.thawer.thaw_for_time_s, params.thawing_time)
-        yield from wait_for_smargon_not_disabled(composite.smargon)
-
         yield from take_robot_snapshots(
             composite.oav, composite.webcam, params.snapshot_directory
         )
diff --git a/tests/unit_tests/experiment_plans/test_wait_for_robot_load_then_centre.py b/tests/unit_tests/experiment_plans/test_wait_for_robot_load_then_centre.py
index 5648c9940..66915a111 100644
--- a/tests/unit_tests/experiment_plans/test_wait_for_robot_load_then_centre.py
+++ b/tests/unit_tests/experiment_plans/test_wait_for_robot_load_then_centre.py
@@ -398,6 +398,62 @@ def get_read(axis, msg):
         )
 
 
+@patch(
+    "hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan"
+)
+@patch(
+    "hyperion.experiment_plans.robot_load_then_centre_plan.set_energy_plan",
+    MagicMock(return_value=iter([])),
+)
+def test_when_plan_run_then_lower_gonio_moved_before_robot_loads_and_back_after_smargon_enabled(
+    mock_centring_plan: MagicMock,
+    robot_load_composite: RobotLoadThenCentreComposite,
+    robot_load_then_centre_params_no_energy: RobotLoadThenCentre,
+    sim_run_engine: RunEngineSimulator,
+):
+    initial_values = {"x": 0.11, "y": 0.12, "z": 0.13}
+
+    def get_read(axis, msg):
+        return {f"lower_gonio-{axis}": {"value": initial_values[axis]}}
+
+    for axis in initial_values.keys():
+        sim_run_engine.add_handler(
+            "read", partial(get_read, axis), f"lower_gonio-{axis}"
+        )
+
+    messages = sim_run_engine.simulate_plan(
+        robot_load_then_centre(
+            robot_load_composite,
+            robot_load_then_centre_params_no_energy,
+        )
+    )
+
+    assert_message_and_return_remaining(
+        messages, lambda msg: msg.command == "set" and msg.obj.name == "robot"
+    )
+
+    for axis in initial_values.keys():
+        messages = assert_message_and_return_remaining(
+            messages,
+            lambda msg: msg.command == "set"
+            and msg.obj.name == f"lower_gonio-{axis}"
+            and msg.args == (0,),
+        )
+
+    assert_message_and_return_remaining(
+        messages,
+        lambda msg: msg.command == "read" and msg.obj.name == "smargon-disabled",
+    )
+
+    for axis, initial in initial_values.items():
+        messages = assert_message_and_return_remaining(
+            messages,
+            lambda msg: msg.command == "set"
+            and msg.obj.name == f"lower_gonio-{axis}"
+            and msg.args == (initial,),
+        )
+
+
 @patch(
     "hyperion.experiment_plans.robot_load_then_centre_plan.pin_centre_then_xray_centre_plan"
 )