initial commit

Author: rerpha

src/ibex_bluesky_core/plans/reflectometry/_autoalign.py

@@ -17,7 +17,7 @@
 logger = logging.getLogger(__name__)
 
 
-def _print_and_log(msg):
+def _print_and_log(msg: str) -> None:
     logger.info(msg)
     print(msg)
 
@@ -96,8 +96,9 @@ def _optimise_axis_over_range(  # noqa: PLR0913 PLR0917
             A callback for what to do if the optimised value is not found to be sensible.
 
     Returns:
-        Tuple of [ISISCallbacks, success (bool)]. The ISISCallbacks are a reference to a set of standard callbacks,
-        including the fit. The success flag indicates whether the optimization was successful.
+        Tuple of [ISISCallbacks, success (bool)]. The ISISCallbacks are a reference to a set of
+        standard callbacks, including the fit.
+        The success flag indicates whether the optimization was successful.
 
     """
     _print_and_log(
@@ -158,7 +159,7 @@ def _optimise_axis_over_range(  # noqa: PLR0913 PLR0917
     return icc, True
 
 
-def optimise_axis_against_intensity(
+def optimise_axis_against_intensity(  # noqa: PLR0913
     dae: SimpleDae,
     alignment_param: NamedMovable[float],
     fit_method: FitMethod,
@@ -205,9 +206,9 @@ def optimise_axis_against_intensity(
         Instance of :obj:`ibex_bluesky_core.callbacks.ISISCallbacks`.
 
     """
-    problem_found_plan = problem_found_plan or (lambda: bps.null())
+    problem_found_plan = problem_found_plan or bps.null
 
-    logging.info(
+    logger.info(
         "Starting optimise_axis_against_intensity with param=%s, ranges=%s",
         alignment_param.name,
         rel_scan_ranges,
@@ -232,11 +233,11 @@ def optimise_axis_against_intensity(
             )
 
             if not all_ok:
-                logging.info("Problem found during _optimise_axis_over_range, restarting scan loop")
+                logger.info("Problem found during _optimise_axis_over_range, restarting scan loop")
                 break
 
         if all_ok:
-            logging.info(
+            logger.info(
                 "Finished optimise_axis_against_intensity with param=%s, ranges=%s",
                 alignment_param.name,
                 rel_scan_ranges,

tests/plans/test_reflectometry.py

@@ -19,6 +19,7 @@
 )
 from ibex_bluesky_core.plans.reflectometry._autoalign import (
     _check_parameter,
+    _optimise_axis_over_range,
 )
 
 
@@ -104,10 +105,6 @@ def plan(mock) -> Generator[Msg, None, None]:
 
     mock = MagicMock()
 
-    def _mock_optimise_axis_over_range(*args, **kwargs):
-        yield from bps.null()
-        return None, True
-
     with (
         patch("ibex_bluesky_core.devices.reflectometry.get_pv_prefix", return_value=prefix),
         patch("ibex_bluesky_core.plans.reflectometry._autoalign.scan", return_value=_fake_scan()),
@@ -120,10 +117,6 @@ def _mock_optimise_axis_over_range(*args, **kwargs):
             "ibex_bluesky_core.plans.reflectometry._autoalign._check_parameter",
             side_effect=["Test!", None],
         ),
-        patch(
-            "ibex_bluesky_core.plans.reflectometry._autoalign._optimise_axis_over_range",
-            new=_mock_optimise_axis_over_range,
-        ),
     ):
         param = ReflParameter(prefix=prefix, name="S1VG", changing_timeout_s=60)
 
@@ -143,10 +136,10 @@ def _mock_optimise_axis_over_range(*args, **kwargs):
 
 
 @pytest.mark.parametrize(
-    "param_value, problem_str",
+    ("param_value", "problem_str"),
     [
-        (-5, "Optimised value found to be to be outside, to the right, of scan range"),
-        (5, "Optimised value found to be to be outside, to the left, of scan range"),
+        (-5, "Optimised value found to be to be outside, to the left, of scan range"),
+        (5, "Optimised value found to be to be outside, to the right, of scan range"),
         (0, None),
     ],
 )
@@ -155,11 +148,9 @@ def test_alignment_param_value_outside_of_scan_range_returns_problem(param_value
     with (
         patch("lmfit.model.ModelResult") as mr,
     ):
-        mr.values = {"x0": param_value}
-
         assert (
             _check_parameter(
-                alignment_param_value=0.0,
+                alignment_param_value=param_value,
                 result=mr,
                 init_mot_pos=0.0,
                 rel_scan_range=1.0,
@@ -182,31 +173,26 @@ def my_check(model: ModelResult, param_val: float) -> str | None:
     with patch("lmfit.model.ModelResult") as mr:
         mr.values = {"x0": 0.0}
 
-        if problem:
-            with pytest.raises(ValueError):
-                _check_parameter(
-                    alignment_param_value=0.0,
-                    result=mr,
-                    init_mot_pos=0.0,
-                    rel_scan_range=1.0,
-                    is_good_fit=my_check,
-                )
-
-        else:  # Check that does not raise
-            _check_parameter(
-                alignment_param_value=0.0,
-                result=mr,
-                init_mot_pos=0.0,
-                rel_scan_range=1.0,
-                is_good_fit=my_check,
-            )
+        assert ("problem" if problem else None) == _check_parameter(
+            alignment_param_value=0.0,
+            result=mr,
+            init_mot_pos=0.0,
+            rel_scan_range=1.0,
+            is_good_fit=my_check,
+        )
 
 
-def test_that_if_no_problem_found_then_motor_is_moved_and_rezeroed(RE, prefix, simpledae):
+def test_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 and redefine this as 0
+    value then move the motor to it
     """
-    sp = 5.0
+
+    icc = MagicMock(spec=ISISCallbacks)
+    icc.live_fit.fit_result.params[""] = 0.0
+
+    def mock_scan(*a, **k):
+        yield from bps.null()
+        return icc
 
     with (
         patch("ibex_bluesky_core.devices.reflectometry.get_pv_prefix", return_value=prefix),
@@ -223,19 +209,24 @@ def test_that_if_no_problem_found_then_motor_is_moved_and_rezeroed(RE, prefix, s
             return_value=_plan_return_0(),
         ),
         patch(
-            "ibex_bluesky_core.plans.reflectometry._autoalign._optimise_axis_over_range",
-            return_value=sp,
+            "ibex_bluesky_core.plans.reflectometry._autoalign.scan",
+            new=mock_scan,
         ),
     ):
         param = ReflParameter(prefix=prefix, name="S1VG", changing_timeout_s=60)
 
         RE(
-            optimise_axis_against_intensity(
+            _optimise_axis_over_range(
                 simpledae,
                 alignment_param=param,
-                rel_scan_ranges=[0.0],
+                rel_scan_range=0.0,
                 fit_method=SlitScan().fit(),
                 fit_param="",
+                num_points=10,
+                periods=True,
+                save_run=True,
+                is_good_fit=lambda *a, **k: None,
+                problem_found_plan=bps.null,
             )
         )
 
@@ -247,11 +238,13 @@ def test_that_if_problem_found_and_type_1_then_re_scan(RE, prefix, simpledae, mo
     """Test that if a problem is found, and the user types 1, then rescan.
     Then if they type 2, moves to value.
     """
+    call_count = 0
 
     def counter(str: str):
-        counter.call_count += 1  # type: ignore
+        nonlocal call_count
+        call_count += 1  # type: ignore
 
-        if counter.call_count == 1:  # type: ignore
+        if call_count == 1:  # type: ignore
             return "1"
         else:
             return "2"
@@ -271,13 +264,8 @@ def counter(str: str):
             "ibex_bluesky_core.plans.reflectometry._autoalign.bps.rd",
             return_value=_plan_return_0(),
         ),
-        patch(
-            "ibex_bluesky_core.plans.reflectometry._autoalign._optimise_axis_over_range",
-            return_value=0,
-        ),
     ):
         param = ReflParameter(prefix=prefix, name="S1VG", changing_timeout_s=60)
-        counter.call_count = 0  # type: ignore
 
         monkeypatch.setattr("builtins.input", counter)
         RE(
@@ -297,13 +285,15 @@ def test_that_if_problem_found_and_type_random_then_re_ask(RE, prefix, simpledae
     """Test that if a problem is found, and the user types gibberish, then ask again.
     Then if they type 1, it rescans. If they type 2, it moves to value.
     """
+    call_count = 0
 
     def counter(str: str):
-        counter.call_count += 1  # type: ignore
+        nonlocal call_count
+        call_count += 1
 
-        if counter.call_count == 1:  # type: ignore
+        if call_count == 1:  # type: ignore
             return "platypus"
-        elif counter.call_count == 2:  # type: ignore
+        elif call_count == 2:  # type: ignore
             return "1"
         else:
             return "2"
@@ -323,20 +313,15 @@ def counter(str: str):
             "ibex_bluesky_core.plans.reflectometry._autoalign.bps.rd",
             return_value=_plan_return_0(),
         ),
-        patch(
-            "ibex_bluesky_core.plans.reflectometry._autoalign._optimise_axis_over_range",
-            return_value=0,
-        ),
     ):
         param = ReflParameter(prefix=prefix, name="S1VG", changing_timeout_s=60)
-        counter.call_count = 0  # type: ignore
 
         monkeypatch.setattr("builtins.input", counter)
         RE(
             optimise_axis_against_intensity(
                 simpledae,
                 alignment_param=param,
-                rel_scan_ranges=[0.0],
+                rel_scan_ranges=[10.0],
                 fit_method=SlitScan().fit(),
                 fit_param="",
             )

tests/test_plan_stubs.py

@@ -11,6 +11,7 @@
     CALL_QT_AWARE_MSG_KEY,
     call_qt_aware,
     call_sync,
+    prompt_user_for_choice,
 )
 from ibex_bluesky_core.run_engine._msg_handlers import call_sync_handler
 
@@ -123,3 +124,12 @@ def plan():
         RE(plan())
 
     mock.assert_not_called()
+
+
+def test_get_user_input(RE):
+    with patch("ibex_bluesky_core.plan_stubs.input") as mock_input:
+        mock_input.__name__ = "mock"
+        mock_input.side_effect = ["foo", "bar", "baz"]
+
+        result = RE(prompt_user_for_choice(prompt="choice?", choices=["bar", "baz"]))
+        assert result.plan_result == "bar"