diff --git a/src/hyperion/device_setup_plans/setup_zebra.py b/src/hyperion/device_setup_plans/setup_zebra.py index fc9e4e7bb..add7a5328 100644 --- a/src/hyperion/device_setup_plans/setup_zebra.py +++ b/src/hyperion/device_setup_plans/setup_zebra.py @@ -1,4 +1,8 @@ +from functools import wraps +from typing import Callable + import bluesky.plan_stubs as bps +import bluesky.preprocessors as bpp from dodal.devices.zebra import ( DISCONNECT, IN1_TTL, @@ -20,6 +24,31 @@ from hyperion.log import LOGGER +def bluesky_retry(func: Callable): + """Decorator that will retry the decorated plan if it fails. + + Use this with care as it knows nothing about the state of the world when things fail. + If it is possible that your plan fails when the beamline is in a transient state that + the plan could not act on do not use this decorator without doing some more intelligent + clean up. + + You should avoid using this decorator often in general production as it hides errors, + instead it should be used only for debugging these underlying errors. + """ + + @wraps(func) + def newfunc(*args, **kwargs): + def log_and_retry(exception): + LOGGER.error(f"Function {func.__name__} failed with {exception}, retrying") + yield from func(*args, **kwargs) + + yield from bpp.contingency_wrapper( + func(*args, **kwargs), except_plan=log_and_retry, auto_raise=False + ) + + return newfunc + + def arm_zebra(zebra: Zebra): yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True) @@ -28,6 +57,7 @@ def disarm_zebra(zebra: Zebra): yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True) +@bluesky_retry def setup_zebra_for_rotation( zebra: Zebra, axis: I03Axes = I03Axes.OMEGA, @@ -93,6 +123,7 @@ def setup_zebra_for_rotation( yield from bps.wait(group) +@bluesky_retry def setup_zebra_for_gridscan( zebra: Zebra, group="setup_zebra_for_gridscan", wait=False ): @@ -105,6 +136,7 @@ def setup_zebra_for_gridscan( yield from bps.wait(group) +@bluesky_retry def set_zebra_shutter_to_manual( zebra: Zebra, group="set_zebra_shutter_to_manual", wait=False ): @@ -115,10 +147,12 @@ def set_zebra_shutter_to_manual( yield from bps.wait(group) +@bluesky_retry def make_trigger_safe(zebra: Zebra, group="make_zebra_safe", wait=False): yield from bps.abs_set(zebra.inputs.soft_in_1, 0, wait=wait, group=group) +@bluesky_retry def setup_zebra_for_panda_flyscan( zebra: Zebra, group="setup_zebra_for_panda_flyscan", wait=False ): diff --git a/tests/unit_tests/device_setup_plans/test_zebra_setup.py b/tests/unit_tests/device_setup_plans/test_zebra_setup.py index b69bea212..beffb083e 100644 --- a/tests/unit_tests/device_setup_plans/test_zebra_setup.py +++ b/tests/unit_tests/device_setup_plans/test_zebra_setup.py @@ -1,7 +1,8 @@ from functools import partial -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest +from bluesky import plan_stubs as bps from dodal.beamlines import i03 from dodal.devices.zebra import ( IN3_TTL, @@ -17,6 +18,7 @@ from hyperion.device_setup_plans.setup_zebra import ( arm_zebra, + bluesky_retry, disarm_zebra, set_zebra_shutter_to_manual, setup_zebra_for_gridscan, @@ -81,3 +83,37 @@ def side_effect(set_armed_to: int, _): with pytest.raises(Exception): zebra.pc.arm.armed.set(1) RE(disarm_zebra(zebra, 0.2)) + + +class MyException(Exception): + pass + + +def test_when_first_try_fails_then_bluesky_retry_tries_again(RE, done_status): + mock_device = MagicMock() + + @bluesky_retry + def my_plan(value): + yield from bps.abs_set(mock_device, value) + + mock_device.set.side_effect = [MyException(), done_status] + + RE(my_plan(10)) + + assert mock_device.set.mock_calls == [call(10), call(10)] + + +def test_when_all_tries_fail_then_bluesky_retry_throws_error(RE, done_status): + mock_device = MagicMock() + + @bluesky_retry + def my_plan(value): + yield from bps.abs_set(mock_device, value) + + exception_2 = MyException() + mock_device.set.side_effect = [MyException(), exception_2] + + with pytest.raises(MyException) as e: + RE(my_plan(10)) + + assert e.value == exception_2