allow 'exit now' as an option for optimise_axis_over_range()

src/ibex_bluesky_core/plans/reflectometry/_autoalign.py

@@ -148,12 +148,18 @@ def _optimise_axis_over_range(  # noqa: PLR0913 PLR0917
         yield from problem_found_plan()
 
         choice = yield from prompt_user_for_choice(
-            prompt=f"Type '1' if you would like to re-scan or type '2' to "
-            f"move {alignment_param.name} to {alignment_param_value} and keep going.",
-            choices=["1", "2"],
+            prompt=f"Type '1' if you would like to re-scan,\n '2' to "
+            f"move {alignment_param.name} to {alignment_param_value} and keep going, "
+            f"or\n '3' to pause now - use RE.resume() to rescan if you wish.\n",
+            choices=["1", "2", "3"],
         )
         if choice == "1":
             return icc, False
+        elif choice == "3":
+            _print_and_log("Plan paused.")
+            yield from bps.checkpoint()
+            yield from bps.pause()
+            return icc, False
 
     _print_and_log(f"Moving {alignment_param.name} to {alignment_param_value}.")
     yield from bps.mv(alignment_param, alignment_param_value)
@@ -221,7 +227,6 @@ def optimise_axis_against_intensity(  # noqa: PLR0913
     while True:  # If a problem is found, then start the alignment again
         all_ok = True
         icc = None
-
         for rel_scan_range in rel_scan_ranges:
             icc, all_ok = yield from _optimise_axis_over_range(
                 dae=dae,

tests/conftest.py

@@ -1,3 +1,5 @@
+from collections.abc import Generator
+
 import pytest
 from bluesky.run_engine import RunEngine
 
@@ -8,14 +10,19 @@
 
 
 @pytest.fixture
-def RE() -> RunEngine:
+def RE() -> Generator[RunEngine, None, None]:
     get_run_engine.cache_clear()
     RE = get_run_engine()
     # Clear the preprocessors as the rb number metadata injector relies
     # on a real signal which won't exist during testing
     RE.preprocessors.clear()
 
-    return RE
+    yield RE
+
+    # This is to avoid allowing one test to leave the RE in running/paused state
+    # which may affect further tests
+    if RE.state != "idle":
+        RE.abort()
 
 
 @pytest.fixture

tests/plans/test_reflectometry.py

@@ -5,7 +5,7 @@
 
 import bluesky.plan_stubs as bps
 import pytest
-from bluesky.utils import Msg
+from bluesky.utils import Msg, RunEngineInterrupted
 from lmfit.model import ModelResult
 
 from ibex_bluesky_core.callbacks import ISISCallbacks
@@ -95,7 +95,9 @@ def test_refl_adaptive_scan_creates_refl_param_device_and_simpledae(RE, prefix):
 # Auto-Alignment Utils
 
 
-def test_found_problem_callback_is_called_if_problem(RE, simpledae, prefix, monkeypatch):
+def test_optimise_axis_against_intensity_found_problem_callback_is_called_if_problem(
+    RE, simpledae, prefix, monkeypatch
+):
     """Test that if a problem is found then the problem callback plan is run"""
 
     def plan(mock) -> Generator[Msg, None, None]:
@@ -147,7 +149,9 @@ def plan(mock) -> Generator[Msg, None, None]:
         (0, None),
     ],
 )
-def test_alignment_param_value_outside_of_scan_range_returns_problem(param_value, problem_str):
+def test_optimise_axis_against_intensity_alignment_param_value_outside_scan_range_returns_problem(
+    param_value, problem_str
+):
     """Test that if the optimised value is outside of the scan range then it is reported"""
     with (
         patch("lmfit.model.ModelResult") as mr,
@@ -164,7 +168,7 @@ def test_alignment_param_value_outside_of_scan_range_returns_problem(param_value
 
 
 @pytest.mark.parametrize("problem", [False, True])
-def test_that_user_checks_are_called_when_provided(problem):
+def test_check_parameter_that_user_checks_are_called_when_provided(problem):
     """Test that a user provided check function on the optimised value is always ran"""
     mock = MagicMock()
 
@@ -186,7 +190,9 @@ def my_check(model: ModelResult, param_val: float) -> str | None:
         )
 
 
-def test_that_if_no_problem_found_then_motor_is_moved(RE, prefix, simpledae):
+def test_optimise_axis_against_intensity_that_if_no_problem_found_then_motor_is_moved(
+    RE, prefix, simpledae
+):
     """Test that if no problems are found with the optimised
     value then move the motor to it
     """
@@ -238,7 +244,9 @@ def mock_scan(*a, **k):
         mv.assert_called_once()
 
 
-def test_that_if_problem_found_and_type_1_then_re_scan(RE, prefix, simpledae, monkeypatch):
+def test_optimise_axis_against_intensity_if_problem_found_and_type_1_then_re_scan(
+    RE, prefix, simpledae, monkeypatch
+):
     """Test that if a problem is found, and the user types 1, then rescan.
     Then if they type 2, moves to value.
     """
@@ -340,3 +348,75 @@ def counter(str: str):
         )
 
         assert scan.call_count == 2
+
+
+async def test_optimise_axis_against_intensity_exits_if_three_selected_when_optimising_axis(
+    RE, prefix, simpledae, monkeypatch
+):
+    with (
+        patch("ibex_bluesky_core.devices.reflectometry.get_pv_prefix", return_value=prefix),
+        patch(
+            "ibex_bluesky_core.plans.reflectometry._autoalign.scan",
+            side_effect=[_fake_scan(), _fake_scan()],
+        ) as scan,
+        patch(
+            "ibex_bluesky_core.plans.reflectometry._autoalign.ensure_connected",
+            return_value=bps.null(),
+        ),
+    ):
+        param = ReflParameter(prefix=prefix, name="S1VG", changing_timeout_s=60)
+        await param.connect(mock=True)
+        monkeypatch.setattr("builtins.input", lambda str: "3")
+        with pytest.raises(RunEngineInterrupted):
+            RE(
+                optimise_axis_against_intensity(
+                    simpledae,
+                    alignment_param=param,
+                    rel_scan_ranges=[10.0],
+                    fit_method=SlitScan().fit(),
+                    fit_param="",
+                )
+            )
+
+        assert RE.state == "paused"
+
+        monkeypatch.setattr("builtins.input", lambda str: "2")
+        RE.resume()
+        assert scan.call_count == 2
+        assert RE.state == "idle"
+
+
+async def test_optimise_axis_against_intensity_pauses_then_rescans_on_resume_and_can_pause_again(
+    RE, prefix, simpledae, monkeypatch
+):
+    with (
+        patch("ibex_bluesky_core.devices.reflectometry.get_pv_prefix", return_value=prefix),
+        patch(
+            "ibex_bluesky_core.plans.reflectometry._autoalign.scan",
+            side_effect=[_fake_scan(), _fake_scan()],
+        ),
+        patch(
+            "ibex_bluesky_core.plans.reflectometry._autoalign.ensure_connected",
+            return_value=bps.null(),
+        ),
+    ):
+        param = ReflParameter(prefix=prefix, name="S1VG", changing_timeout_s=60)
+        await param.connect(mock=True)
+        monkeypatch.setattr("builtins.input", lambda str: "3")
+        with pytest.raises(RunEngineInterrupted):
+            RE(
+                optimise_axis_against_intensity(
+                    simpledae,
+                    alignment_param=param,
+                    rel_scan_ranges=[10.0],
+                    fit_method=SlitScan().fit(),
+                    fit_param="",
+                )
+            )
+
+        assert RE.state == "paused"
+
+        # pause again when prompted
+        monkeypatch.setattr("builtins.input", lambda str: "3")
+        with pytest.raises(RunEngineInterrupted):
+            RE.resume()