Fix i04 bugs

src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py

@@ -226,21 +226,22 @@ def grid_detect_then_xray_centre_with_callbacks():
                 oav_config=oav_config,
             )
 
-        try:
-            yield from grid_detect_then_xray_centre_with_callbacks()
-            assert isinstance(
-                grid_common_params.specified_grid_params, SpecifiedThreeDGridScan
-            ), "Specified grid params couldn't be found after grid detection"
-            yield from get_results_and_move_to_xtal(
-                composite,
-                grid_common_params.specified_grid_params,
-                flyscan_event_handler,
-            )
-        except CrystalNotFoundError:
-            yield from bps.mv(
-                smargon.x, initial_x, smargon.y, initial_y, smargon.z, initial_z
-            )
-            raise
+            try:
+                assert isinstance(
+                    grid_common_params.specified_grid_params, SpecifiedThreeDGridScan
+                ), "Specified grid params couldn't be found after grid detection"
+                yield from get_results_and_move_to_xtal(
+                    composite,
+                    grid_common_params.specified_grid_params,
+                    flyscan_event_handler,
+                )
+            except CrystalNotFoundError:
+                yield from bps.mv(
+                    smargon.x, initial_x, smargon.y, initial_y, smargon.z, initial_z
+                )
+                raise
+
+        yield from grid_detect_then_xray_centre_with_callbacks()
 
     yield from _change_beamsize(
         transfocator, DEFAULT_XRC_BEAMSIZE_MICRONS, grid_common_params

src/mx_bluesky/beamlines/i04/oav_centering_plans/oav_imaging.py

@@ -35,10 +35,10 @@ class FindBeamCentresComposite:
     scintillator: Scintillator
     xbpm_feedback: XBPMFeedback
     max_pixel: MaxPixel
-    centre_ellipse: CentreEllipseMethod
+    beam_centre: CentreEllipseMethod
     attenuator: BinaryFilterAttenuator
     zoom_controller: ZoomControllerWithBeamCentres
-    shutter: ZebraShutter
+    sample_shutter: ZebraShutter
 
 
 def take_oav_image_with_scintillator_in(
@@ -299,7 +299,7 @@ def find_beam_centres(
         composite.backlight,
         composite.scintillator,
         composite.xbpm_feedback,
-        composite.shutter,
+        composite.sample_shutter,
         OAV_PREPARE_BEAMLINE_FOR_SCINT_WAIT,
     )
     LOGGER.info("Waiting for prepare beamline plan to complete...")
@@ -320,9 +320,9 @@ def find_beam_centres(
                     xbpm_feedback=composite.xbpm_feedback,
                 )
 
-            yield from bps.trigger(composite.centre_ellipse, wait=True)
-            centre_x = yield from bps.rd(composite.centre_ellipse.center_x_val)
-            centre_y = yield from bps.rd(composite.centre_ellipse.center_y_val)
+            yield from bps.trigger(composite.beam_centre, wait=True)
+            centre_x = yield from bps.rd(composite.beam_centre.center_x_val)
+            centre_y = yield from bps.rd(composite.beam_centre.center_y_val)
             centre_x = round(centre_x)
             centre_y = round(centre_y)
             LOGGER.info(

tests/conftest.py

@@ -1239,7 +1239,7 @@ def default_raw_gridscan_params(
     return raw_params_from_file(json_file, tmp_path)
 
 
-def _dummy_params(tmp_path):
+def dummy_params(tmp_path):
     dummy_params = SpecifiedThreeDGridScan(
         **raw_params_from_file(
             "tests/test_data/parameter_json_files/test_gridscan_param_defaults.json",
@@ -1354,7 +1354,7 @@ def test_grid_detect_and_gridscan_start_document(self) -> RunStart:
             "scan_id": 1,
             "plan_type": "generator",
             "subplan_name": PlanNameConstants.GRID_DETECT_AND_DO_GRIDSCAN,
-            "mx_bluesky_parameters": _dummy_params(self._tmp_path).model_dump_json(),
+            "mx_bluesky_parameters": dummy_params(self._tmp_path).model_dump_json(),
         }
 
     @property
@@ -1388,7 +1388,7 @@ def test_gridscan_outer_start_document(self):
             "plan_name": PlanNameConstants.GRIDSCAN_OUTER,
             "subplan_name": PlanNameConstants.GRIDSCAN_OUTER,
             "zocalo_environment": EnvironmentConstants.ZOCALO_ENV,
-            "mx_bluesky_parameters": _dummy_params(self._tmp_path).model_dump_json(),
+            "mx_bluesky_parameters": dummy_params(self._tmp_path).model_dump_json(),
         }
 
     @property

tests/unit_tests/beamlines/i04/test_i04_grid_detect_then_xray_centre_plan.py

@@ -50,7 +50,12 @@
 )
 from mx_bluesky.common.utils.exceptions import CrystalNotFoundError
 from mx_bluesky.common.utils.xrc_result import XRayCentreResult
-from tests.conftest import TEST_RESULT_LARGE, fake_generator, simulate_xrc_result
+from tests.conftest import (
+    TEST_RESULT_LARGE,
+    dummy_params,
+    fake_generator,
+    simulate_xrc_result,
+)
 from tests.unit_tests.common.experiment_plans.test_common_flyscan_xray_centre_plan import (
     CompleteError,
 )
@@ -424,6 +429,10 @@ async def test_i04_grid_detect_then_xrc_sets_beamsize_before_grid_detect_then_re
     ]
 
 
+@patch(
+    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal",
+    autospec=True,
+)
 @patch(
     "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_ready_for_oav_and_close_shutter",
     autospec=True,
@@ -435,16 +444,23 @@ async def test_i04_grid_detect_then_xrc_sets_beamsize_before_grid_detect_then_re
 async def test_given_no_diffraction_found_i04_grid_detect_then_xrc_returns_sample_to_initial_position(
     mock_grid_detect_then_xray_centre: MagicMock,
     mock_get_ready_for_oav_and_close_shutter: MagicMock,
+    mock_get_results_and_move_to_xtal: MagicMock,
     run_engine: RunEngine,
     i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
     smargon: Smargon,
+    tmp_path,
 ):
     initial_x, initial_y, initial_z = 1, 2, 3
     set_mock_value(smargon.x.user_readback, initial_x)
     set_mock_value(smargon.y.user_readback, initial_y)
     set_mock_value(smargon.z.user_readback, initial_z)
 
-    mock_grid_detect_then_xray_centre.side_effect = CrystalNotFoundError
+    def fake_xray_centre(parameters: GenericGrid, **__):
+        parameters.set_specified_grid_params(dummy_params(tmp_path))
+        yield Msg(command="open_run")
+
+    mock_grid_detect_then_xray_centre.side_effect = fake_xray_centre
+    mock_get_results_and_move_to_xtal.side_effect = CrystalNotFoundError
 
     with pytest.raises(CrystalNotFoundError):
         run_engine(i04_grid_detect_then_xrc_default_params())

tests/unit_tests/beamlines/i04/test_oav_imaging.py

@@ -704,10 +704,10 @@ def find_beam_centre_devices(
         scintillator=scintillator,
         xbpm_feedback=xbpm_feedback,
         max_pixel=max_pixel,
-        centre_ellipse=centre_ellipse,
+        beam_centre=centre_ellipse,
         attenuator=attenuator,
         zoom_controller=zoom_controller_with_centres,
-        shutter=sample_shutter,
+        sample_shutter=sample_shutter,
     )
 
 
@@ -750,7 +750,8 @@ def test_find_beam_centres_starts_by_prepping_scintillator(
 ):
     run_engine(find_beam_centres(composite=find_beam_centre_devices))
     mock_prepare_scintillator.assert_called_once()
-    assert find_beam_centre_devices.centre_ellipse.trigger.call_count == 4  # type: ignore
+    assert isinstance(find_beam_centre_devices.beam_centre.trigger, MagicMock)
+    assert find_beam_centre_devices.beam_centre.trigger.call_count == 4
 
 
 def mock_centre_ellipse_with_given_centres(
@@ -781,7 +782,7 @@ async def test_find_beam_centres_iterates_and_sets_centres(
     new_centres = [(100, 100), (200, 200), (300, 300), (400, 400)]
     expected_centres = new_centres.copy()
 
-    centre_ellipse = find_beam_centre_devices.centre_ellipse
+    centre_ellipse = find_beam_centre_devices.beam_centre
     zoom_controller_with_centres: ZoomControllerWithBeamCentres = (
         find_beam_centre_devices.zoom_controller
     )
@@ -817,7 +818,7 @@ async def test_if_only_some_levels_given_then_find_beam_centres_iterates_and_set
 ):
     new_centres = [(100, 100), (200, 200), (300, 300), (400, 400)]
 
-    centre_ellipse = find_beam_centre_devices.centre_ellipse
+    centre_ellipse = find_beam_centre_devices.beam_centre
     zoom_controller_with_centres: ZoomControllerWithBeamCentres = (
         find_beam_centre_devices.zoom_controller
     )