Pass new arguments to wrapper

src/mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py

@@ -2,7 +2,7 @@
 from bluesky.preprocessors import finalize_wrapper
 from bluesky.utils import make_decorator
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
-from dodal.devices.undulator_dcm import UndulatorDCM
+from dodal.devices.undulator import Undulator
 from dodal.devices.xbpm_feedback import Pause, XBPMFeedback
 
 from mx_bluesky.common.utils.log import LOGGER
@@ -51,7 +51,7 @@ def _unpause_xbpm_feedback_and_set_transmission_to_1(
 def transmission_and_xbpm_feedback_for_collection_wrapper(
     plan,
     energy_ev: float,
-    undulator_dcm: UndulatorDCM,
+    undulator: Undulator,
     xbpm_feedback: XBPMFeedback,
     attenuator: BinaryFilterAttenuator,
     desired_transmission_fraction: float,
@@ -69,6 +69,9 @@ def transmission_and_xbpm_feedback_for_collection_wrapper(
     mostly accounts for slow thermal drift so it is safe to assume that the beam is
     stable during a collection.
 
+    In the case of a beam dump, undulator gap may not return. Therefore, we check here
+    that the undulator gap is correct after XBPM is stable, and before collection.
+
     Args:
         plan: The plan performing the data collection
         xbpm_feedback (XBPMFeedback): The XBPM device that is responsible for keeping
@@ -81,7 +84,7 @@ def _inner_plan():
         yield from _check_and_pause_feedback(
             xbpm_feedback, attenuator, desired_transmission_fraction
         )
-        yield from bps.abs_set(undulator_dcm, energy_ev / 1000, wait=True)
+        yield from bps.abs_set(undulator, energy_ev / 1000, wait=True)
         return (yield from plan)
 
     return (

src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -155,6 +155,8 @@ def flyscan_xray_centre_no_move(
     )
     @bpp.finalize_decorator(lambda: feature_controlled.tidy_plan(composite))
     @transmission_and_xbpm_feedback_for_collection_decorator(
+        parameters.demand_energy_ev,
+        composite.undulator,
         composite.xbpm_feedback,
         composite.attenuator,
         parameters.transmission_frac,

src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py

@@ -378,6 +378,8 @@ def rotation_scan(
         }
     )
     @transmission_and_xbpm_feedback_for_collection_decorator(
+        parameters.demand_energy_ev,
+        composite.undulator,
         composite.xbpm_feedback,
         composite.attenuator,
         parameters.transmission_frac,
@@ -431,6 +433,8 @@ def multi_rotation_scan(
     )
     @bpp.stage_decorator([eiger])
     @transmission_and_xbpm_feedback_for_collection_decorator(
+        parameters.demand_energy_ev,
+        composite.undulator,
         composite.xbpm_feedback,
         composite.attenuator,
         parameters.transmission_frac,

tests/system_tests/hyperion/experiment_plans/test_fgs_plan.py

@@ -159,6 +159,8 @@ async def test_xbpm_feedback_decorator(
     # in S03
 
     @transmission_and_xbpm_feedback_for_collection_decorator(
+        params.demand_energy_ev,
+        fxc_composite.undulator,
         fxc_composite.xbpm_feedback,
         fxc_composite.attenuator,
         params.transmission_frac,

tests/unit_tests/hyperion/device_setup_plans/test_xbpm_feedback.py

@@ -16,12 +16,13 @@
 async def test_given_xpbm_checks_pass_when_plan_run_with_decorator_then_run_as_expected(
     RE,
     xbpm_feedback,
+    undulator,
     attenuator,
 ):
     expected_transmission = 0.3
 
     @transmission_and_xbpm_feedback_for_collection_decorator(
-        xbpm_feedback, attenuator, expected_transmission
+        100, undulator, xbpm_feedback, attenuator, expected_transmission
     )
     def my_collection_plan():
         read_transmission = yield from bps.rd(attenuator.actual_transmission)
@@ -41,12 +42,13 @@ def my_collection_plan():
 async def test_given_xbpm_checks_fail_when_plan_run_with_decorator_then_plan_not_run(
     RE,
     xbpm_feedback,
+    undulator,
     attenuator,
 ):
     mock = MagicMock()
 
     @transmission_and_xbpm_feedback_for_collection_decorator(
-        xbpm_feedback, attenuator, 0.1
+        100, undulator, xbpm_feedback, attenuator, 0.1
     )
     def my_collection_plan():
         mock()
@@ -68,6 +70,7 @@ def my_collection_plan():
 async def test_given_xpbm_checks_pass_and_plan_fails_when_plan_run_with_decorator_then_cleaned_up(
     RE,
     xbpm_feedback,
+    undulator,
     attenuator,
 ):
     set_mock_value(xbpm_feedback.pos_stable, True)  # type: ignore
@@ -76,7 +79,7 @@ class MyException(Exception):
         pass
 
     @transmission_and_xbpm_feedback_for_collection_decorator(
-        xbpm_feedback, attenuator, 0.1
+        100, undulator, xbpm_feedback, attenuator, 0.1
     )
     def my_collection_plan():
         yield from bps.null()

tests/unit_tests/hyperion/experiment_plans/test_rotation_scan_plan.py

@@ -179,7 +179,6 @@ def test_rotation_scan(
 ):
     composite = fake_create_rotation_devices
     RE(rotation_scan(composite, test_rotation_params, oav_parameters_for_rotation))
-
     composite.eiger.do_arm.set.assert_called()  # type: ignore
     composite.eiger.unstage.assert_called()  # type: ignore