Fix extruder tests and add one for idy up plan

DiamondLightSource · Jul 12, 2024 · bd139c9 · bd139c9
1 parent e97e35b
commit bd139c9
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 41 deletions.
diff --git a/src/mx_bluesky/I24/serial/extruder/i24ssx_Extruder_Collect_py3v2.py b/src/mx_bluesky/I24/serial/extruder/i24ssx_Extruder_Collect_py3v2.py
@@ -405,9 +405,9 @@ def tidy_up_at_collection_end_plan(
  """A plan to tidy up at the end of a collection.
 
  Args:
- zebra (Zebra): _description_
- shutter (HutchShutter): _description_
- parameters (ExtruderParameters): _description_
+ zebra (Zebra): The Zebra device.
+ shutter (HutchShutter): The HutchShutter device.
+ parameters (ExtruderParameters): Collection parameters.
  """
  yield from reset_zebra_when_collection_done_plan(zebra)
 
@@ -428,7 +428,7 @@ def tidy_up_at_collection_end_plan(
  sup.pilatus("return-to-normal")
  elif parameters.detector_name == "eiger":
  sup.eiger("return-to-normal")
- logger.debug(parameters.filename + "_" + caget(pv.eiger_seqID))
+ logger.debug(f"{parameters.filename}_{caget(pv.eiger_seqID)}")
  logger.debug("End of Run")
  logger.debug("Close hutch shutter")
  yield from bps.abs_set(shutter, ShutterDemand.CLOSE, wait=True)

diff --git a/tests/I24/serial/extruder/test_extruder_collect.py b/tests/I24/serial/extruder/test_extruder_collect.py
@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, call, patch
+from unittest.mock import ANY, MagicMock, call, patch
 
 import bluesky.plan_stubs as bps
 import pytest
@@ -12,6 +12,7 @@
  initialise_extruder,
  laser_check,
  run_main_extruder_plan,
+ tidy_up_at_collection_end_plan,
 )
 from mx_bluesky.I24.serial.parameters import ExtruderParameters
 from mx_bluesky.I24.serial.setup_beamline import Eiger, Pilatus
@@ -111,37 +112,27 @@ async def test_laser_check(
  assert await zebra.output.out_pvs[TTL].get_value() == expected_out
 
 
-@patch(
- "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.write_parameter_file",
-)
-@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.shutil")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.sleep")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.DCID")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.call_nexgen")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.caput")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.caget")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.sup")
-@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.get_detector_type")
 @patch(
  "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.setup_zebra_for_quickshot_plan"
 )
-@patch(
- "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.ExtruderParameters"
-)
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.setup_logging")
+@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.bps.rd")
 def test_run_extruder_quickshot_with_eiger(
+ fake_read,
  fake_log_setup,
- mock_params,
  mock_quickshot_plan,
- fake_det,
  fake_sup,
  fake_caget,
  fake_caput,
  fake_nexgen,
  fake_dcid,
  fake_sleep,
- fake_shutil,
- fake_write_params,
  RE,
  zebra,
  shutter,
@@ -151,9 +142,9 @@ def test_run_extruder_quickshot_with_eiger(
  detector_stage,
  dummy_params,
 ):
- mock_params.from_file.return_value = dummy_params
- fake_det.return_value = Eiger()
  fake_start_time = MagicMock()
+ # Mock end of data collection (zebra disarmed)
+ fake_read.side_effect = [fake_generator(0)]
  RE(
  run_main_extruder_plan(
  zebra,
@@ -168,46 +159,32 @@ def test_run_extruder_quickshot_with_eiger(
  )
  )
  assert fake_nexgen.call_count == 1
- assert fake_dcid.call_count == 1
+ assert fake_dcid.notify_start.call_count == 1
  assert fake_sup.setup_beamline_for_collection_plan.call_count == 1
  # Check temporary piilatus hack is in there
  assert fake_sup.pilatus.call_count == 2
  mock_quickshot_plan.assert_called_once()
 
 
-@patch(
- "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.write_parameter_file"
-)
-@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.shutil")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.sleep")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.DCID")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.caput")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.caget")
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.sup")
-@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.get_detector_type")
 @patch(
  "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.setup_zebra_for_extruder_with_pump_probe_plan"
 )
-@patch(
- "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.reset_zebra_when_collection_done_plan"
-)
-@patch(
- "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.ExtruderParameters"
-)
 @patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.setup_logging")
+@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.bps.rd")
 def test_run_extruder_pump_probe_with_pilatus(
+ fake_read,
  fake_log_setup,
- mock_params,
- mock_reset_zebra_plan,
  mock_pp_plan,
- fake_det,
  fake_sup,
  fake_caget,
  fake_caput,
  fake_dcid,
  fake_sleep,
- fake_shutil,
- fake_write_params,
  RE,
  zebra,
  shutter,
@@ -217,9 +194,9 @@ def test_run_extruder_pump_probe_with_pilatus(
  detector_stage,
  dummy_params_pp,
 ):
- mock_params.from_file.return_value = dummy_params_pp
- fake_det.return_value = Pilatus()
  fake_start_time = MagicMock()
+ # Mock end of data collection (zebra disarmed)
+ fake_read.side_effect = [fake_generator(0)]
  RE(
  run_main_extruder_plan(
  zebra,
@@ -233,15 +210,47 @@ def test_run_extruder_pump_probe_with_pilatus(
  fake_start_time,
  )
  )
- assert fake_dcid.call_count == 1
+ assert fake_dcid.notify_start.call_count == 1
  assert fake_sup.move_detector_stage_to_position_plan.call_count == 1
  mock_pp_plan.assert_called_once()
- mock_reset_zebra_plan.assert_called_once()
 
  shutter_call_list = [
  call("Reset", wait=True, timeout=10.0),
  call("Open", wait=True, timeout=10.0),
- call("Close", wait=True, timeout=10.0),
  ]
  mock_shutter = get_mock_put(shutter.control)
  mock_shutter.assert_has_calls(shutter_call_list)
+
+
+@patch(
+ "mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.reset_zebra_when_collection_done_plan"
+)
+@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.DCID")
+@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.caput")
+@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.caget")
+@patch("mx_bluesky.I24.serial.extruder.i24ssx_Extruder_Collect_py3v2.sup")
+def test_tidy_up_at_collection_end_plan_with_eiger(
+ fake_sup,
+ fake_caget,
+ fake_caput,
+ fake_dcid,
+ mock_reset_zebra_plan,
+ RE,
+ zebra,
+ shutter,
+ dummy_params,
+):
+ RE(tidy_up_at_collection_end_plan(zebra, shutter, dummy_params, fake_dcid))
+
+ mock_reset_zebra_plan.assert_called_once()
+ mock_shutter = get_mock_put(shutter.control)
+ mock_shutter.assert_has_calls([call("Close", wait=True, timeout=10.0)])
+
+ assert fake_dcid.collection_complete.call_count == 1
+ assert fake_dcid.notify_end.call_count == 1
+ assert fake_caget.call_count == 1
+
+ call_list = [call(ANY, 0), call(ANY, "Done")]
+ fake_caput.assert_has_calls(call_list)
+
+ fake_sup.eiger.assert_called_once_with("return-to-normal")