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" )