use enum type as soon as args are parsed

DiamondLightSource · Nov 10, 2023 · 8994d57 · 8994d57
1 parent 0586ac1
commit 8994d57
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/src/mx_bluesky/I24/serial/setup_beamline/setup_detector.py b/src/mx_bluesky/I24/serial/setup_beamline/setup_detector.py
@@ -13,7 +13,7 @@
 from mx_bluesky.I24.serial import log
 from mx_bluesky.I24.serial.parameters.constants import SSXType
 from mx_bluesky.I24.serial.setup_beamline import pv
-from mx_bluesky.I24.serial.setup_beamline.ca import caget  # , caput
+from mx_bluesky.I24.serial.setup_beamline.ca import caget
 from mx_bluesky.I24.serial.setup_beamline.pv_abstract import (
     Detector,
     Eiger,
@@ -22,6 +22,11 @@
 
 logger = logging.getLogger("I24ssx.sup_det")
 
+EXPT_TYPE_DETECTOR_PVS = {
+    SSXType.FIXED: pv.me14e_gp101,
+    SSXType.EXTRUDER: pv.ioc12_gp15,
+}
+
 
 def setup_logging():
     logfile = time.strftime("SSXdetectorOps_%d%B%y.log").lower()
@@ -56,10 +61,10 @@ def _move_detector_stage(detector_stage: DetectorMotion, target: float):
     )
 
 
-def setup_detector_stage(detector_stage: DetectorMotion, expt_type: str):
+def setup_detector_stage(detector_stage: DetectorMotion, expt_type: SSXType):
     # Grab the correct PV depending on experiment
     # Its value is set with MUX on edm screen
-    det_type = pv.me14e_gp101 if expt_type == SSXType.FIXED.value else pv.ioc12_gp15
+    det_type = EXPT_TYPE_DETECTOR_PVS[expt_type]
     requested_detector = caget(det_type)
     logger.info(f"Requested detector: {requested_detector}.")
     det_y_target = (
@@ -71,17 +76,16 @@ def setup_detector_stage(detector_stage: DetectorMotion, expt_type: str):
 
 if __name__ == "__main__":
     setup_logging()
-    RE = RunEngine()
-    # Use dodal device for move
-    detector_stage = i24.detector_motion()
-
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "expt",
         type=str,
-        choices=["Serial Jet", "Serial Fixed"],
+        choices=[expt.value for expt in SSXType],
         help="Type of serial experiment being run.",
     )
-
     args = parser.parse_args()
-    RE(setup_detector_stage(detector_stage, args.expt))
+    expt_type = SSXType(args.expt)
+    RE = RunEngine()
+    # Use dodal device for move
+    detector_stage = i24.detector_motion()
+    RE(setup_detector_stage(detector_stage, expt_type))
diff --git a/tests/I24/serial/setup_beamline/test_setup_detector.py b/tests/I24/serial/setup_beamline/test_setup_detector.py
@@ -7,7 +7,8 @@
 from dodal.devices.i24.I24_detector_motion import DetectorMotion
 from ophyd.status import Status
 
-from mx_bluesky.I24.serial.setup_beamline import Eiger
+from mx_bluesky.I24.serial.parameters.constants import SSXType
+from mx_bluesky.I24.serial.setup_beamline import Eiger, Pilatus
 from mx_bluesky.I24.serial.setup_beamline.setup_detector import (
     get_detector_type,
     setup_detector_stage,
@@ -46,7 +47,11 @@ def test_get_detector_type_finds_pilatus(fake_caget):
 @patch("mx_bluesky.I24.serial.setup_beamline.setup_detector.caget")
 def test_setup_detector_stage_for_eiger(fake_caget, fake_detector_motion):
     RE = RunEngine()
-    fake_caget.return_value = "eiger"
 
-    RE(setup_detector_stage(fake_detector_motion, "Serial Fixed"))
+    fake_caget.return_value = "eiger"
+    RE(setup_detector_stage(fake_detector_motion, SSXType.FIXED))
     assert fake_detector_motion.y.user_readback.get() == Eiger.det_y_target
+
+    fake_caget.return_value = "pilatus"
+    RE(setup_detector_stage(fake_detector_motion, SSXType.EXTRUDER))
+    assert fake_detector_motion.y.user_readback.get() == Pilatus.det_y_target