Change shutter to auto when configuring the zebra

src/mx_bluesky/hyperion/device_setup_plans/setup_zebra.py

@@ -22,6 +22,7 @@
     RotationDirection,
     Zebra,
 )
+from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl
 
 from mx_bluesky.hyperion.log import LOGGER
 
@@ -57,25 +58,34 @@ def arm_zebra(zebra: Zebra):
     yield from bps.abs_set(zebra.pc.arm, ArmDemand.ARM, wait=True)
 
 
-def disarm_zebra(zebra: Zebra):
-    yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, wait=True)
+def tidy_up_zebra_after_rotation_scan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="tidy_up_zebra_after_rotation",
+    wait=True,
+):
+    yield from bps.abs_set(zebra.pc.arm, ArmDemand.DISARM, group=group)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
+    )
+    if wait:
+        yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
 
 
-def set_shutter_auto_input(
-    zebra: Zebra, input: int, group="set_shutter_trigger", wait=True
-):
+def set_shutter_auto_input(zebra: Zebra, input: int, group="set_shutter_trigger"):
     """Set the input that the shutter uses when set to auto.
 
     For more details see the ZebraShutter device."""
     auto_shutter_control = zebra.logic_gates.and_gates[AUTO_SHUTTER_GATE]
     yield from bps.abs_set(
-        auto_shutter_control.sources[AUTO_SHUTTER_INPUT], input, group, wait
+        auto_shutter_control.sources[AUTO_SHUTTER_INPUT], input, group
     )
 
 
 @bluesky_retry
 def setup_zebra_for_rotation(
     zebra: Zebra,
+    zebra_shutter: ZebraShutter,
     axis: EncEnum = I03Axes.OMEGA,
     start_angle: float = 0,
     scan_width: float = 360,
@@ -128,7 +138,10 @@ def setup_zebra_for_rotation(
     # Set gate position to be angle of interest
     yield from bps.abs_set(zebra.pc.gate_trigger, axis.value, group=group)
     # Trigger the shutter with the gate
-    yield from set_shutter_auto_input(zebra, PC_GATE)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+    yield from set_shutter_auto_input(zebra, PC_GATE, group=group)
     # Trigger the detector with a pulse
     yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], PC_PULSE, group=group)
     # Don't use the fluorescence detector
@@ -140,9 +153,17 @@ def setup_zebra_for_rotation(
 
 
 @bluesky_retry
-def setup_zebra_for_gridscan(zebra: Zebra, group="setup_zebra_for_gridscan", wait=True):
+def setup_zebra_for_gridscan(
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="setup_zebra_for_gridscan",
+    wait=True,
+):
     yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN3_TTL, group=group)
-    yield from set_shutter_auto_input(zebra, IN4_TTL)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+    yield from set_shutter_auto_input(zebra, IN4_TTL, group=group)
     yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group)
     yield from bps.abs_set(zebra.output.pulse_1.input, DISCONNECT, group=group)
 
@@ -152,24 +173,36 @@ def setup_zebra_for_gridscan(zebra: Zebra, group="setup_zebra_for_gridscan", wai
 
 @bluesky_retry
 def tidy_up_zebra_after_gridscan(
-    zebra: Zebra, group="tidy_up_zebra_after_gridscan", wait=True
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="tidy_up_zebra_after_gridscan",
+    wait=True,
 ) -> MsgGenerator:
     yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], PC_PULSE, group=group)
-    yield from set_shutter_auto_input(zebra, PC_GATE)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.MANUAL, group=group
+    )
+    yield from set_shutter_auto_input(zebra, PC_GATE, group=group)
 
     if wait:
         yield from bps.wait(group, timeout=ZEBRA_STATUS_TIMEOUT)
 
 
 @bluesky_retry
 def setup_zebra_for_panda_flyscan(
-    zebra: Zebra, group="setup_zebra_for_panda_flyscan", wait=True
+    zebra: Zebra,
+    zebra_shutter: ZebraShutter,
+    group="setup_zebra_for_panda_flyscan",
+    wait=True,
 ):
     # Forwards eiger trigger signal from panda
     yield from bps.abs_set(zebra.output.out_pvs[TTL_DETECTOR], IN1_TTL, group=group)
 
     # Forwards signal from PPMAC to fast shutter. High while panda PLC is running
-    yield from set_shutter_auto_input(zebra, IN4_TTL)
+    yield from bps.abs_set(
+        zebra_shutter.control_mode, ZebraShutterControl.AUTO, group=group
+    )
+    yield from set_shutter_auto_input(zebra, IN4_TTL, group=group)
 
     yield from bps.abs_set(zebra.output.out_pvs[TTL_XSPRESS3], DISCONNECT, group=group)
 

src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -35,7 +35,7 @@
 from dodal.devices.undulator import Undulator
 from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra import Zebra
-from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl
+from dodal.devices.zebra_controlled_shutter import ZebraShutter
 from dodal.devices.zocalo.zocalo_results import (
     ZOCALO_READING_PLAN_NAME,
     ZOCALO_STAGE_GROUP,
@@ -275,13 +275,6 @@ def run_gridscan(
     LOGGER.info("Waiting for gridscan validity check")
     yield from wait_for_gridscan_valid(feature_controlled.fgs_motors)
 
-    LOGGER.info("Set shutter to auto")
-    yield from bps.abs_set(
-        fgs_composite.sample_shutter.control_mode,
-        ZebraShutterControl.AUTO,
-        group=CONST.WAIT.GRID_READY_FOR_DC,
-    )
-
     LOGGER.info("Waiting for arming to finish")
     yield from bps.wait(CONST.WAIT.GRID_READY_FOR_DC)
     yield from bps.stage(fgs_composite.eiger)
@@ -446,10 +439,9 @@ def _get_feature_controlled(
 def _generic_tidy(
     fgs_composite: FlyScanXRayCentreComposite, group, wait=True
 ) -> MsgGenerator:
-    LOGGER.info("Tidying up Zebra and shutter")
-    yield from tidy_up_zebra_after_gridscan(fgs_composite.zebra, group=group, wait=wait)
-    yield from bps.abs_set(
-        fgs_composite.sample_shutter.control_mode, ZebraShutterControl.MANUAL
+    LOGGER.info("Tidying up Zebra")
+    yield from tidy_up_zebra_after_gridscan(
+        fgs_composite.zebra, fgs_composite.sample_shutter, group=group, wait=wait
     )
     LOGGER.info("Tidying up Zocalo")
     # make sure we don't consume any other results
@@ -470,7 +462,9 @@ def _zebra_triggering_setup(
     parameters: ThreeDGridScan,
     initial_xyz: np.ndarray,
 ):
-    yield from setup_zebra_for_gridscan(fgs_composite.zebra, wait=True)
+    yield from setup_zebra_for_gridscan(
+        fgs_composite.zebra, fgs_composite.sample_shutter, wait=True
+    )
 
 
 def _panda_triggering_setup(
@@ -528,4 +522,6 @@ def _panda_triggering_setup(
     )
 
     LOGGER.info("Setting up Zebra for panda flyscan")
-    yield from setup_zebra_for_panda_flyscan(fgs_composite.zebra, wait=True)
+    yield from setup_zebra_for_panda_flyscan(
+        fgs_composite.zebra, fgs_composite.sample_shutter, wait=True
+    )

src/mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py

@@ -20,7 +20,7 @@
 from dodal.devices.synchrotron import Synchrotron
 from dodal.devices.undulator import Undulator
 from dodal.devices.zebra import RotationDirection, Zebra
-from dodal.devices.zebra_controlled_shutter import ZebraShutter, ZebraShutterControl
+from dodal.devices.zebra_controlled_shutter import ZebraShutter
 from dodal.plans.check_topup import check_topup_and_wait_if_necessary
 
 from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import (
@@ -37,8 +37,8 @@
 )
 from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
     arm_zebra,
-    disarm_zebra,
     setup_zebra_for_rotation,
+    tidy_up_zebra_after_rotation_scan,
 )
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     start_preparing_data_collection_then_do_plan,
@@ -218,12 +218,9 @@ def _rotation_scan_plan(
             group=CONST.WAIT.ROTATION_READY_FOR_DC,
         )
 
-        yield from bps.abs_set(
-            composite.sample_shutter.control_mode, ZebraShutterControl.AUTO
-        )
-
         yield from setup_zebra_for_rotation(
             composite.zebra,
+            composite.sample_shutter,
             start_angle=motion_values.start_scan_deg,
             scan_width=motion_values.scan_width_deg,
             direction=motion_values.direction,
@@ -289,12 +286,10 @@ def _cleanup_plan(composite: RotationScanComposite, **kwargs):
     max_vel = yield from bps.rd(composite.smargon.omega.max_velocity)
     yield from cleanup_sample_environment(composite.detector_motion, group="cleanup")
     yield from bps.abs_set(composite.smargon.omega.velocity, max_vel, group="cleanup")
-    yield from bps.abs_set(
-        composite.sample_shutter.control_mode,
-        ZebraShutterControl.MANUAL,
-        group="cleanup",
+    yield from tidy_up_zebra_after_rotation_scan(
+        composite.zebra, composite.sample_shutter, group="cleanup", wait=False
     )
-    yield from bpp.finalize_wrapper(disarm_zebra(composite.zebra), bps.wait("cleanup"))
+    yield from bps.wait("cleanup")
 
 
 def _move_and_rotation(

tests/unit_tests/hyperion/device_setup_plans/test_zebra_setup.py

@@ -2,7 +2,6 @@
 
 import pytest
 from bluesky import plan_stubs as bps
-from bluesky.run_engine import RunEngine
 from dodal.beamlines import i03
 from dodal.devices.zebra import (
     AUTO_SHUTTER_GATE,
@@ -17,6 +16,7 @@
     I03Axes,
     Zebra,
 )
+from dodal.devices.zebra_controlled_shutter import ZebraShutter
 
 from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
     bluesky_retry,
@@ -28,11 +28,15 @@
 
 
 @pytest.fixture
-def zebra():
-    RunEngine()
+def zebra(RE):
     return i03.zebra(fake_with_ophyd_sim=True)
 
 
+@pytest.fixture
+def zebra_shutter(RE):
+    return i03.sample_shutter(fake_with_ophyd_sim=True)
+
+
 async def _get_shutter_input(zebra: Zebra):
     return (
         await zebra.logic_gates.and_gates[AUTO_SHUTTER_GATE]
@@ -41,27 +45,29 @@ async def _get_shutter_input(zebra: Zebra):
     )
 
 
-async def test_zebra_set_up_for_panda_gridscan(RE, zebra: Zebra):
-    RE(setup_zebra_for_panda_flyscan(zebra, wait=True))
+async def test_zebra_set_up_for_panda_gridscan(
+    RE, zebra: Zebra, zebra_shutter: ZebraShutter
+):
+    RE(setup_zebra_for_panda_flyscan(zebra, zebra_shutter, wait=True))
     assert await zebra.output.out_pvs[TTL_DETECTOR].get_value() == IN1_TTL
     assert await zebra.output.out_pvs[TTL_PANDA].get_value() == IN3_TTL
     assert await _get_shutter_input(zebra) == IN4_TTL
 
 
-async def test_zebra_set_up_for_gridscan(RE, zebra: Zebra):
-    RE(setup_zebra_for_gridscan(zebra, wait=True))
+async def test_zebra_set_up_for_gridscan(RE, zebra: Zebra, zebra_shutter: ZebraShutter):
+    RE(setup_zebra_for_gridscan(zebra, zebra_shutter, wait=True))
     assert await zebra.output.out_pvs[TTL_DETECTOR].get_value() == IN3_TTL
     assert await _get_shutter_input(zebra) == IN4_TTL
 
 
-async def test_zebra_set_up_for_rotation(RE, zebra: Zebra):
-    RE(setup_zebra_for_rotation(zebra, wait=True))
+async def test_zebra_set_up_for_rotation(RE, zebra: Zebra, zebra_shutter: ZebraShutter):
+    RE(setup_zebra_for_rotation(zebra, zebra_shutter, wait=True))
     assert await zebra.pc.gate_trigger.get_value() == I03Axes.OMEGA.value
     assert await zebra.pc.gate_width.get_value() == pytest.approx(360, 0.01)
 
 
-async def test_zebra_cleanup(RE, zebra: Zebra):
-    RE(tidy_up_zebra_after_gridscan(zebra, wait=True))
+async def test_zebra_cleanup(RE, zebra: Zebra, zebra_shutter: ZebraShutter):
+    RE(tidy_up_zebra_after_gridscan(zebra, zebra_shutter, wait=True))
     assert await zebra.output.out_pvs[TTL_DETECTOR].get_value() == PC_PULSE
     assert await _get_shutter_input(zebra) == PC_GATE