Updated test_plans

Author: thopkins32

src/blop/plans.py

@@ -1,20 +1,21 @@
 import functools
 import warnings
 from collections.abc import Callable, Generator, Mapping, Sequence
-from typing import Any, cast
+from typing import Any, TypeVar, cast
 
 import bluesky.plan_stubs as bps
 import bluesky.plans as bp
-from ax.api.types import TParameterization
 from bluesky.protocols import Movable, NamedMovable, Readable, Reading
 from bluesky.utils import Msg, MsgGenerator, plan
 from ophyd import Signal  # type: ignore[import-untyped]
 
 from .dofs import DOF
 from .protocols import OptimizationProblem
 
+_NamedMovableT = TypeVar("_NamedMovableT", bound=NamedMovable)
 
-def _unpack_for_list_scan(movables: Mapping[NamedMovable, Sequence[Any]]) -> list[NamedMovable | Any]:
+
+def _unpack_for_list_scan(movables: Mapping[_NamedMovableT, Sequence[Any]]) -> list[_NamedMovableT | Any]:
     """Unpack the movables and inputs into Bluesky list_scan plan arguments."""
     unpacked_list = []
     for movable, values in movables.items():
@@ -26,7 +27,7 @@ def _unpack_for_list_scan(movables: Mapping[NamedMovable, Sequence[Any]]) -> lis
 
 @plan
 def default_acquire(
-    movables: Mapping[NamedMovable, Sequence[Any]],
+    movables: Mapping[_NamedMovableT, Sequence[Any]],
     readables: Sequence[Readable] | None = None,
     *,
     per_step: bp.PerStep | None = None,
@@ -37,12 +38,10 @@ def default_acquire(
 
     Parameters
     ----------
-    movables: Mapping[NamedMovable, Sequence[Any]]
+    movables: Mapping[_NamedMovableT, Sequence[Any]]
         The movables to move and the inputs to move them to.
     readables: Sequence[Readable]
         The readables to trigger and read.
-    trials: dict[int, TParameterization]
-        A dictionary mapping trial indices to their suggested parameterizations. Typically only a single trial is provided.
     per_step: bp.PerStep | None = None
         The plan to execute for each step of the scan.
     **kwargs: Any
@@ -310,7 +309,7 @@ def per_step_background_read(
 
 @plan
 def acquire_with_background(
-    movables: Mapping[NamedMovable, Sequence[Any]],
+    movables: Mapping[_NamedMovableT, Sequence[Any]],
     readables: Sequence[Readable] | None = None,
     *,
     block_beam: Callable[[], MsgGenerator[None]],
@@ -322,7 +321,7 @@ def acquire_with_background(
 
     Parameters
     ----------
-    movables: Mapping[NamedMovable, Sequence[Any]]
+    movables: Mapping[_NamedMovableT, Sequence[Any]]
         The movables and the inputs to move them to.
     readables: Sequence[Readable] | None = None
         The readables that produce data to evaluate.

src/blop/tests/unit/test_plans.py

@@ -5,7 +5,7 @@
 from bluesky.run_engine import RunEngine
 
 from blop.dofs import DOF
-from blop.plans import acquire, acquire_with_background
+from blop.plans import acquire_with_background, default_acquire
 
 from .conftest import MovableSignal, ReadableSignal
 
@@ -15,20 +15,20 @@ def RE():
     return RunEngine({})
 
 
-def test_acquire_single_dof(RE):
-    dof = DOF(movable=MovableSignal("x1", initial_value=-1.0), search_domain=(-5.0, 5.0))
+def test_default_acquire_single_movable_readable(RE):
+    movable = MovableSignal("x1", initial_value=-1.0)
     readable = ReadableSignal("objective")
+    movable_and_input = {movable: [0.0]}
     with patch.object(readable, "read", wraps=readable.read) as mock_read:
         RE(
-            acquire(
-                readables=[readable],
-                dofs={"x1": dof},
-                trials={0: {"x1": 0.0}},
+            default_acquire(
+                movable_and_input,
+                [readable],
             )
         )
         assert mock_read.call_count == 1
 
-    assert dof.movable.read()["x1"]["value"] == 0.0
+    assert movable.read()["x1"]["value"] == 0.0
 
 
 def test_acquire_with_background(RE):
@@ -38,18 +38,17 @@ def block_beam():
     def unblock_beam():
         yield from bps.null()
 
-    dof = DOF(movable=MovableSignal("x1", initial_value=-1.0), search_domain=(-5.0, 5.0))
+    movable = MovableSignal("x1", initial_value=-1.0)
     readable = ReadableSignal("objective")
-
+    movable_and_input = {movable: [0.0]}
     with patch.object(readable, "read", wraps=readable.read) as mock_read:
         RE(
             acquire_with_background(
+                movables=movable_and_input,
                 readables=[readable],
-                dofs={"x1": dof},
-                trials={0: {"x1": 0.0}},
                 block_beam=block_beam,
                 unblock_beam=unblock_beam,
             )
         )
         assert mock_read.call_count == 2
-    assert dof.movable.read()["x1"]["value"] == 0.0
+    assert movable.read()["x1"]["value"] == 0.0