diff --git a/src/mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py b/src/mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py index 0a183e5f2..14f94ed12 100644 --- a/src/mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py +++ b/src/mx_bluesky/hyperion/device_setup_plans/xbpm_feedback.py @@ -50,10 +50,10 @@ def _unpause_xbpm_feedback_and_set_transmission_to_1( def transmission_and_xbpm_feedback_for_collection_wrapper( plan, - energy_ev: float | None, undulator: Undulator, xbpm_feedback: XBPMFeedback, attenuator: BinaryFilterAttenuator, + energy_ev: float | None, desired_transmission_fraction: float, ): """Sets the transmission for the data collection, ensuring the xbpm feedback is valid diff --git a/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py b/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py index 62bf9416c..8102d3f4b 100755 --- a/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py +++ b/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py @@ -155,10 +155,10 @@ 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.demand_energy_ev, parameters.transmission_frac, ) def run_gridscan_and_fetch_and_tidy( 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 9c3ac91fa..d0134f362 100644 --- a/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py +++ b/src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py @@ -378,10 +378,10 @@ def rotation_scan( } ) @transmission_and_xbpm_feedback_for_collection_decorator( - parameters.demand_energy_ev, composite.undulator, composite.xbpm_feedback, composite.attenuator, + parameters.demand_energy_ev, parameters.transmission_frac, ) def rotation_scan_plan_with_stage_and_cleanup( @@ -433,10 +433,10 @@ 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.demand_energy_ev, parameters.transmission_frac, ) @bpp.finalize_decorator(lambda: _cleanup_plan(composite)) diff --git a/src/mx_bluesky/hyperion/experiment_plans/set_energy_plan.py b/src/mx_bluesky/hyperion/experiment_plans/set_energy_plan.py index 17455f231..c5df7234d 100644 --- a/src/mx_bluesky/hyperion/experiment_plans/set_energy_plan.py +++ b/src/mx_bluesky/hyperion/experiment_plans/set_energy_plan.py @@ -54,9 +54,9 @@ def set_energy_plan( if energy_ev: yield from transmission_and_xbpm_feedback_for_collection_wrapper( _set_energy_plan(energy_ev / 1000, composite), - energy_ev, composite.undulator_dcm.undulator_ref(), composite.xbpm_feedback, composite.attenuator, + energy_ev, DESIRED_TRANSMISSION_FRACTION, ) diff --git a/tests/conftest.py b/tests/conftest.py index 31c4b3485..7b9ed1a44 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -705,6 +705,12 @@ def fake_create_rotation_devices( ) +@pytest.fixture +def fake_undulator_set(undulator, done_status): + undulator.set = MagicMock(return_value=done_status) + return undulator + + @pytest.fixture def zocalo(done_status): zoc = i03.zocalo(fake_with_ophyd_sim=True) diff --git a/tests/system_tests/hyperion/experiment_plans/test_fgs_plan.py b/tests/system_tests/hyperion/experiment_plans/test_fgs_plan.py index e4ab7813f..2f7db94ec 100644 --- a/tests/system_tests/hyperion/experiment_plans/test_fgs_plan.py +++ b/tests/system_tests/hyperion/experiment_plans/test_fgs_plan.py @@ -159,10 +159,10 @@ 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.demand_energy_ev, params.transmission_frac, ) def decorated_plan(): diff --git a/tests/unit_tests/hyperion/device_setup_plans/test_xbpm_feedback.py b/tests/unit_tests/hyperion/device_setup_plans/test_xbpm_feedback.py index c0813b446..6fcee1506 100644 --- a/tests/unit_tests/hyperion/device_setup_plans/test_xbpm_feedback.py +++ b/tests/unit_tests/hyperion/device_setup_plans/test_xbpm_feedback.py @@ -16,13 +16,13 @@ async def test_given_xpbm_checks_pass_when_plan_run_with_decorator_then_run_as_expected( RE, xbpm_feedback, - undulator, + fake_undulator_set, attenuator, ): expected_transmission = 0.3 @transmission_and_xbpm_feedback_for_collection_decorator( - 100, undulator, xbpm_feedback, attenuator, expected_transmission + fake_undulator_set, xbpm_feedback, attenuator, 100, expected_transmission ) def my_collection_plan(): read_transmission = yield from bps.rd(attenuator.actual_transmission) @@ -42,13 +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, + fake_undulator_set, attenuator, ): mock = MagicMock() @transmission_and_xbpm_feedback_for_collection_decorator( - 100, undulator, xbpm_feedback, attenuator, 0.1 + fake_undulator_set, xbpm_feedback, attenuator, 100, 0.1 ) def my_collection_plan(): mock() @@ -70,7 +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, + fake_undulator_set, attenuator, ): set_mock_value(xbpm_feedback.pos_stable, True) # type: ignore @@ -79,7 +79,7 @@ class MyException(Exception): pass @transmission_and_xbpm_feedback_for_collection_decorator( - 100, undulator, xbpm_feedback, attenuator, 0.1 + fake_undulator_set, xbpm_feedback, attenuator, 100, 0.1 ) def my_collection_plan(): yield from bps.null()