From 62e47db2bd226dd6fcee11e4b2284248208a2fae Mon Sep 17 00:00:00 2001
From: Dominic Oram <dominic.oram@diamond.ac.uk>
Date: Wed, 18 Sep 2024 11:06:09 +0100
Subject: [PATCH] Move transmission change in rotation to after detector
 unstaged

---
 .../experiment_plans/rotation_scan_plan.py    | 10 ++++----
 .../test_rotation_scan_plan.py                | 25 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py b/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py
index 98a6a27d3..54c02f8fc 100644
--- a/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py
+++ b/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py
@@ -205,11 +205,6 @@ def rotation_scan_plan(
             "scan_points": [params.scan_points],
         }
     )
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite.xbpm_feedback,
-        composite.attenuator,
-        params.transmission_frac,
-    )
     def _rotation_scan_plan(
         motion_values: RotationMotionProfile,
         composite: RotationScanComposite,
@@ -360,6 +355,11 @@ def rotation_scan(
             ],
         }
     )
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite.xbpm_feedback,
+        composite.attenuator,
+        parameters.transmission_frac,
+    )
     def rotation_scan_plan_with_stage_and_cleanup(
         params: RotationScan,
     ):
diff --git a/tests/unit_tests/hyperion/experiment_plans/test_rotation_scan_plan.py b/tests/unit_tests/hyperion/experiment_plans/test_rotation_scan_plan.py
index 88221509e..68877ac37 100644
--- a/tests/unit_tests/hyperion/experiment_plans/test_rotation_scan_plan.py
+++ b/tests/unit_tests/hyperion/experiment_plans/test_rotation_scan_plan.py
@@ -13,6 +13,7 @@
 from dodal.devices.oav.oav_parameters import OAVParameters
 from dodal.devices.smargon import Smargon
 from dodal.devices.synchrotron import SynchrotronMode
+from dodal.devices.xbpm_feedback import Pause
 from dodal.devices.zebra import PC_GATE, Zebra
 from dodal.devices.zebra_controlled_shutter import ZebraShutterControl
 from ophyd_async.core import get_mock_put
@@ -384,7 +385,7 @@ def test_rotation_scan_triggers_xbpm_then_pauses_xbpm_and_sets_transmission(
         msgs,
         lambda msg: msg.command == "set"
         and msg.obj.name == "xbpm_feedback-pause_feedback"
-        and msg.args[0] == "Paused",
+        and msg.args[0] == Pause.PAUSE.value,
     )
     msgs = assert_message_and_return_remaining(
         msgs,
@@ -394,6 +395,28 @@ def test_rotation_scan_triggers_xbpm_then_pauses_xbpm_and_sets_transmission(
     )
 
 
+def test_rotation_scan_does_not_change_transmission_back_until_after_data_collected(
+    rotation_scan_simulated_messages,
+    test_rotation_params: RotationScan,
+):
+    msgs = assert_message_and_return_remaining(
+        rotation_scan_simulated_messages,
+        lambda msg: msg.command == "unstage" and msg.obj.name == "eiger",
+    )
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "set"
+        and msg.obj.name == "xbpm_feedback-pause_feedback"
+        and msg.args[0] == Pause.RUN.value,
+    )
+    msgs = assert_message_and_return_remaining(
+        msgs,
+        lambda msg: msg.command == "set"
+        and msg.obj.name == "attenuator"
+        and msg.args[0] == 1.0,
+    )
+
+
 def test_rotation_scan_moves_gonio_to_start_before_snapshots(
     rotation_scan_simulated_messages,
 ):