From 08ce47de9c4b31dcecd61077e93e8e3b8a1402d4 Mon Sep 17 00:00:00 2001 From: Tom Cobb Date: Tue, 10 Nov 2026 12:18:13 +0000 Subject: [PATCH] Rework StandardMovable and StandardFlyable around a single logic, add fly() and ensure_stopped() StandardMovable and StandardFlyable take one `logic` property instead of movable_logic and flyable_logic (ADR 0021). FlyableLogic's hooks become a prepare_to_fly() generator, a fly() generator that yields KICKED_OFF, and an abstract ensure_stopped() (ADR 0022). Standard devices raise TransitionError for a verb in the wrong state and StoppedError when stop() or a new set() ends a move, prepare or fly early, and unstage() cancels whatever is running before it stops the hardware. The motor, pmac, panda and sim motor ports move onto it. The old detector keeps running on it unchanged. Assisted-by: Claude Opus 5.5 --- CHANGELOG.md | 45 + CLAUDE.md | 4 +- docs/conf.py | 6 +- .../decisions/0017-standard-movable.md | 4 + .../decisions/0021-single-logic-property.md | 82 ++ .../0022-fly-generator-and-ensure-stopped.md | 44 + .../explanations/declarative-vs-procedural.md | 4 +- docs/explanations/when-to-extend-movable.md | 7 +- docs/how-to/write-a-flyer.md | 148 ++- pyproject.toml | 3 +- src/ophyd_async/core/__init__.py | 11 +- src/ophyd_async/core/_device.py | 8 +- src/ophyd_async/core/_flyable.py | 432 ++++--- src/ophyd_async/core/_movable.py | 210 +-- src/ophyd_async/core/_readable.py | 4 +- src/ophyd_async/core/_signal.py | 54 +- src/ophyd_async/core/_standard_base.py | 165 ++- src/ophyd_async/core/_status.py | 82 +- src/ophyd_async/core/_utils.py | 87 +- src/ophyd_async/epics/demo/_motor.py | 5 +- src/ophyd_async/epics/motor.py | 145 ++- .../epics/pmac/_pmac_trajectory.py | 209 ++- src/ophyd_async/fastcs/panda/_fly_logic.py | 44 +- src/ophyd_async/sim/_blob_acquire_logic.py | 9 +- src/ophyd_async/sim/_motor.py | 144 ++- src/ophyd_async/tango/demo/_motor.py | 6 +- tests/unit_tests/core/test_cancel_swallow.py | 71 ++ tests/unit_tests/core/test_flyer.py | 1132 +++++++++++++++-- tests/unit_tests/core/test_movable.py | 272 +++- .../core/test_multi_derived_signal.py | 4 +- tests/unit_tests/core/test_readable.py | 26 +- tests/unit_tests/core/test_signal.py | 34 + .../unit_tests/core/test_standard_detector.py | 2 +- tests/unit_tests/core/test_status.py | 24 + tests/unit_tests/core/test_utils.py | 203 +++ .../core/test_watchable_async_status.py | 193 +++ .../unit_tests/epics/adcore/test_use_cases.py | 3 +- .../epics/pmac/test_pmac_trajectory.py | 66 +- tests/unit_tests/epics/test_motor.py | 196 ++- .../unit_tests/fastcs/panda/test_hdf_panda.py | 4 +- tests/unit_tests/fastcs/panda/test_trigger.py | 17 +- tests/unit_tests/sim/test_sim_motor.py | 200 ++- uv.lock | 78 +- 43 files changed, 3576 insertions(+), 911 deletions(-) create mode 100644 docs/explanations/decisions/0021-single-logic-property.md create mode 100644 docs/explanations/decisions/0022-fly-generator-and-ensure-stopped.md create mode 100644 tests/unit_tests/core/test_cancel_swallow.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e4b5f36a9..908c10b569 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,51 @@ ## Unreleased +### Breaking changes + +- `StandardMovable.movable_logic` and `StandardFlyable.flyable_logic` are replaced by a single `logic` property, declared by each mix-in with its own type, so a Device that moves and flies implements it once with a logic inheriting both `MovableLogic` and `FlyableLogic` (see ADR 0021). Rename your `movable_logic`/`flyable_logic` `@cached_property` to `logic`; no deprecated aliases are left. Forgetting it is now a `TypeError` at instantiation rather than an error the first time it is used. `StandardFlyable` is now abstract, use `FlyableLogic.with_device()` for an ephemeral flyer. +- `FlyableLogic` hooks are renamed and reshaped; a subclass that still defines `on_prepare`, `on_kickoff`, `on_complete`, `on_stage`, `on_unstage` or (unless it is also a `MovableLogic`) `stop` raises `TypeError` naming the replacement: + - `on_prepare(value)` is now `prepare_to_fly(name, value, previous_ctx)`, an async generator that yields the context first, then optionally `WatcherUpdate`s that report the slow part (such as moving to the start) to the watchers of `prepare()`. A logic with no context yields `None`. `name` is the name of the device and `previous_ctx` is the context of the last finished `prepare()` since `stage()`, or `None` for the first. `prepare()` keeps the context only once the generator has finished; it raises `TypeError` if the first yield is a `WatcherUpdate` or a later one is not, and `RuntimeError` if it yields nothing. The generator is closed if `prepare()` fails or is cancelled, so it must own anything it starts (for example with `async with AsyncStatus(...)`). + - `on_kickoff(ctx)` and `on_complete(ctx)` are replaced by one `fly(ctx)`, an async generator that `StandardFlyable` runs in the task of the status `complete()` returns. It yields `KICKED_OFF` once the scan is up to speed, which finishes `kickoff()`, then optionally `WatcherUpdate`s, which go to the watchers of `complete()`; `complete()` finishes when it returns. A failure before `KICKED_OFF` fails `kickoff()`, one after it fails `complete()`. It is cancelled on `unstage()`, when `complete()`'s status is cancelled and when `kickoff()` is cancelled before `KICKED_OFF`, so `fly` must wait for everything it starts (for example with `async with AsyncStatus(...)`). + - `stop()` is now the abstract `ensure_stopped()`, which must wait until the hardware has stopped. `on_stage()` is now `ensure_ready()`, which defaults to `ensure_stopped()` and is called on `stage()`. `unstage()` cancels a `prepare()` or `fly` in progress, then always calls `ensure_stopped()`. +- "Device was stopped." is now raised as `StoppedError`, a `RuntimeError`, so `except RuntimeError` still catches it. It is raised by `StandardMovable` and `StandardFlyable` themselves, not by the logic: a `MovableLogic.move()` may end early without an error when stopped, as `Motor` and `SimMotor` now do. +- `StandardMovable.stop()` ends the `set()` in progress, then calls `ensure_stopped()` without waiting for it to finish failing. `stop(success=True)` lets it finish normally, `stop(success=False)` fails it with `StoppedError`, even if its move has not started yet, in which case the hardware is not touched. A second `stop()` only calls `ensure_stopped()` again. It does not end a `prepare()` or fly, only `unstage()` does, but `ensure_stopped()` still halts the hardware, so a `stop()` during the run-up of `prepare()` can leave it finishing short of the start; a future `pause()` hook is the place to fix this. `StandardMovable.set()` is now a plain method returning a `WatchableAsyncStatus`. +- A new `set()` on a `StandardMovable` ends the one in progress, which fails with `StoppedError` saying it was superseded; the hardware is not stopped, the new move retargets it. +- `StandardMovable` is now stageable, but a `StandardMovable` alone does nothing on `stage()` or `unstage()`. A Device that also flies runs the `FlyableLogic`'s `ensure_ready()` on `stage()` and `ensure_stopped()` on `unstage()`, so it stops a `set()` in progress at both; the RunEngine already stops everything it set when a run ends or pauses. +- `Motor` sends `.STOP` and waits for `.DMOV` every time `ensure_stopped()` is called, even if the motor looks stopped, as `.DMOV` can flicker at standstill on some motors; so every staged `Motor` gets one `.STOP` at `unstage()`. `Motor.stage()` stops a motion in progress rather than waiting for it. +- `FlyableLogic.prepare_to_fly` and `FlyableLogic.fly` are annotated as returning `AsyncGenerator`; annotate overrides the same way. +- `complete()` fails with `TypeError` if `fly` yields `KICKED_OFF` more than once, as for any other yield that is not a `WatcherUpdate`. +- `MovableLogic.stop()` is now the abstract `ensure_stopped()`, which must wait until the motion has stopped, and `StandardMovable.stop()` returns once it has. A subclass that still defines `stop()` gets a `DeprecationWarning`, and its `stop()` is used as `ensure_stopped()`. `Motor` now waits for `.DMOV` after writing `.STOP`, allowing its acceleration time to decelerate, and `InstantMotorMock` finishes a move on `.STOP`. +- `ophyd_async` now needs `typing-extensions>=4.14`, for `KICKED_OFF`. Type-checking a `fly` annotated with `KICKED_OFF` needs pyright 1.1.410 or later. +- Renamed `MotorFlyCtx` to `MotorPrepareCtx`, `PmacFlyCtx` and `PmacPrepareContext` to `PmacPrepareCtx`, and `SimMotorMoveLogic` to `SimMotorLogic`. +- A flyer must now be staged before it is prepared: `prepare()` is only allowed after `stage()`, `kickoff()` after `prepare()` and `complete()` after `kickoff()`, and a verb in the wrong state raises `TransitionError` (a `RuntimeError`). Staging an already-staged flyer, even one with a verb running, used to silently reset it and now raises (see below). A fly that fails leaves the flyer staged. `complete()` returns the same status each time it is called after a `kickoff()`, until the next `prepare()` or `stage()`. After `complete()` a flyer can be prepared again without restaging. Stock bluesky plans that stage their devices (`bp.count`, `bp.scan`) keep working; `bp.fly` neither stages nor prepares, so stage the flyer and prepare it before `bp.fly`, for example with `bpp.stage_wrapper`. +- Staging any already-staged Standard device (`StandardReadable`, `StandardMovable`, `StandardFlyable`, ...) now raises `TransitionError`, as ophyd raises `RedundantStaging`. For example, nesting `bpp.stage_wrapper` around a plan that also stages the same device now fails. +- `unstage()` now cancels the `prepare()` or fly that is running. `unstage()` waits for it before calling `ensure_stopped()`, then raises any error it raised while being cancelled. +- `MovableLogic` and `FlyableLogic` now share the base `Logic`, an `ABC`, so a subclass missing an abstract hook (`ensure_stopped`, `prepare_to_fly` or `fly`) raises `TypeError` when it is constructed. See ADR 0022 for the reasoning behind these hooks. +- `SimMotor` is now a `StandardFlyable`, so like any flyer it must be staged before `prepare()`. +- `SimMotor.stage()` now calls `SimMotorLogic.ensure_ready()`, which calls `ensure_stopped()` and so stops the simulated motion, waits for it to finish, and sets the setpoint to the current readback. `unstage()` cancels a fly in progress, then calls `ensure_stopped()`. The simulated motion runs in the task of the verb that moves, so `stop()` during a `prepare()` or fly, or `stage()` during a `set()`, cancels that verb, which raises `CancelledError`. +- The progress of `StandardFlyable.prepare()` now comes from the `WatcherUpdate`s the logic's `prepare_to_fly` yields, replacing the heuristic that watched setpoint writes (#1448). +- `WatchableAsyncStatus` closes the async generator it drives as soon as it is cancelled or fails, rather than when the generator is collected, and a watcher that raises is logged and ignored rather than failing the status. + ### Added +- `KICKED_OFF`, the sentinel `FlyableLogic.fly` yields once the scan is up to speed. +- `watch_readback()` reports the progress of a move by watching a `MovableLogic` readback, for use from `FlyableLogic` hooks. +- `SimMotor.prepare()` and `SimMotor.complete()` report progress to watchers, as `Motor` does. +- `TransitionError`, a `RuntimeError` raised by Standard devices for a verb called in the wrong state. +- `StoppedError`, the `RuntimeError` raised when `stop()` or a new `set()` ends a move, prepare or fly early. +- `cancel_and_wait()` cancels tasks and waits for all of them to end, even through a cancel of its caller, then raises the error of the first task, in the order given, that ended with one. If its caller was cancelled, it raises `CancelledError` instead. - `plugins_with_proc()` adds the process plugin made from a `proc_suffix` to a detector's `plugins`, and raises `ValueError` if `plugins` already has a `"proc"`. Every areaDetector vendor detector and `ContAcqDetector` uses it, so `"proc"` is reserved: pass `proc_suffix` instead. +- `Logic`, the base of `MovableLogic` and `FlyableLogic`, with the abstract `ensure_stopped()`. `FlyableLogic` adds `ensure_ready()`, which defaults to it. + +### Fixed + +- `complete()` finishing after `unstage()` no longer makes the device look prepared or flown to a new fly. +- `StandardReadable.add_children_as_readables()` registers the children it adds in the order they were declared, so `hints`, `read()` and `describe()` no longer vary from run to run. +- `unstage()` now cancels a running fly, so the PMAC trajectory stops appending points and nothing is left running once the hardware is stopped. +- `unstage()` now ends a running `complete()`, such as for the PandA. +- A `prepare()` that fails or is cancelled now stops the move to the run-up start of a `Motor` or `SimMotor`. +- A failed `prepare()` no longer leaves the flyer prepared with the failed context, so there is nothing to kick off. +- `set_and_wait_for_other_value()` cancelled before the first value arrives, such as by `stop()` during a move, no longer leaves its value-watching task running. +- A signal `get_value()`, `set()` or `trigger()` cancelled just as it finishes no longer swallows the cancellation on Python 3.11. +- Leaving an `async with AsyncStatus(...)` block, `set_and_wait_for_other_value()` and the sim `BlobAcquireLogic.ensure_stopped()` no longer swallow a cancellation of the caller that arrives while they wait for the task they cancelled: they wait for that task to end, then raise the cancellation. diff --git a/CLAUDE.md b/CLAUDE.md index e1032b8b40..b4c8d61c40 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,7 +51,9 @@ tox -p # all envs in parallel (CI equivalent) ## Code conventions - **No `ABC` base needed** — `@abstractmethod` without `ABC`/`ABCMeta` is intentional; pyright enforces implementation statically. -- **A new `TypeVar` needs an entry in `docs/conf.py`.** Docs build with `nitpicky = True` and warnings-as-errors, and TypeVars never resolve as references, so every one is listed by qualified name (`ophyd_async.core._utils.T`, …) in the `nitpick_ignore` list near the end of `conf.py`. Miss it and `tox -e docs` fails on `reference target not found` while every other env stays green. Prefer declaring TypeVars in the `_utils.py` of their package, next to the existing ones, so the list stays in one place. +- **A new `TypeVar` needs an entry in `docs/conf.py`.** Docs build with `nitpicky = True` and warnings-as-errors, and TypeVars never resolve as references, so every one is listed by qualified name (`ophyd_async.core._utils.T`, …) in the `nitpick_ignore` list near the end of `conf.py`. Miss it and `tox -e docs` fails on `reference target not found` while every other env stays green. Generic TypeVars (`T`, `V`, …) live in the `_utils.py` of their package; a TypeVar specific to one class is declared in that class's module, next to it. +- **Plain `Enum` members take their own name as their value** (`UNSTAGED = "UNSTAGED"`), never `auto()`. This is for enums that do not represent a control system value. `StrictEnum`, `SubsetEnum` and `SupersetEnum` are different: their values are the strings the control system uses (e.g. `ON = "On"`), so they keep those. +- **A lone marker value is a `typing_extensions.Sentinel`** (`KICKED_OFF = Sentinel("KICKED_OFF")`, usable as its own type in hints), not a one-member `Enum`, unless it crosses a serialisation boundary or is compared `==` to plain values: `CALCULATE_TIMEOUT` stays a string because pydantic fields and YAML/JSON plan arguments carry it. ## Updating this guide diff --git a/docs/conf.py b/docs/conf.py index fcd2422bff..695b7f8ea4 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -171,9 +171,11 @@ def setup(app: application.Sphinx): "ophyd_async.epics.adcore._io.ADBaseIOT", "ophyd_async.epics.adcore._io.NDPluginBaseIOT", "ophyd_async.epics.adcore._io.NDPluginFileIOT", - "ophyd_async.epics.motor.MotorFlyCtx", + "ophyd_async.epics.motor.MotorPrepareCtx", "ophyd_async.epics.motor.MotorFlyableMovableLogic", - "ophyd_async.epics.pmac._pmac_trajectory.PmacFlyCtx", + "ophyd_async.epics.pmac._pmac_trajectory.PmacPrepareCtx", + "ophyd_async.sim._motor.SimMotorLogic", + "ophyd_async.sim._motor.SimMotorPrepareCtx", "ophyd_async.tango.core._base_device.T", "ophyd_async.tango.core._tango_transport.P", "ophyd_async.tango.core._tango_transport.R", diff --git a/docs/explanations/decisions/0017-standard-movable.md b/docs/explanations/decisions/0017-standard-movable.md index 3d70c39da9..4994dc3cfa 100644 --- a/docs/explanations/decisions/0017-standard-movable.md +++ b/docs/explanations/decisions/0017-standard-movable.md @@ -6,6 +6,10 @@ Date: 2026-03-10 Accepted +Partly superseded by [ADR 0021](0021-single-logic-property.md), which replaced `movable_logic` with a single `logic` property. + +The `stop()` hook below has since been replaced by an abstract `ensure_stopped()` that waits until the motion has stopped (see [ADR 0022](0022-fly-generator-and-ensure-stopped.md)). + ## Context Several device types share a common movement pattern: write a setpoint, then observe a diff --git a/docs/explanations/decisions/0021-single-logic-property.md b/docs/explanations/decisions/0021-single-logic-property.md new file mode 100644 index 0000000000..028ab0449b --- /dev/null +++ b/docs/explanations/decisions/0021-single-logic-property.md @@ -0,0 +1,82 @@ +# 21. One `logic` property per Device + +Date: 2026-08-18 + +## Status + +Accepted + +Supersedes part of ADR 0017, which introduced `movable_logic`. + +## Context + +A Device that both moves and flies had to expose the same logic object twice. `Motor` +built one `MotorFlyableMovableLogic` and returned it from a `movable_logic` property for +`StandardMovable` and a `flyable_logic` property for `StandardFlyable`, plus a +private `_logic` to avoid building it twice. Every future mix-in would add another name. + +The obstacle is that Python has no intersection type: there is no way to write "the logic +must be both a `MovableLogic[float]` and a `FlyableLogic[FlyMotorInfo, MotorFlyCtx]`" as a +single annotation. + +ADR 0017 also chose `@cached_property` stacked on `@abstractmethod` over an +`add_movable_logic()` method, on the grounds that the property is "checkable by static +analysis". That only half held. pyright did report a Device that forgot its logic, but +`functools.cached_property` does not forward `__isabstractmethod__` from the function it +wraps, so `StandardMovable.__abstractmethods__` was empty and at runtime the Device failed +with `AttributeError: 'NoneType' object has no attribute 'readback'` from inside +`Device.__init__`. + +## Decision + +Each mix-in declares the **same** attribute name, `logic`, with **its own** +required type: + +```python +class StandardMovable(Device, ..., Generic[SignalDatatypeT]): + @abstract_cached_property + @abstractmethod + def logic(self) -> MovableLogic[SignalDatatypeT]: ... + +class StandardFlyable(_StandardBase, ..., Generic[PrepareT, CtxT]): + @abstract_cached_property + @abstractmethod + def logic(self) -> FlyableLogic[PrepareT, CtxT]: ... +``` + +A Device inheriting both provides one implementation, and the type checker verifies it +against each base independently. The intersection is expressed by the logic class's own +MRO: `MotorFlyableMovableLogic` inherits both, so it satisfies both declarations, and the +concrete type is what callers see back. + +This catches, at type-check time, a logic that satisfies one mix-in but not the other, a +wrong datatype parameter, and a Device with no implementation. That last is also caught at +instantiation, by stacking `abstract_cached_property` on `@abstractmethod`: pyright reads +the `@abstractmethod`, and `abstract_cached_property` is a `cached_property` that sets +`__isabstractmethod__` on the descriptor so that `Device`'s metaclass collects it. It is +declared as returning +`cached_property[T]` rather than as a subclass, because a subclass breaks the covariance +check on the override and would reject the correct case too. + +Because `StandardFlyable` is now genuinely abstract, `FlyableLogic.with_device` builds +a concrete `_EphemeralFlyable`, which takes the logic in its `__init__`. + +`logic` is public, so mock helpers such as `InstantMovableMock` and subclasses can reach +the signals it holds. + +An alternative of a single `logics` collection searched by type was rejected: membership +cannot be proved statically, so it would reintroduce exactly the runtime-only failure +this change fixes. + +## Consequences + +- One property per Device however many mix-ins it combines; `Motor` loses two properties. +- Where the move and fly logic are genuinely unrelated implementations they must be + combined into one class, since one property returns one object. Nothing in-tree pays + this cost, as `MotorFlyableMovableLogic` was already that shape. +- Forgetting the logic now fails with `TypeError: Can't instantiate abstract class ... + with abstract method logic` rather than an `AttributeError` on `None`. +- `movable_logic` and `flyable_logic` are removed rather than deprecated. A shim would + only have helped code that *reads* the attribute: a subclass that still *defines* + `movable_logic` would override the shim and leave `logic` unimplemented, so the + Device would fail at instantiation anyway. diff --git a/docs/explanations/decisions/0022-fly-generator-and-ensure-stopped.md b/docs/explanations/decisions/0022-fly-generator-and-ensure-stopped.md new file mode 100644 index 0000000000..9c06490351 --- /dev/null +++ b/docs/explanations/decisions/0022-fly-generator-and-ensure-stopped.md @@ -0,0 +1,44 @@ +# 22. `fly` is a generator, and `ensure_stopped` waits + +Date: 2026-10-20 + +## Status + +Accepted + +Amends ADR 0017: `MovableLogic.stop()` is now `ensure_stopped()`, and the fly hooks are different from the ones that ADR describes. + +## Context + +`FlyableLogic` had separate `on_kickoff` and `on_complete` hooks, and `stop()` returned as soon as it had asked the hardware to stop. This allowed three kinds of bug. A scan started in `on_kickoff` could outlive `unstage()`, because nothing owned it. A new `stage()` or `set()` could start while the previous motion was still decelerating. A `set()` caught by `stop()` before its move had started was not ended, and ran on after it. + +## Decision + +**`prepare_to_fly()` and `fly()` are async generators.** `prepare_to_fly` yields the context first, then any progress; `fly` yields `KICKED_OFF` once flying, then any progress. Everything a scan starts lives inside the generator, so closing or cancelling it is enough to stop it. Writing it as `async with AsyncStatus(...)` plus `yield` makes the ownership visible in the code, and a failure before or after the first yield maps cleanly on to the verb that failed. Two hooks that shared state through the logic, or a hook that returned a `(context, progress)` tuple, could not guarantee this: the progress iterator was not closed if `prepare()` failed, and a context that was itself a tuple was ambiguous. A failed `prepare()` leaves the device staged with no context, so there is nothing to kick off and a retry is a first `prepare()`, with `previous_ctx` of `None`. The hook names are not bluesky verbs, so a logic passed to a plan stub does not look like a device. The first yield is checked at runtime, as a `WatcherUpdate` yielded first is an easy mistake that the type checker cannot see. + +**The fly runs in the foreground of the status that `complete()` returns.** `kickoff()` makes a `WatchableAsyncStatus` that drives `fly`, and returns once it has yielded `KICKED_OFF`; `complete()` returns that same status. It already replays the latest update to a late watcher, and `unstage()` cancels the running verb, so no separate task, update forwarding or cleanup path is needed. `WatchableAsyncStatus` closes the generator it drives as soon as it is cancelled, rather than when it is collected. Cancelling `complete()`'s status therefore cancels the fly, and a fly cleanup that fails is reported by that status instead of its `CancelledError`. + +**`ensure_stopped()` waits until the motion has ended.** `unstage()` cancels `fly`, then calls `ensure_stopped()`, so when `unstage()` returns nothing is running and the hardware is still. Returning early would leave the next `stage()` or `set()` racing the deceleration of the last one. + +**One stop future ends a set.** Each `set()` passes a fresh future to the `stop` argument of its `WatchableAsyncStatus`, which races the operation against the future: when the future resolves first the operation is cancelled and the future's value becomes the outcome; statuses still have no public `stop()`, as bluesky treats a status as a task the device stops. The device keeps one future, the latest `set()`'s. `stop()` resolves it with `None` for `stop(success=True)`, so the `set()` finishes normally, or with `StoppedError("Device was stopped.")`, a `RuntimeError`, which it raises. `RunEngine.pause` calls `stop(success=True)` on everything it has set and expects that to finish without an error. A new `set()` resolves the previous future with a `StoppedError` saying it was superseded, so only one `set()` runs at a time; its move is cancelled but the hardware is not stopped, as the new move retargets it. The future is resolved with a value rather than an exception, so a future nobody awaits raises nothing. `stop()` ends only a `set()`, never a `prepare()` or fly: those are ended by `unstage()` through the state machine, so a `stop()` cannot reach a verb it was not meant for. A stop that arrives before the move starts ends the `set()` without touching the hardware, and a second `stop()` finds the future already resolved and only calls `ensure_stopped()` again. `stop()` does not wait for the `set()` to finish failing before `ensure_stopped()`: the hardware should stop as soon as possible. Logics no longer track whether they were stopped: a `move()` may end early without an error when stopped, as an EPICS motor's put completion does, and the stop future decides the outcome. + +**Only a flyer stages, and `unstage()` always stops.** A `StandardMovable` alone does nothing on `stage()` or `unstage()`: `stop()` is how it is stopped. `ensure_ready()` is a hook of `FlyableLogic` only, which defaults to `ensure_stopped()`, and only `StandardFlyable` calls it, on `stage()`; a device that moves and flies therefore stops a `set()` in progress at `stage()`, no worse than `caput MOTOR.STOP`. `StandardFlyable.unstage()` cancels a `prepare()` or fly, then always calls `ensure_stopped()`, whether or not the device was prepared. Sparing a `set()` there would only defer its stop: the bluesky RunEngine stops every movable it set when a run ends or pauses. The cost is that every staged `Motor` gets one `.STOP` at `unstage()`, because `Motor.ensure_stopped()` sends it unconditionally: `.DMOV` can flicker at standstill on some motors, so it is not read to decide, and the motor record does not write `.VAL` from `.RBV` on `.STOP`. + +The resulting model: + +| Verb | Mix-in | Does | +|------|--------|------| +| `stage()` | `StandardFlyable` | calls `logic.ensure_ready()` | +| `unstage()` | `StandardFlyable` | cancels a `prepare()` or fly in progress, then calls `logic.ensure_stopped()` | +| `stop()` | `StandardMovable` | ends the `set()` via the stop future, then calls `logic.ensure_stopped()` | +| `set()` | `StandardMovable` | supersedes the `set()` in progress | + +The logic hooks `prepare_to_fly` and `fly` are async generators with a first-yield marker: the context for `prepare_to_fly`, `KICKED_OFF` for `fly`. + +## Consequences + +- The stage and unstage of the previous plan are the only thing that keeps a motion from overlapping the next one. `stage()` can start `ensure_ready()` while the `ensure_stopped()` from the previous `unstage()` is still running, and we accept that. We assume every plan stages its devices with the stage/unstage decorator, as the stock plans do, and the tests run plans back to back and after an interruption to show that this behaves. A plan that does not stage gets none of this. +- A `FlyableLogic` that also moves gets `ensure_stopped()` called on `stage()` by default, which stops a `set()` in progress. Override `ensure_ready()` if stopping is wrong there. +- A plan that fails while a `set()` or `complete()` is in flight sees bluesky raise a `FailedStatus` for it from `unstage()`, where it would have seen the plan's own exception first. +- `StandardFlyable` has no `stop(success)` of its own, and `stop()` on a device that moves and flies does not end a `prepare()` or fly. `ensure_stopped` still halts the hardware, so a `stop()` during the run-up (for example `bps.stop`, or a `RunEngine` pause of a device that was also `set()`, with `rewindable=False`) can leave `prepare()` finishing short of the start. We do not check the position, as motors have a tolerance and none is baked in here. A future `pause()` hook is the better place to fix this. +- `complete()` returns the fly status even after the fly has failed, until the next `prepare()` or `stage()`, so the failure is reported rather than a `TransitionError`. The last run is forgotten at `stage()` rather than `unstage()`, which `RE.stop`, abort or halt may have failed or cancelled. diff --git a/docs/explanations/declarative-vs-procedural.md b/docs/explanations/declarative-vs-procedural.md index a83025751b..f9481cce7b 100644 --- a/docs/explanations/declarative-vs-procedural.md +++ b/docs/explanations/declarative-vs-procedural.md @@ -6,8 +6,8 @@ Ophyd async has two styles of creating Devices, Declarative and Procedural. This The procedural style mirrors how you would create a traditional python class, you define an `__init__` method, add some class members, then call the superclass `__init__` method. In the case of ophyd async those class members are likely to be Signals and other Devices. For example, in the `ophyd_async.sim.SimMotor` we create its soft signal children in an `__init__` method: ```{literalinclude} ../../src/ophyd_async/sim/_motor.py -:start-after: class SimMotor( -:end-before: self._fly_info +:start-after: For usage when simulating a motor. +:end-before: super().__init__(name=name) ``` It is explicit and obvious, but verbose. It also allows you to embed arbitrary python logic in the creation of signals, so is required for making soft signals and DeviceVectors with contents based on an argument passed to `__init__`. It also allows you to use the [](#StandardReadable.add_children_as_readables) context manager which can save some typing. diff --git a/docs/explanations/when-to-extend-movable.md b/docs/explanations/when-to-extend-movable.md index 4862845c86..0f3178e46e 100644 --- a/docs/explanations/when-to-extend-movable.md +++ b/docs/explanations/when-to-extend-movable.md @@ -34,7 +34,8 @@ You only need to describe what is device-specific. ### Implementing MovableLogic Create a `@dataclass` subclass of [](#MovableLogic) that adds any extra signals your -device needs, and override the hook methods that differ from the defaults: +device needs, implement `ensure_stopped`, and override the hook methods that differ from +the defaults: ```{literalinclude} ../../src/ophyd_async/epics/demo/_motor.py :pyobject: DemoMotorMoveLogic @@ -44,9 +45,9 @@ The available hooks are: | Method | Default behaviour | Override to… | |--------|-------------------|--------------| -| `stop()` | no-op | trigger a stop signal or cancel an in-flight command | +| `ensure_stopped()` | *abstract* | stop the motion and wait until it has stopped | | `check_move(new)` | no-op | validate soft limits before the move begins | -| `calculate_timeout(old, new)` | `DEFAULT_TIMEOUT` | derive a velocity-based or distance-based timeout | +| `calculate_timeout(old, new)` | `None`, so no timeout | derive a velocity-based or distance-based timeout | | `get_units_precision()` | reads from `readback.describe()` | supply units and precision from dedicated signals | | `move(new_position, timeout)` | `set_and_wait_for_other_value` | write the setpoint and wait in a device-specific way | diff --git a/docs/how-to/write-a-flyer.md b/docs/how-to/write-a-flyer.md index 344e0f76c3..4f9b28cd59 100644 --- a/docs/how-to/write-a-flyer.md +++ b/docs/how-to/write-a-flyer.md @@ -4,96 +4,128 @@ A *flyer* is a Device that hardware-triggers a scan: bluesky prepares it, kicks off, then waits for it to complete while the hardware runs autonomously. Rather than implement the `Preparable`, `Flyable`, and `Stageable` protocols by hand, compose a [](#FlyableLogic) instance into a [](#StandardFlyable). `StandardFlyable` drives the -three phases in order and threads a context object between them: - -- `prepare(value)` → `FlyableLogic.on_prepare(value)` moves to the start and arms the - hardware, returning a context object. -- `kickoff()` → `FlyableLogic.on_kickoff(ctx)` starts the scan, returning the - (possibly updated) context. -- `complete()` → `FlyableLogic.on_complete(ctx)` blocks until the scan is done. - -`stage()`/`unstage()` run the logic's `on_stage`/`on_unstage`, which default to `stop`. +phases in order and threads a context object between them: + +- `prepare(value)` → `FlyableLogic.prepare_to_fly(name, value, previous_ctx)`, an async generator that works + out the scan and yields a context object, then sets the hardware up and moves to the + start. It may report progress. +- `kickoff()` → starts `FlyableLogic.fly(ctx)`, an async generator that `StandardFlyable` + runs in the task of a status. `kickoff()` returns when it yields `KICKED_OFF`, once the + scan is up to speed. +- `complete()` returns that status, which finishes when `fly` does and reports the + progress it yields after `KICKED_OFF`. + +The flyer must be staged first: `prepare()` is only allowed after `stage()`, `kickoff()` +only after `prepare()`, and `complete()` only after `kickoff()`; a verb called out of +order raises `TransitionError`. After `complete()` the flyer can `prepare()` again for the +next point without restaging. `stage()` raises unless the flyer is unstaged, as does a +second `prepare()` before `complete()`. If `fly` fails before yielding `KICKED_OFF`, `kickoff()` +fails; after it, `complete()` fails. Either leaves the flyer staged. Cancelling +`complete()`'s status cancels `fly`, and so does cancelling `kickoff()` before +`KICKED_OFF`. + +`stage()` runs the logic's `ensure_ready`, which defaults to `ensure_stopped`. +`unstage()` works from any state: it cancels a `prepare()` or `fly` in flight (which sees +`CancelledError`) and waits for it, then always runs `ensure_stopped` to stop the +hardware. If the cancelled verb raised while being cancelled, `unstage()` raises it after `ensure_stopped`. On a device that also moves, +`stop()` ends only a `set()`: it does not end a `prepare()` or `fly`, but its +`ensure_stopped` halts the hardware, so a `stop()` during the run-up can leave +`prepare()` finishing short of the start. The hook names are deliberately not bluesky verbs, so a +logic passed to a plan stub does not look like a device. ## Implementing FlyableLogic Subclass [](#FlyableLogic) — it is generic over two types: `PrepareT`, the per-scan -info passed to `prepare()`, and `CtxT`, the context threaded between the phases (use -`None` if there is no state to carry). Fill in the hooks that your hardware needs: +info passed to `prepare()`; and `CtxT`, the context `prepare_to_fly` passes to `fly`. +`StandardFlyable` holds the context until the next `stage()`, and passes it to the next +`prepare_to_fly` as `previous_ctx` (`None` for the first since `stage()`). It defaults +to `None`, so `FlyableLogic[PrepareT]` carries no state and its `prepare_to_fly` yields +`None` first. `prepare()` keeps the context only if `prepare_to_fly` finishes, so a +failed one leaves nothing to kick off. Fill in the hooks that your hardware needs: | Method | Default behaviour | Override to… | |--------|-------------------|--------------| -| `on_prepare(value)` | *abstract* | move to the start, arm the hardware, and return a context | -| `on_kickoff(ctx)` | *abstract* | start the scan and return the context for `on_complete` | -| `on_complete(ctx)` | *abstract* | block until the scan finishes | -| `stop()` | no-op | disarm the hardware and wait for it to stop | -| `on_stage()` | calls `stop()` | set the hardware up on `stage()` if it differs from stop | -| `on_unstage()` | calls `stop()` | clean the hardware up on `unstage()` if it differs from stop | +| `prepare_to_fly(name, value, previous_ctx)` | *abstract* | as an async generator: yield the context first, then set up the hardware, yielding any progress (see [](#flyer-progress)) | +| `fly(ctx)` | *abstract* | run the scan as an async generator: yield `KICKED_OFF` once up to speed, then any progress, and return when it is done | +| `ensure_stopped()` | *abstract* | stop the hardware and wait until it has stopped | +| `ensure_ready()` | calls `ensure_stopped()` | set the hardware up on `stage()` if it differs from stopping it | + +`prepare_to_fly` and `fly` must not leave anything running that they do not wait for, +because `StandardFlyable` cancels them on `unstage()` and expects that to stop +everything they started. Own a status with `async with AsyncStatus(...) as status:`, which cancels it on +the way out, and run concurrent work with [](#watch_readback) or an +`asyncio.TaskGroup`. There are two ways to get a Device out of a `FlyableLogic`: mix `StandardFlyable` into a Device class (for hardware that is always a flyer, like a motor), or call [](#FlyableLogic.with_device) to wrap a bare logic object in an ephemeral flyer inside a plan (for hardware configured per scan, like a PandA block). +(flyer-progress)= +### Reporting progress + +`prepare()` and `complete()` can report progress to watchers as [](#WatcherUpdate)s, +stamped with the flyer's name. + +- `prepare_to_fly` yields `WatcherUpdate`s after the context, for example the progress + of the move to the start. They go to the watchers of `prepare()`. +- `fly` yields `WatcherUpdate`s after `KICKED_OFF`. They go to the watchers of + `complete()`, which also gets the latest one yielded before it was called. + ## A movable device that also flies An EPICS [](#Motor) is the first case: it is a [](#StandardMovable) *and* a `StandardFlyable` (and a [](#StandardReadable)). A single logic object, `MotorFlyableMovableLogic`, implements both [](#MovableLogic) and `FlyableLogic`, so its fly hooks can reuse the same `move`/`check_move` methods as a normal `set()`. The -per-scan info is a [](#FlyMotorInfo) and the threaded context is a small dataclass that -carries the move status from `on_kickoff` to `on_complete`: +per-scan info is a [](#FlyMotorInfo) and the threaded context is a small frozen +dataclass of everything `prepare_to_fly` works out: ```{literalinclude} ../../src/ophyd_async/epics/motor.py :language: python -:pyobject: MotorFlyCtx +:pyobject: MotorPrepareCtx ``` -`on_prepare` validates the requested velocity against the motor's limits, moves to the -run-up start position, and returns the context: +`prepare_to_fly` validates the requested velocity and the run-up and run-down positions +against the motor's limits and yields the context. It then owns the move to the run-up +start with `async with AsyncStatus(...)`, and yields its progress from +[](#watch_readback), which watches the readback of the logic until the move is done: ```{literalinclude} ../../src/ophyd_async/epics/motor.py :language: python -:pyobject: MotorFlyableMovableLogic.on_prepare +:pyobject: MotorFlyableMovableLogic.prepare_to_fly ``` -`on_kickoff` starts the move to the far end of the run-down and stashes the resulting -status on the context so `on_complete` can await it: +`fly` starts the move to the far end of the run-down, waits out the run-up so the motor +is at constant velocity when `kickoff()` returns, then yields the progress of the rest of +the move: ```{literalinclude} ../../src/ophyd_async/epics/motor.py :language: python -:pyobject: MotorFlyableMovableLogic.on_kickoff +:pyobject: MotorFlyableMovableLogic.fly ``` -`on_complete` simply blocks on that status: +`ensure_stopped` is shared with `MovableLogic`, so it serves both `stop()` and +`unstage()`. It sends `.STOP` every time and waits for the motor to be done moving: ```{literalinclude} ../../src/ophyd_async/epics/motor.py :language: python -:pyobject: MotorFlyableMovableLogic.on_complete +:pyobject: MotorFlyableMovableLogic.ensure_stopped ``` -Because the logic is also a `MovableLogic`, `prepare()` and `complete()` report progress -to watchers (e.g. progress bars) automatically by observing the readback — reusing the same -[](#WatcherUpdate) machinery as [](#StandardMovable.set). A flyer whose logic is *not* -movable (see the next section) just blocks with no progress updates. See +Opting in to progress is up to each hook: the same [](#WatcherUpdate) machinery as +[](#StandardMovable.set) is used, but a flyer whose logic is *not* movable (see the next +section) just waits with no progress updates. See [](../explanations/when-to-extend-movable.md) for more on `MovableLogic`. Wire the logic into the Device with a `@cached_property`. `Motor` builds one shared -instance and returns it from both `movable_logic` and `flyable_logic` so the move and -fly paths act on the same signals: +instance and returns it as its single `logic`, which both the movable and flyable +mix-ins declare (see [ADR 0021](../explanations/decisions/0021-single-logic-property.md) for why), so the move and fly paths act +on the same signals: ```{literalinclude} ../../src/ophyd_async/epics/motor.py :language: python -:pyobject: Motor._logic -``` - -```python -@cached_property -def movable_logic(self) -> MotorFlyableMovableLogic: - return self._logic - -@cached_property -def flyable_logic(self) -> MotorFlyableMovableLogic: - return self._logic +:pyobject: Motor.logic ``` ## An ephemeral flyer created in a plan @@ -101,37 +133,33 @@ def flyable_logic(self) -> MotorFlyableMovableLogic: Some hardware is only a flyer for the duration of a scan and is configured differently each time. The PandA position-compare block is an example: [](#StaticPcompFlyableLogic) wraps a single `PcompBlock` and carries no state between phases, so it is a -`FlyableLogic[PcompInfo, None]`. Its `on_prepare` disarms then programs the block from -the per-scan [](#PcompInfo): +`FlyableLogic[PcompInfo]`. Its `prepare_to_fly` yields its `None` context, then disarms +and programs the block from the per-scan [](#PcompInfo): ```{literalinclude} ../../src/ophyd_async/fastcs/panda/_fly_logic.py :language: python -:pyobject: StaticPcompFlyableLogic.on_prepare +:pyobject: StaticPcompFlyableLogic.prepare_to_fly ``` -`on_kickoff` enables the block and waits for it to go active, while `on_complete` waits -for it to go inactive again (the context is `None`, so it is ignored): - -```{literalinclude} ../../src/ophyd_async/fastcs/panda/_fly_logic.py -:language: python -:pyobject: StaticPcompFlyableLogic.on_kickoff -``` +`fly` enables the block, yields `KICKED_OFF` once it is active, then waits for it to go +inactive again. It reports no progress: ```{literalinclude} ../../src/ophyd_async/fastcs/panda/_fly_logic.py :language: python -:pyobject: StaticPcompFlyableLogic.on_complete +:pyobject: StaticPcompFlyableLogic.fly ``` -`stop()` disarms the block; it is called by the default `on_stage`/`on_unstage`: +`ensure_stopped` disarms the block and waits for it to go inactive; it is called on +`unstage()` and, through the default `ensure_ready`, on `stage()`: ```{literalinclude} ../../src/ophyd_async/fastcs/panda/_fly_logic.py :language: python -:pyobject: StaticPcompFlyableLogic.stop +:pyobject: StaticPcompFlyableLogic.ensure_stopped ``` Rather than declare this logic on a Device class, construct it in the plan from the -PandA block you want to drive and call [](#FlyableLogic.with_device) to get a ready -flyer to prepare, kickoff, and complete: +PandA block you want to drive and call [](#FlyableLogic.with_device) to get a flyer +to stage, prepare, kickoff, complete and unstage: ```python flyer = StaticPcompFlyableLogic(panda.pcomp[1]).with_device(name="pcomp") diff --git a/pyproject.toml b/pyproject.toml index b387728898..d9d7ca5652 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "stamina>=23.1.0", "scanspec>=1.0.0", "velocity-profile", + "typing-extensions>=4.14", ] dynamic = ["version"] license.file = "LICENSE" @@ -59,7 +60,7 @@ dev = [ "pre-commit", "pydata-sphinx-theme>=0.12", "pyepics>=3.4.2", - "pyright", + "pyright>=1.1.410", "pytest", "pytest-asyncio", "pytest-cov", diff --git a/src/ophyd_async/core/__init__.py b/src/ophyd_async/core/__init__.py index 1dde570c3d..74e1115149 100644 --- a/src/ophyd_async/core/__init__.py +++ b/src/ophyd_async/core/__init__.py @@ -51,7 +51,7 @@ OnOff, YesNo, ) -from ._flyable import FlyableLogic, FlyMotorInfo, StandardFlyable +from ._flyable import KICKED_OFF, FlyableLogic, FlyMotorInfo, StandardFlyable from ._log import config_ophyd_async_logging, logger from ._mock_signal_backend import MockSignalBackend from ._mock_signal_utils import ( @@ -75,6 +75,7 @@ MovableLogic, StandardMovable, TimeoutCalculator, + watch_readback, ) from ._path_providers import ( AutoIncrementFilenameProvider, @@ -133,6 +134,7 @@ make_datakey, ) from ._soft_signal_backend import SoftSignalBackend +from ._standard_base import Logic, StoppedError, TransitionError from ._status import ( AsyncStatus, WatchableAsyncStatus, @@ -152,6 +154,7 @@ SubsetEnum, SupersetEnum, WatcherUpdate, + cancel_and_wait, error_if_none, gather_dict, get_dtype, @@ -198,9 +201,12 @@ def __getattr__(name): "print_connection_status", # Movable "MovableLogic", + "StoppedError", + "Logic", "StandardMovable", "InstantMovableMock", "TimeoutCalculator", + "watch_readback", # Protocols "AsyncReadable", "AsyncConfigurable", @@ -232,6 +238,7 @@ def __getattr__(name): "SupersetEnum", "EnumTypes", "Table", + "TransitionError", "SignalMetadata", "Primitive", # Soft signal @@ -300,6 +307,7 @@ def __getattr__(name): "StandardFlyable", "FlyMotorInfo", "FlyableLogic", + "KICKED_OFF", # Settings "Settings", "SettingsProvider", @@ -314,6 +322,7 @@ def __getattr__(name): "ConfinedModel", "NotConnectedError", "Reference", + "cancel_and_wait", "error_if_none", "gather_dict", "get_dtype", diff --git a/src/ophyd_async/core/_device.py b/src/ophyd_async/core/_device.py index ac844259ff..b5f523a1af 100644 --- a/src/ophyd_async/core/_device.py +++ b/src/ophyd_async/core/_device.py @@ -3,7 +3,13 @@ import asyncio import os import sys -from collections.abc import Awaitable, Callable, Iterator, Mapping, MutableMapping +from collections.abc import ( + Awaitable, + Callable, + Iterator, + Mapping, + MutableMapping, +) from functools import cached_property from logging import LoggerAdapter, getLogger from typing import Any, Generic, TypeVar diff --git a/src/ophyd_async/core/_flyable.py b/src/ophyd_async/core/_flyable.py index 70fb85a2cb..3a5d22912b 100644 --- a/src/ophyd_async/core/_flyable.py +++ b/src/ophyd_async/core/_flyable.py @@ -1,92 +1,141 @@ import asyncio +import contextlib from abc import abstractmethod -from collections.abc import AsyncIterator -from enum import Enum +from collections.abc import AsyncGenerator, AsyncIterator +from dataclasses import replace from functools import cached_property -from typing import Generic, TypeVar, cast +from typing import Any, Generic, TypeVar, cast from bluesky.protocols import Flyable, Preparable from pydantic import Field +from typing_extensions import Sentinel +from typing_extensions import TypeVar as TypeVarWithDefault from ._movable import MovableLogic -from ._signal import observe_signals_value, observe_value -from ._standard_base import _StandardBase +from ._standard_base import ( + Logic, + TransitionError, + _StandardBase, + _State, + _StateMachine, +) from ._status import AsyncStatus, WatchableAsyncStatus from ._utils import ( CALCULATE_TIMEOUT, CalculatableTimeout, ConfinedModel, WatcherUpdate, + abstract_cached_property, + cancel_and_wait, ) -#: The per-scan info passed to `FlyableLogic.on_prepare` (often a +#: The per-scan info passed to `FlyableLogic.prepare_to_fly` (often a #: [`scanspec.core.Path`](inv:scanspec#scanspec.core.Path), but any type). PrepareT = TypeVar("PrepareT") -#: The context object threaded from `on_prepare` through `on_kickoff` to -#: `on_complete`. Use `None` for flyers that carry no state between stages. -CtxT = TypeVar("CtxT") +#: The context yielded by `FlyableLogic.prepare_to_fly` and passed on to `fly`. +#: Defaults to `None`, for flyers that carry no state from prepare to fly. +CtxT = TypeVarWithDefault("CtxT", default=None) + +#: Yielded by `FlyableLogic.fly` to mark the flyer as flying. +KICKED_OFF = Sentinel("KICKED_OFF") -class FlyableLogic(Generic[PrepareT, CtxT]): + +class FlyableLogic(Logic, Generic[PrepareT, CtxT]): """Minimum logic needed for controlling a `StandardFlyable`. Inherit and fill in the hooks for a particular flyable (e.g. a motion or trigger system). This base holds no state; concrete subclasses are typically `@dataclass`es that hold whatever signals or sub-devices they need as fields. - State that must be carried between stages is threaded through an explicit - context object rather than stored on the logic: `on_prepare` returns it, - `on_kickoff` receives and returns it, and `on_complete` receives it. - `StandardFlyable` owns that context and enforces the call ordering, so - subclasses do not need to guard against being called out of order. A flyer - with no cross-stage state uses `None` for the context (see the PandA trigger - logics). + What a scan needs is carried by a context object: `prepare_to_fly` yields it, + and it is passed to `fly` and to the next `prepare_to_fly`. `StandardFlyable` + holds it from a finished `prepare()` until the next `stage()` and enforces the + call ordering, so subclasses do not need to guard against being called out of + order. + + The type parameters are the per-scan info `PrepareT` and the context `CtxT`, + which defaults to `None`, so `FlyableLogic[PrepareT]` is a logic without one. """ - @abstractmethod - async def on_prepare(self, value: PrepareT) -> CtxT: - """Move to the start of the fly scan, set it up, and return its context. + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + renamed = { + "on_prepare": "prepare_to_fly", + "on_kickoff": "fly", + "on_complete": "fly", + "on_stage": "ensure_ready", + "on_unstage": "ensure_stopped", + # A logic that is also a MovableLogic gets its deprecation alias + **({} if issubclass(cls, MovableLogic) else {"stop": "ensure_stopped"}), + } + if stale := [name for name in renamed if name in cls.__dict__]: + raise TypeError( + f"{cls.__name__} defines {', '.join(stale)}, which FlyableLogic no " + "longer calls: " + + ", ".join(f"{name} -> {renamed[name]}" for name in stale) + ) - :param value: the per-scan info for this fly scan. - """ + async def ensure_ready(self) -> None: + """Get the device ready on `stage()`. Defaults to `ensure_stopped`.""" + await self.ensure_stopped() @abstractmethod - async def on_kickoff(self, ctx: CtxT) -> CtxT: - """Start the fly scan. - - :param ctx: the context returned by `on_prepare`. - - Return the (possibly updated) context to be passed to `on_complete`. + def prepare_to_fly( + self, name: str, value: PrepareT, previous_ctx: CtxT | None + ) -> AsyncGenerator[CtxT | WatcherUpdate, None]: + """Work out the scan, set up the hardware, and yield the scan's context. + + Write it as an async generator whose first yield is the context; any + later yields are `WatcherUpdate`s reporting the progress of the slow part, + e.g. moving to the start of the scan. `prepare()` finishes when the + generator does, and only then is the context kept. Own anything it starts + with `async with`, so it is stopped if `prepare()` fails or is cancelled: + + ```python + async def prepare_to_fly(self, name, value, previous_ctx): + ctx = MyCtx(start=value.start) + yield ctx + async with AsyncStatus(self.move(ctx.start, timeout)) as status: + async for update in watch_readback(self, status, target=ctx.start): + yield update + ``` + + :param name: the name of the device. + :param value: the per-scan info for this fly scan. + :param previous_ctx: the context of the last finished `prepare()` since + `stage()`, or `None` if there was none. """ @abstractmethod - async def on_complete(self, ctx: CtxT) -> None: - """Block until the fly scan is done. - - :param ctx: the context returned by `on_kickoff`. - """ - - async def stop(self) -> None: - """Stop/disarm the flyer and wait for everything to be stopped. - - Called by `on_stage` and `on_unstage` by default; override those instead - if stage and unstage need to differ. + def fly(self, ctx: CtxT) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + """Run the fly scan: yield `KICKED_OFF` once flying, then its progress. + + Write it as an async generator. `kickoff()` finishes at the first yield, + which must be `KICKED_OFF`; later yields go to the watchers of + `complete()`, which finishes when the generator does. + + `StandardFlyable` runs it in the task of the status `complete()` returns, + and cancels that task on `unstage()`. Do not leave anything + running that this generator does not wait for: own a status with + `async with`, and run concurrent work through `watch_readback` or an + `asyncio.TaskGroup`. For example: + + ```python + async def fly(self, ctx): + async with AsyncStatus(self.move(ctx.end, timeout)) as status: + await asyncio.sleep(ctx.acceleration_time) + yield KICKED_OFF + async for update in watch_readback(self, status, target=ctx.end): + yield update + ``` + + :param ctx: the context yielded by `prepare_to_fly`. """ - pass - async def on_stage(self) -> None: - """Set the flyer up on `stage()`. Defaults to `stop`.""" - await self.stop() - - async def on_unstage(self) -> None: - """Clean the flyer up on `unstage()`. Defaults to `stop`.""" - await self.stop() - - def with_device(self, name: str = "") -> "StandardFlyable": + def with_device(self, name: str = "") -> "StandardFlyable[Any, Any]": """Wrap this logic in an ephemeral `StandardFlyable` for use in a plan.""" - flyer = StandardFlyable(name=name) - flyer.flyable_logic = self - return flyer + return _EphemeralFlyable(self, name=name) class FlyMotorInfo(ConfinedModel): @@ -130,14 +179,6 @@ def _ramp_distance(self, acceleration_time: float) -> float: ) -class _FlyStage(Enum): - """Lifecycle stage of a `StandardFlyable`, used to enforce call ordering.""" - - IDLE = "IDLE" - PREPARED = "PREPARED" - KICKED_OFF = "KICKED_OFF" - - class StandardFlyable( _StandardBase, Preparable, @@ -146,137 +187,194 @@ class StandardFlyable( ): """Device that provides standard logic for flying. - This class must be inherited and have a `flyable_logic` @cached_property. + This class must be inherited and have a `logic` @cached_property. For an ephemeral flyer in a plan, call `FlyableLogic.with_device` instead of - inheriting. It owns the context threaded between the logic's stages and - enforces prepare -> kickoff -> complete ordering. `stage()`/`unstage()` run - the logic's `on_stage`/`on_unstage` (which default to `stop`) and reset the - context. + inheriting. It must be staged before it is prepared. `prepare()`, `kickoff()` + and `complete()` must then be called in that order, otherwise they raise + `TransitionError`; it can be prepared again only once complete. It holds the + context yielded by the logic's `prepare_to_fly` once that `prepare()` has + finished, until the next `stage()`. A `prepare()` that fails leaves no context. + + `kickoff()` starts the logic's `fly` generator in the task of a status and + finishes when it yields `KICKED_OFF`; `complete()` returns that status, which + finishes when the generator does and reports its later yields to watchers. + A failure before `KICKED_OFF` fails `kickoff()`, one after it fails + `complete()`; either leaves the device staged. Cancelling `complete()`'s + status cancels the fly, cancelling `kickoff()` before `KICKED_OFF` does too. + + `stage()` runs the logic's `ensure_ready`; it raises unless the device is + unstaged. On a device that is also a `StandardMovable` this stops a `set()` in + progress, as `ensure_ready` defaults to `ensure_stopped`. `unstage()` cancels a + running `prepare()` or fly, then always runs the logic's `ensure_stopped`, + which also stops a `set()` in progress. On a device that is also a + `StandardMovable`, `stop()` does not end a `prepare()` or fly, only + `unstage()` does, but it does halt the hardware, so it can leave `prepare()` + finishing short of the start. + `prepare()` and `complete()` report progress to watchers named after this + device; `kickoff()` reports none. """ + # The context of the last finished prepare(), until the next stage() + _ctx: CtxT | None = None + # The fly started by kickoff(), returned by complete() until the next prepare() + # or stage() + _fly_status: WatchableAsyncStatus | None = None + def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) - # Contribute the flyer's stage/unstage hooks (composes with any other - # _StandardBase mix-in, e.g. StandardReadable on a Motor). - self._stage_funcs += (self._on_stage,) - self._unstage_funcs += (self._on_unstage,) - # Context threaded prepare -> kickoff -> complete. Only meaningful once - # prepared; access is guarded by _fly_stage. - self._fly_ctx: CtxT = cast(CtxT, None) - self._fly_stage = _FlyStage.IDLE + # Composes with any other _StandardBase mix-in, e.g. StandardReadable + self._stage_funcs += (self.__on_stage,) + self._unstage_funcs += (self.__on_unstage,) @cached_property - def flyable_logic(self) -> FlyableLogic[PrepareT, CtxT]: + def _state(self) -> _StateMachine: + # super()._state caches the base machine in the instance first, then this + # property overwrites it with the combined one + return super()._state | { + _State.STAGED: {_State.PREPARING}, + _State.PREPARING: {_State.PREPARED, _State.STAGED}, + _State.PREPARED: {_State.FLYING}, + _State.FLYING: {_State.FLOWN, _State.STAGED}, + _State.FLOWN: {_State.PREPARING}, + } + + @abstract_cached_property + @abstractmethod + def logic(self) -> FlyableLogic[PrepareT, CtxT]: """The logic object that describes how this device flies. A static flyer (e.g. `Motor`) provides this as a `@cached_property` that builds a `FlyableLogic` from its signals. An ephemeral flyer created by - `FlyableLogic.with_device` has it set directly on the instance. + `FlyableLogic.with_device` returns the logic it was created with. + + A Device that is flyable *and* movable implements this once with a logic + object inheriting both `FlyableLogic` and `MovableLogic`. """ - raise NotImplementedError( - f"{type(self).__name__} has no flyable_logic; override it as a " - "@cached_property or create the flyer via FlyableLogic.with_device()" + raise NotImplementedError + + async def __on_stage(self) -> None: + # Forget the last run + self._ctx = None + self._fly_status = None + await self.logic.ensure_ready() + + async def __on_unstage(self) -> None: + # Always stop, a set() in progress included + await self.logic.ensure_stopped() + + def _prepared_ctx(self, verb: str) -> CtxT: + """The context of the last successful `prepare()`, for the called `verb`.""" + self._state.check_one_of( + _State.PREPARED, + _State.FLYING, + _State.FLOWN, + _State.TRIGGERING, + _State.TRIGGERED, + verb=verb, ) - - def _reset_fly_state(self) -> None: - self._fly_ctx = cast(CtxT, None) - self._fly_stage = _FlyStage.IDLE - - @AsyncStatus.wrap - async def _on_stage(self) -> None: - await self.flyable_logic.on_stage() - self._reset_fly_state() - - @AsyncStatus.wrap - async def _on_unstage(self) -> None: - await self.flyable_logic.on_unstage() - self._reset_fly_state() + return cast(CtxT, self._ctx) @WatchableAsyncStatus.wrap async def prepare(self, value: PrepareT) -> AsyncIterator[WatcherUpdate]: - """Move to the start and set up the fly scan. + """Set up the fly scan and move to its start. - If the logic is also a `MovableLogic` (e.g. a flying `Motor`), report - progress of the move to the start to watchers by observing its readback, - starting once `on_prepare` writes the setpoint it is moving to. + Reports the progress the logic's `prepare_to_fly` yields after the + context. """ - logic = self.flyable_logic - - async def on_prepare() -> None: - self._fly_ctx = await logic.on_prepare(value) - - if isinstance(logic, MovableLogic): - initial, (units, precision) = await asyncio.gather( - logic.readback.get_value(), logic.get_units_precision() - ) - current, target = initial, None - setpoint_seen = False - async with AsyncStatus(on_prepare()) as preparing: - async for signal, signal_value in observe_signals_value( - logic.readback, logic.setpoint, done_status=preparing - ): - if signal is logic.readback: - current = signal_value - elif setpoint_seen: - target = signal_value - else: - # The first setpoint value is the stale one from before the move - setpoint_seen = True - if target is not None: - yield WatcherUpdate( - current=current, - initial=initial, - target=target, - name=self.name, - unit=units, - precision=precision, + # A failed prepare must leave nothing to kick off + async with self._state.during( + _State.PREPARING, end=_State.PREPARED, on_error=_State.STAGED + ): + self._fly_status = None + previous_ctx, self._ctx = self._ctx, None + hook = f"{type(self.logic).__name__}.prepare_to_fly()" + steps = self.logic.prepare_to_fly(self.name, value, previous_ctx) + async with contextlib.aclosing(steps): + try: + ctx = await anext(steps) + except StopAsyncIteration: + raise RuntimeError( + f"{hook} finished without yielding its context" + ) from None + if isinstance(ctx, WatcherUpdate): + raise TypeError(f"{hook} must yield its context first, not {ctx!r}") + async for update in steps: + if not isinstance(update, WatcherUpdate): + raise TypeError( + f"{hook} must yield only WatcherUpdates after its " + f"context, not {update!r}" ) - else: - await on_prepare() - self._fly_stage = _FlyStage.PREPARED + yield replace(update, name=self.name) + self._ctx = ctx @AsyncStatus.wrap async def kickoff(self) -> None: - """Start the fly scan.""" - if self._fly_stage is not _FlyStage.PREPARED: - raise RuntimeError( - f"{self.name}: prepare() must be called before kickoff()" + """Start the logic's `fly`, returning once it has yielded `KICKED_OFF`.""" + kicked_off: asyncio.Future[None] = asyncio.get_running_loop().create_future() + status = WatchableAsyncStatus(self._fly(kicked_off), name=self.name) + try: + await asyncio.wait( + [kicked_off, status.task], return_when=asyncio.FIRST_COMPLETED ) - self._fly_ctx = await self.flyable_logic.on_kickoff(self._fly_ctx) - self._fly_stage = _FlyStage.KICKED_OFF + except BaseException: + # Cancelled, even just after KICKED_OFF: end the fly, report our cancel + with contextlib.suppress(Exception): + await cancel_and_wait(status.task) + raise + if not kicked_off.done(): + # The fly ended before KICKED_OFF, raise why + await status + self._fly_status = status + + async def _fly( + self, kicked_off: asyncio.Future[None] + ) -> AsyncIterator[WatcherUpdate]: + async with self._state.during( + _State.FLYING, end=_State.FLOWN, on_error=_State.STAGED + ): + hook = f"{type(self.logic).__name__}.fly()" + fly = self.logic.fly(self._prepared_ctx("kickoff")) + async with contextlib.aclosing(fly): + try: + first = await anext(fly) + except StopAsyncIteration: + raise RuntimeError( + f"{hook} finished without yielding KICKED_OFF" + ) from None + if first is not KICKED_OFF: + raise TypeError( + f"{hook} must yield KICKED_OFF first, not {first!r}" + ) + kicked_off.set_result(None) + async for update in fly: + if not isinstance(update, WatcherUpdate): + raise TypeError( + f"{hook} must yield only WatcherUpdates after " + f"KICKED_OFF, not {update!r}" + ) + yield replace(update, name=self.name) - @WatchableAsyncStatus.wrap - async def complete(self) -> AsyncIterator[WatcherUpdate]: - """Block until the fly scan is done. + def complete(self) -> WatchableAsyncStatus[Any]: + """Return the status of the fly the last `kickoff()` started. - If the logic is also a `MovableLogic` (e.g. a flying `Motor`), report - progress to watchers by observing its readback while the fly scan runs, - reusing the same watcher-update stream as `StandardMovable.set`. Flyers - whose logic is not movable simply block with no progress updates. + It reports the progress `fly` yields after `KICKED_OFF`, starting with the + latest, and finishes when `fly` does. Cancelling it cancels the fly. """ - if self._fly_stage is not _FlyStage.KICKED_OFF: - raise RuntimeError( - f"{self.name}: kickoff() must be called before complete()" + if self._fly_status is None: + raise TransitionError( + f"{self.name}: cannot complete from {self._state.state.name}, " + "allowed only after kickoff()" ) - logic = self.flyable_logic - if isinstance(logic, MovableLogic): - initial, target, (units, precision) = await asyncio.gather( - logic.readback.get_value(), - logic.setpoint.get_value(), - logic.get_units_precision(), - ) - async with AsyncStatus(logic.on_complete(self._fly_ctx)) as completing: - async for current_position in observe_value( - logic.readback, done_status=completing - ): - yield WatcherUpdate( - current=current_position, - initial=initial, - target=target, - name=self.name, - unit=units, - precision=precision, - ) - else: - await self.flyable_logic.on_complete(self._fly_ctx) - self._fly_stage = _FlyStage.IDLE + return self._fly_status + + +class _EphemeralFlyable(StandardFlyable[PrepareT, CtxT]): + """A `StandardFlyable` with the logic it was created with.""" + + def __init__(self, logic: FlyableLogic[PrepareT, CtxT], name: str = "") -> None: + self._logic = logic + super().__init__(name=name) + + @cached_property + def logic(self) -> FlyableLogic[PrepareT, CtxT]: + return self._logic diff --git a/src/ophyd_async/core/_movable.py b/src/ophyd_async/core/_movable.py index beeb1a1ee8..9b0d72ce81 100644 --- a/src/ophyd_async/core/_movable.py +++ b/src/ophyd_async/core/_movable.py @@ -1,10 +1,9 @@ import asyncio import time +import warnings from abc import abstractmethod -from collections.abc import Callable -from dataclasses import dataclass, field -from enum import Enum -from functools import cached_property +from collections.abc import AsyncIterator, Callable +from dataclasses import dataclass, field, replace from typing import Generic from bluesky.protocols import ( @@ -16,16 +15,18 @@ Subscribable, ) -from ._device import Device, DeviceMock, default_mock_class +from ._device import DeviceMock, default_mock_class from ._mock_signal_utils import callback_on_mock_put, set_mock_value from ._signal import SignalR, SignalRW, observe_value, set_and_wait_for_other_value from ._signal_backend import SignalDatatypeT +from ._standard_base import Logic, StoppedError, _StandardBase from ._status import AsyncStatus, WatchableAsyncStatus from ._utils import ( CALCULATE_TIMEOUT, CalculatableTimeout, Callback, WatcherUpdate, + abstract_cached_property, ) @@ -48,21 +49,28 @@ def __call__(self) -> float | None: @dataclass -class MovableLogic(Generic[SignalDatatypeT]): +class MovableLogic(Logic, Generic[SignalDatatypeT]): """Minimum logic needed for controlling a `StandardMovable`. - Can be inherited to add specialised logic for stop, checking if a move is valid, - calculate a valid timeout for a move, and add units and precision. Sub classes can - also override the `get_move_status` and `move` methods if more control is needed - to provide a custom AsyncStatus and to handle the move logic. + Inherit it to say how to stop the motion, and optionally to check a move is + valid, calculate a timeout for a move, and add units and precision. Subclasses + can also override `move` if more control is needed over the move itself. """ setpoint: SignalRW[SignalDatatypeT] readback: SignalR[SignalDatatypeT] - async def stop(self) -> None: - """Optional hook to add logic on how to stop the motion.""" - pass + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if "stop" in cls.__dict__: + warnings.warn( + f"{cls.__name__}.stop() is deprecated: rename it to " + "ensure_stopped(), which must wait until the motion has stopped", + DeprecationWarning, + stacklevel=2, + ) + if "ensure_stopped" not in cls.__dict__: + cls.ensure_stopped = cls.__dict__["stop"] async def check_move(self, new_position: SignalDatatypeT) -> None: """Optional hook to validate the move. @@ -111,20 +119,47 @@ async def connect(self, device: "StandardMovable") -> None: """Mock signals to do an instant move on setpoint write.""" def _instant_move(value): - set_mock_value(device.movable_logic.readback, value) # Arrive instantly + set_mock_value(device.logic.readback, value) # Arrive instantly + + callback_on_mock_put(device.logic.setpoint, _instant_move) - callback_on_mock_put(device.movable_logic.setpoint, _instant_move) +async def watch_readback( + logic: MovableLogic[SignalDatatypeT], + status: AsyncStatus, + target: SignalDatatypeT, + initial: SignalDatatypeT | None = None, +) -> AsyncIterator[WatcherUpdate[SignalDatatypeT]]: + """Report the progress of a move by watching the readback of a logic. -class StopState(Enum): - NONE = "NONE" - USER_SUCCESS = "USER_SUCCESS" - USER_FAILURE = "USER_FAILURE" + Yields a `WatcherUpdate` for each readback value until the status is done, + then re-raises any error the status raised. Use it from a `FlyableLogic` + hook that moves the logic's device, returning it from the hook as its + progress. + + :param logic: the logic whose readback is watched, and which supplies the + units and precision. + :param status: the status of the move, which is left running. + :param target: the value the move is going to. + :param initial: where the move started from. Defaults to the readback now. + """ + if initial is None: + initial = await logic.readback.get_value() + units, precision = await logic.get_units_precision() + async for current in observe_value(logic.readback, done_status=status): + yield WatcherUpdate( + current=current, + initial=initial, + target=target, + name=logic.readback.name, + unit=units, + precision=precision, + ) @default_mock_class(InstantMovableMock) class StandardMovable( - Device, + _StandardBase, Locatable[SignalDatatypeT], Checkable[SignalDatatypeT], Stoppable, @@ -133,102 +168,101 @@ class StandardMovable( ): """Device that provides standard logic for moving. - This class must be inherited and have a `movable_logic` @cached_property. - If stop is called while moving, it will raise a RuntimeError. + This class must be inherited and have a `logic` @cached_property. + One `set()` runs at a time: a new one ends the one in progress, which raises + `StoppedError`. `stop()` ends the `set()` in progress, which raises + `StoppedError` unless it was `stop(success=True)`, then runs the logic's + `ensure_stopped` without waiting for the `set()` to finish. A + `StandardMovable` alone does nothing on `stage()` or `unstage()`; combined with + a `StandardFlyable`, `stage()` and `unstage()` run its `ensure_ready` and + `ensure_stopped`, which stop a `set()` in progress. + + `stop()` ends only a `set()`: it does not end a `prepare()` or fly of a + `StandardFlyable`, only `unstage()` does. `ensure_stopped` still halts the + hardware, so a `stop()` during the run-up of a `prepare()` (for example from + `bps.stop`, or a `RunEngine` pause of a device that was also `set()`, with + `rewindable=False`) can leave `prepare()` finishing short of the start. """ - # Whether set() should complete successfully or raise an error due to stop called or - # due to timeout. - _stop_state: Enum = StopState.NONE - _move_status: AsyncStatus | None = None + # Ends the latest set(): resolved by stop(), or by the next set() + _set_stop: asyncio.Future[Exception | None] | None = None - @cached_property + @abstract_cached_property @abstractmethod - def movable_logic(self) -> MovableLogic: + def logic(self) -> MovableLogic[SignalDatatypeT]: """The logic object that describes how this device moves. - This is intentionally public so that mock helpers (e.g. - `InstantMovableMock`) and subclasses can access the `setpoint` and - `readback` signals directly. Subclasses must implement this as a - `@cached_property` that returns a `MovableLogic` instance. + Subclasses must implement this as a `@cached_property` returning a + `MovableLogic`. + + A Device that is movable *and* flyable (e.g. `Motor`) implements this + once with a logic object that inherits both `MovableLogic` and + `FlyableLogic`. """ + raise NotImplementedError async def check_value(self, value: SignalDatatypeT) -> None: """Check the move is valid before doing it.""" - await self.movable_logic.check_move(value) + await self.logic.check_move(value) - @WatchableAsyncStatus.wrap - async def set( + def set( self, new_position: SignalDatatypeT, timeout: CalculatableTimeout = CALCULATE_TIMEOUT, - ): - """Move to the given value.""" - self._stop_state = StopState.NONE - old_position, (units, precision) = await asyncio.gather( - self.movable_logic.readback.get_value(), - self.movable_logic.get_units_precision(), + ) -> WatchableAsyncStatus[SignalDatatypeT]: + """Move to the given value, ending any `set()` in progress.""" + self._end_set(StoppedError(f"Device {self.name}: superseded by a new set()")) + self._set_stop = asyncio.get_running_loop().create_future() + return WatchableAsyncStatus( + self._set(new_position, timeout), + name=self.name, + stop=self._set_stop, ) - await self.movable_logic.check_move(new_position) + async def _set( + self, new_position: SignalDatatypeT, timeout: CalculatableTimeout + ) -> AsyncIterator[WatcherUpdate[SignalDatatypeT]]: + old_position = await self.logic.readback.get_value() + await self.logic.check_move(new_position) if timeout == CALCULATE_TIMEOUT: - move_timeout = await self.movable_logic.calculate_timeout( + move_timeout = await self.logic.calculate_timeout( old_position, new_position ) else: move_timeout = timeout + move = self.logic.move( + new_position=new_position, timeout=MoveTimeout(move_timeout) + ) + async with AsyncStatus(move) as move_status: + async for update in watch_readback( + self.logic, move_status, target=new_position, initial=old_position + ): + yield replace(update, name=self.name) + + async def stop(self, success: bool = False) -> None: + """End any `set()` in progress, then run `ensure_stopped`. + + A `prepare()` or fly of a `StandardFlyable` is not ended, only halted. + """ + self._end_set( + None if success else StoppedError(f"Device {self.name} was stopped.") + ) + await self.logic.ensure_stopped() - try: - async with AsyncStatus( - self.movable_logic.move( - new_position=new_position, timeout=MoveTimeout(move_timeout) - ) - ) as self._move_status: - async for current_position in observe_value( - self.movable_logic.readback, - done_status=self._move_status, - ): - yield WatcherUpdate( - current=current_position, - initial=old_position, - target=new_position, - name=self.name, - unit=units, - precision=precision, - ) - # Convert cancellation into the appropriate stop result. - # If stop(success=True) was called, complete without raising. - # If stop(success=False) was called, raise a RuntimeError below. - # If not stopped but times out, show the CancelledError from the timeout. - except asyncio.CancelledError: - if self._stop_state == StopState.NONE: - raise - finally: - self._move_status = None - - # Raise error if needed, reset stop state. - stop_state = self._stop_state - self._stop_state = StopState.NONE - if stop_state == StopState.USER_FAILURE: - raise RuntimeError(f"Device {self.name} was stopped.") - - async def stop(self, success=False): - """Request to stop moving and return immediately.""" - self._stop_state = StopState.USER_SUCCESS if success else StopState.USER_FAILURE - await self.movable_logic.stop() - if self._move_status: - self._move_status.task.cancel() + def _end_set(self, error: StoppedError | None) -> None: + if self._set_stop is not None and not self._set_stop.done(): + self._set_stop.set_result(error) def set_name(self, name: str, *, child_name_separator: str | None = None) -> None: super().set_name(name, child_name_separator=child_name_separator) # Readback should be named the same as its parent in read() - self.movable_logic.readback.set_name(name) + self.logic.readback.set_name(name) async def locate(self) -> Location[SignalDatatypeT]: """Return the current setpoint and readback of the device.""" setpoint, readback = await asyncio.gather( - self.movable_logic.setpoint.get_value(), - self.movable_logic.readback.get_value(), + self.logic.setpoint.get_value(), + self.logic.readback.get_value(), ) return Location(setpoint=setpoint, readback=readback) @@ -236,7 +270,7 @@ def subscribe_reading( self, function: Callback[dict[str, Reading[SignalDatatypeT]]] ) -> None: """Subscribe to reading.""" - self.movable_logic.readback.subscribe_reading(function) + self.logic.readback.subscribe_reading(function) subscribe = subscribe_reading @@ -244,4 +278,4 @@ def clear_sub( self, function: Callback[dict[str, Reading[SignalDatatypeT]]] ) -> None: """Unsubscribe.""" - self.movable_logic.readback.clear_sub(function) + self.logic.readback.clear_sub(function) diff --git a/src/ophyd_async/core/_readable.py b/src/ophyd_async/core/_readable.py index 8b6d33eda4..2e378bcb80 100644 --- a/src/ophyd_async/core/_readable.py +++ b/src/ophyd_async/core/_readable.py @@ -182,8 +182,8 @@ def add_children_as_readables( ) raise KeyError(error_msg) - new_keys = dict_copy.keys() ^ new_dict.keys() - new_values = [new_dict[key] for key in new_keys] + # Newly added children, in the order they were declared + new_values = [value for key, value in new_dict.items() if key not in dict_copy] flattened_values = [] for value in new_values: diff --git a/src/ophyd_async/core/_signal.py b/src/ophyd_async/core/_signal.py index 99ab1ce406..94796ec43d 100644 --- a/src/ophyd_async/core/_signal.py +++ b/src/ophyd_async/core/_signal.py @@ -1,12 +1,12 @@ from __future__ import annotations import asyncio -import contextlib import functools import inspect import time import warnings from collections.abc import AsyncGenerator, Callable +from contextlib import aclosing from typing import Any, Generic, TypeVar, cast from bluesky.protocols import ( @@ -477,13 +477,18 @@ async def observe_value( do_something_with(value) ``` """ - async for _, value in observe_signals_value( - signal, - timeout=timeout, - done_status=done_status, - done_timeout=done_timeout, - ): - yield value + # Closing this generator must close the inner one, so its unsubscribe runs + # then, not when the abandoned inner generator is finalised + async with aclosing( + observe_signals_value( + signal, + timeout=timeout, + done_status=done_status, + done_timeout=done_timeout, + ) + ) as updates: + async for _, value in updates: + yield value def _get_iteration_timeout( @@ -600,11 +605,14 @@ def __init__(self, match: SignalDatatypeT | Callable[[SignalDatatypeT], bool]): self._matcher_name = repr(match) async def _wait_for_value(self, signal: SignalR[SignalDatatypeT]): - async for value in observe_value(signal): - self._last_value = value - self.got_first_value.set() - if self._matcher(value): - return + # Closed on return, so the unsubscribe (which cancels any poll task) runs + # now, not when the abandoned generator is finalised after we have gone + async with aclosing(observe_value(signal)) as values: + async for value in values: + self._last_value = value + self.got_first_value.set() + if self._matcher(value): + return async def wait_for_value( self, signal: SignalR[SignalDatatypeT], timeout: float | None @@ -680,13 +688,11 @@ async def set_and_wait_for_other_value( """ # Start monitoring before the set to avoid a race condition checker = _ValueChecker(match_value) - wait_task = asyncio.create_task( + # The wait is cancelled when we exit, whether by error, cancel or success + async with AsyncStatus( checker.wait_for_value(match_signal, timeout=timeout) - ) - - # Put this in a try/except to ensure wait_task is cancelled when we exit - try: - # NOTE: this timeout races wait_task's own internal timeout (same + ) as wait_status: + # NOTE: this timeout races the wait's own internal timeout (same # duration, started moments earlier) - whichever elapses first wins. # That race is harmless because each path reports its own distinct, # helpful message. @@ -704,19 +710,13 @@ async def set_and_wait_for_other_value( if not wait_for_set_completion: # Just wait for the match_signal to be at value - await wait_task + await wait_status else: # Wait for both to complete, so if the set raises an exception it # is surfaced early. Wait for the status task so if either errors # the other is cancelled - await asyncio.gather(wait_task, status.task) + await asyncio.gather(wait_status.task, status.task) return status - except Exception: - # Make sure the wait is not left dangling if there was an error - wait_task.cancel() - with contextlib.suppress(asyncio.CancelledError): - await wait_task - raise async def set_and_wait_for_value( diff --git a/src/ophyd_async/core/_standard_base.py b/src/ophyd_async/core/_standard_base.py index 3c4d69a221..290c6d4169 100644 --- a/src/ophyd_async/core/_standard_base.py +++ b/src/ophyd_async/core/_standard_base.py @@ -1,29 +1,174 @@ import asyncio -from collections.abc import Callable +from abc import ABC, abstractmethod +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from contextlib import asynccontextmanager +from enum import Enum +from functools import cached_property from ._device import Device from ._protocol import AsyncStageable -from ._status import AsyncStatus, AsyncStatusBase +from ._status import AsyncStatus +from ._utils import _gather_or_cancel_all, cancel_and_wait + + +class _State(Enum): + """Every state of `_StandardBase`, `StandardFlyable` and `StandardDetector`. + + Each class in the stack adds the states it needs and the edges between them. + """ + + # _StandardBase + UNSTAGED = "UNSTAGED" + STAGED = "STAGED" + # StandardFlyable + PREPARING = "PREPARING" + PREPARED = "PREPARED" + FLYING = "FLYING" + FLOWN = "FLOWN" + # StandardDetector + TRIGGERING = "TRIGGERING" + TRIGGERED = "TRIGGERED" + + +#: States that mean a verb is running; leaving one is only allowed to the task +#: that entered it (or by `unstage()`, which cancels that task). +_IN_PROGRESS = frozenset({_State.PREPARING, _State.FLYING, _State.TRIGGERING}) + + +class TransitionError(RuntimeError): + """A verb was called in a state that does not allow it.""" + + +class _StateMachine: + """Current state of a device, the edges it may follow and the running verb.""" + + def __init__(self, device: Device, edges: Mapping[_State, set[_State]]) -> None: + self._device = device + self._edges = edges + self.state = _State.UNSTAGED + self._task: asyncio.Task | None = None + + def __or__(self, edges: Mapping[_State, set[_State]]) -> "_StateMachine": + """A new machine for the same device, with these edges added to ours.""" + return _StateMachine( + self._device, + {s: self._edges.get(s, set()) | edges.get(s, set()) for s in _State}, + ) + + def transition(self, to: _State) -> None: + """Move to `to`, or raise `TransitionError` if there is no such edge.""" + # Unstage is the cleanup path: works from any state + allowed = self._edges.get(self.state, set()) | {_State.UNSTAGED} + if to not in allowed: + raise TransitionError( + f"{self._device.name}: cannot go to {to.name} from " + f"{self.state.name}, allowed: " + f"{', '.join(sorted(s.name for s in allowed))}" + ) + if ( + self.state in _IN_PROGRESS + and to is not _State.UNSTAGED + and self._task is not asyncio.current_task() + ): + raise TransitionError( + f"{self._device.name}: cannot go to {to.name} from " + f"{self.state.name}, a verb is running in another task and " + "only unstage can end it" + ) + self._task = asyncio.current_task() if to in _IN_PROGRESS else None + self.state = to + + def running_verb(self) -> asyncio.Task | None: + """The task of the verb in progress, unless it is the current task.""" + running = None + if self.state in _IN_PROGRESS and self._task is not asyncio.current_task(): + running = self._task + return running + + def check_one_of(self, *states: _State, verb: str) -> None: + """Raise `TransitionError` naming `verb` unless the state is in `states`.""" + if self.state not in states: + raise TransitionError( + f"{self._device.name}: {verb} not allowed from {self.state.name}, " + f"allowed from: {', '.join(sorted(s.name for s in states))}" + ) + + @asynccontextmanager + async def during( + self, + state: _State, + *, + end: _State, + on_error: _State, + ) -> AsyncGenerator[None, None]: + """Run a verb in the in-progress `state`, then go to `end` or `on_error`. + + Safe to wrap `yield`s of an async generator. If `unstage()` has moved the + machine on in the meantime, even into a new run of the same verb, this + verb no longer owns the state and is left alone to finish cancelling. + """ + self.transition(state) + try: + yield + except BaseException: + # Can be false: unstage cancelled this verb and moved the state on, + # or an abandoned async generator was closed later from another task + if self.state is state and self._task is asyncio.current_task(): + self.transition(on_error) + raise + else: + if self.state is state and self._task is asyncio.current_task(): + self.transition(end) + + +class StoppedError(RuntimeError): + """A set, prepare or fly was ended early by `stop()` or a new `set()`.""" + + +class Logic(ABC): + """Base of `MovableLogic` and `FlyableLogic`: stopping.""" + + @abstractmethod + async def ensure_stopped(self) -> None: + """Stop the hardware and wait until it has stopped. + + Called by `StandardMovable.stop()`, which bluesky calls on pause and at + the end of a run, and by `StandardFlyable` on `unstage()`. + """ class _StandardBase(Device, AsyncStageable): """Base for `Device` mix-ins that contribute to `stage()`/`unstage()`. `StandardReadable`, `StandardFlyable` and similar mix-ins each need to run - something on `stage()`/`unstage()`. Rather than each defining those methods - and shadowing the others via the MRO when combined on one `Device` (e.g. a - motor that is Readable + Movable + Flyable), they register callables into - `_stage_funcs`/`_unstage_funcs`, which this single base gathers. + something on `stage()`/`unstage()`. They register callables into + `_stage_funcs`/`_unstage_funcs`, which this base gathers. Each mix-in names + its hooks `__on_stage`/`__on_unstage`, name-mangled so they never collide. """ # Immutable defaults to avoid accidental sharing between instances - _stage_funcs: tuple[Callable[[], AsyncStatusBase], ...] = () - _unstage_funcs: tuple[Callable[[], AsyncStatusBase], ...] = () + _stage_funcs: tuple[Callable[[], Awaitable[None]], ...] = () + _unstage_funcs: tuple[Callable[[], Awaitable[None]], ...] = () + + @cached_property + def _state(self) -> _StateMachine: + # cached_property as Device registration runs before super().__init__ + return _StateMachine(self, {_State.UNSTAGED: {_State.STAGED}}) @AsyncStatus.wrap async def stage(self) -> None: - await asyncio.gather(*(func().task for func in self._stage_funcs)) + self._state.transition(_State.STAGED) + await _gather_or_cancel_all([func() for func in self._stage_funcs]) @AsyncStatus.wrap async def unstage(self) -> None: - await asyncio.gather(*(func().task for func in self._unstage_funcs)) + verb = self._state.running_verb() + self._state.transition(_State.UNSTAGED) + # Made now, so a hook that is not an async def takes its state before a + # restage; an async def hook reads its state when it runs, below + unstaging = [func() for func in self._unstage_funcs] + try: + if verb is not None: + await cancel_and_wait(verb) + finally: + await _gather_or_cancel_all(unstaging) diff --git a/src/ophyd_async/core/_status.py b/src/ophyd_async/core/_status.py index 827df3f75e..42233e49bd 100644 --- a/src/ophyd_async/core/_status.py +++ b/src/ophyd_async/core/_status.py @@ -3,11 +3,16 @@ from __future__ import annotations import asyncio -import contextlib import functools import time -from asyncio import CancelledError -from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine +from asyncio import FIRST_COMPLETED, CancelledError +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Coroutine, +) from dataclasses import asdict, replace from typing import Any, Generic @@ -15,7 +20,7 @@ from ._device import Device from ._protocol import Watcher -from ._utils import Callback, P, T, V, WatcherUpdate +from ._utils import Callback, P, T, V, WatcherUpdate, cancel_and_wait class AsyncStatusBase(Status, Awaitable[T]): @@ -131,11 +136,7 @@ async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): - self.task.cancel() - # Need to await the task to suppress teardown warnings, but - # we know it will raise CancelledError as we just cancelled it - with contextlib.suppress(CancelledError): - await self.task + await cancel_and_wait(self.task) # Raise any errors from the block, but not cancellation errors from the status return False @@ -225,22 +226,65 @@ class WatchableAsyncStatus(AsyncStatusBase[None], Generic[T]): """ def __init__( - self, iterator: AsyncIterator[WatcherUpdate[T]], name: str | None = None + self, + iterator: AsyncIterator[WatcherUpdate[T]], + name: str | None = None, + *, + stop: asyncio.Future[Exception | None] | None = None, ): self._watchers: list[Watcher[T]] = [] self._start = time.monotonic() self._last_update: WatcherUpdate[T] | None = None - super().__init__(self._notify_watchers_from(iterator), name) + super().__init__( + self._notify_watchers_from(iterator) + if stop is None + else self._notify_until_stopped(iterator, stop), + name, + ) async def _notify_watchers_from(self, iterator: AsyncIterator[WatcherUpdate[T]]): - async for update in iterator: - self._last_update = ( - update - if update.time_elapsed is not None - else replace(update, time_elapsed=time.monotonic() - self._start) - ) - for watcher in self._watchers: - self._update_watcher(watcher, self._last_update) + try: + async for update in iterator: + self._last_update = ( + update + if update.time_elapsed is not None + else replace(update, time_elapsed=time.monotonic() - self._start) + ) + for watcher in self._watchers: + self._update_watcher(watcher, self._last_update) + except Exception: + # A watcher that raised left the iterator at its yield, so close it + # now rather than at garbage collection + if isinstance(iterator, AsyncGenerator): + await iterator.aclose() + raise + + async def _notify_until_stopped( + self, + iterator: AsyncIterator[WatcherUpdate[T]], + stop: asyncio.Future[Exception | None], + ) -> None: + """Notify watchers until the iterator ends or `stop` resolves. + + `stop` resolves with `None` to finish normally or with the error to raise. + """ + # Stopped before we started, so leave the hardware alone + stopped = stop.done() + if not stopped: + # A task of its own, so the race below can cancel it without cancelling us + notify = asyncio.create_task(self._notify_watchers_from(iterator)) + try: + await asyncio.wait({notify, stop}, return_when=FIRST_COMPLETED) + except asyncio.CancelledError: + await cancel_and_wait(notify) + # Its CancelledError says what it was waiting on + await notify + raise + # If it finished, even in the same iteration as the stop, that end wins + stopped = not notify.done() + await cancel_and_wait(notify) + if stopped and (error := stop.result()) is not None: + raise error def _update_watcher(self, watcher: Watcher[T], update: WatcherUpdate[T]): vals = asdict( diff --git a/src/ophyd_async/core/_utils.py b/src/ophyd_async/core/_utils.py index ca5f8806ca..a25159634b 100644 --- a/src/ophyd_async/core/_utils.py +++ b/src/ophyd_async/core/_utils.py @@ -4,7 +4,7 @@ import logging import time from collections.abc import ( - AsyncIterator, + AsyncGenerator, Awaitable, Callable, Iterable, @@ -13,7 +13,7 @@ ) from dataclasses import dataclass from enum import Enum, EnumMeta, StrEnum -from functools import lru_cache +from functools import cached_property, lru_cache from typing import ( Any, Generic, @@ -232,7 +232,7 @@ def cached_get_origin(tp: Any) -> Any: return get_origin(tp) -def get_dtype(datatype: type) -> np.dtype: +def get_dtype(datatype: Any) -> np.dtype: """Get the runtime dtype from a numpy ndarray type annotation. ```python @@ -307,9 +307,12 @@ def get_enum_cls(datatype: type | None) -> type[EnumTypes] | None: return None -def get_unique(values: dict[str, T], types: str) -> T: +def get_unique(values: Mapping[str, T], types: str) -> T: """If all values are the same, return that value, otherwise raise TypeError. + :param values: the values by label + :param types: what the values are, for the error message + ```python >>> get_unique({"a": 1, "b": 1}, "integers") 1 @@ -339,9 +342,20 @@ async def merge_gathered_dicts( combined_read = await merge_gathered_dicts(s.read() for s in signals) ``` """ + return merge_dicts(await asyncio.gather(*coros)) + + +def merge_dicts(dicts: Iterable[Mapping[str, T]]) -> dict[str, T]: + """Merge dictionaries, raising if any key appears in more than one. + + :example: + ```python + merged = merge_dicts([{"a": 1}, {"b": 2}]) + ``` + """ ret: dict[str, T] = {} duplicates: list[str] = [] - for result in await asyncio.gather(*coros): + for result in dicts: for key in result: if key in ret: duplicates.append(key) @@ -437,9 +451,68 @@ def non_zero(value): return int(value) != 0 +def abstract_cached_property(func: Callable[[Any], T]) -> cached_property[T]: + """A `cached_property` that `ABCMeta` treats as abstract. + + Stack it on `@abstractmethod`. pyright reads the `@abstractmethod`, but + `functools.cached_property` does not forward its `__isabstractmethod__`, so at + runtime the class would stay instantiable. Setting the flag on the descriptor + makes the metaclass collect it, so a subclass that does not implement it + raises `TypeError` at instantiation. + + :param func: The abstract property body. + """ + prop = cached_property(func) + prop.__isabstractmethod__ = True # type: ignore[attr-defined] + return prop + + +async def cancel_and_wait(*tasks: asyncio.Task[Any]) -> None: + """Cancel `tasks`, wait for all of them to end, then raise the first error. + + The first error is that of the first task, in the order given, that ended + with one. If the caller is cancelled while waiting, it keeps waiting until + every task has ended, then raises its own `CancelledError` instead. + + :param tasks: The tasks to cancel. + """ + for task in tasks: + task.cancel() + caller_cancel: asyncio.CancelledError | None = None + while not all(task.done() for task in tasks): + try: + await asyncio.wait(tasks) + except asyncio.CancelledError as cancel: + caller_cancel = cancel + # Retrieve every error, even when re-raising a cancel, so asyncio never logs + # "Task exception was never retrieved" + errors = [task.exception() for task in tasks if not task.cancelled()] + first_error = next((error for error in errors if error is not None), None) + if caller_cancel is not None: + raise caller_cancel + if first_error is not None: + raise first_error + + +async def _gather_or_cancel_all(waits: Sequence[Awaitable[Any]]) -> list[Any]: + """Await `waits` in parallel. + + If one raises, or this is cancelled, the rest are cancelled and waited for + before the error propagates, so no wait outlives the call. + """ + tasks = [asyncio.ensure_future(wait) for wait in waits] + try: + values = await asyncio.gather(*tasks) + finally: + await cancel_and_wait(*tasks) + return values + + async def _wait_for(coro: Awaitable[T], timeout: float | None, source: str) -> T: + """Await `coro`, raising `TimeoutError(source)` if it takes over `timeout`.""" try: - return await asyncio.wait_for(coro, timeout) + async with asyncio.timeout(timeout): + return await coro except TimeoutError as exc: raise TimeoutError(source) from exc @@ -449,7 +522,7 @@ async def simulate_move( new_position: float, velocity: float, acceleration_time: float, -) -> AsyncIterator[float]: +) -> AsyncGenerator[float, None]: """Yield positions at 10Hz along a trapezoidal velocity profile. :param old_position: Where the move starts. diff --git a/src/ophyd_async/epics/demo/_motor.py b/src/ophyd_async/epics/demo/_motor.py index 308bc4f99a..8b6e57c880 100644 --- a/src/ophyd_async/epics/demo/_motor.py +++ b/src/ophyd_async/epics/demo/_motor.py @@ -24,7 +24,8 @@ class DemoMotorMoveLogic(MovableLogic[float]): velocity: SignalRW[float] stop_: TriggerableCommand - async def stop(self): + async def ensure_stopped(self) -> None: + """Stop the motor, which stops as soon as `stop_` has processed.""" await self.stop_.trigger() async def calculate_timeout( @@ -55,7 +56,7 @@ class DemoMotor(EpicsDevice, StandardReadable, StandardMovable[float]): stop_: A[TriggerableCommand, PvSuffix("Stop.PROC")] @cached_property - def movable_logic(self) -> MovableLogic: + def logic(self) -> MovableLogic: return DemoMotorMoveLogic( readback=self.readback, setpoint=self.setpoint, diff --git a/src/ophyd_async/epics/motor.py b/src/ophyd_async/epics/motor.py index 62dbfd0765..cb8c7fbae4 100644 --- a/src/ophyd_async/epics/motor.py +++ b/src/ophyd_async/epics/motor.py @@ -6,13 +6,16 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncGenerator from dataclasses import dataclass from functools import cached_property from ophyd_async.core import ( CALCULATE_TIMEOUT, DEFAULT_TIMEOUT, + KICKED_OFF, AsyncStatus, + CalculatableTimeout, DeviceMock, FlyableLogic, FlyMotorInfo, @@ -25,11 +28,13 @@ StandardReadable, StrictEnum, TimeoutCalculator, + WatcherUpdate, callback_on_mock_put, default_mock_class, - error_if_none, set_mock_value, simulate_move, + wait_for_value, + watch_readback, ) from ophyd_async.core import StandardReadableFormat as Format from ophyd_async.epics.core import epics_signal_r, epics_signal_rw, epics_signal_w @@ -92,26 +97,30 @@ class UseSetMode(StrictEnum): """Change offset (in record or in controller) when setpoint is written.""" -@dataclass -class MotorFlyCtx: - """State threaded through a motor fly scan (prepare -> kickoff -> complete).""" +@dataclass(frozen=True) +class MotorPrepareCtx: + """What a motor fly scan needs after `prepare_to_fly` has worked it out.""" - #: The info the scan was prepared with, produced by `on_prepare`. - fly_info: FlyMotorInfo - #: The move to the end position, produced by `on_kickoff`. - status: AsyncStatus | None = None + #: Where to start from: the start position with the run-up distance added. + start: float + #: Where to end: the end position with the run-down distance added. + end: float + #: How long the motor takes to get up to speed. + acceleration_time: float + #: Maximum time for the move from start to end. + timeout: CalculatableTimeout @dataclass class MotorFlyableMovableLogic( - MovableLogic[float], FlyableLogic[FlyMotorInfo, MotorFlyCtx] + MovableLogic[float], FlyableLogic[FlyMotorInfo, MotorPrepareCtx] ): """Combined move + fly logic for a motor record. - A single logic object backs a `Motor`'s `movable_logic` and `flyable_logic`: + A single logic object backs a `Motor`'s `logic`: it implements `MovableLogic` (the motor-record move, limit and timeout logic) and `FlyableLogic` (the fly hooks, which reuse `move`/`check_move` and carry - all per-scan state in a `MotorFlyCtx` rather than on the logic). + all per-scan state in a `MotorPrepareCtx` rather than on the logic). """ motor_stop: SignalW[int] @@ -125,16 +134,18 @@ class MotorFlyableMovableLogic( max_velocity: SignalR[float] motor_egu: SignalR[str] - async def stop(self): - """Request to stop moving, but only if the motor is currently moving. + async def ensure_stopped(self) -> None: + """Stop the motor and wait until it is done moving. + + This stops a fly too. The stop is sent even if the motor looks stopped, + as DMOV can flicker at standstill on some motors. - This makes stop idempotent when the motor is already stopped, so a motor - that is both moved and staged in a run (e.g. `bp.scan`) is not stopped - twice -- once by the RunEngine (it stops everything it set) and once by - `unstage`. + The wait allows the acceleration time to decelerate, plus `DEFAULT_TIMEOUT`. """ - if not await self.motor_done_move.get_value(): - await self.motor_stop.set(1) + await self.motor_stop.set(1) + # ACCL is the time to go between rest and VELO, so no stop takes longer + timeout = await self.acceleration_time.get_value() + DEFAULT_TIMEOUT + await wait_for_value(self.motor_done_move, 1, timeout=timeout) async def check_move(self, new_position: float): """Check the positions are within limits. @@ -200,8 +211,8 @@ async def calculate_timeout( async def move(self, new_position: float, timeout: TimeoutCalculator) -> None: """Move by setting the setpoint and waiting for put completion. - Raises `MotorAlarmError` if the readback has an alarm severity set once - the put has completed. + The put completes early if the motor is stopped. Raises `MotorAlarmError` + if the readback has an alarm severity set once the put has completed. """ await self.setpoint.set(new_position, timeout=timeout()) reading = await self.readback.read() @@ -212,8 +223,13 @@ async def move(self, new_position: float, timeout: TimeoutCalculator) -> None: f"after a move to {new_position}." ) - async def on_prepare(self, value: FlyMotorInfo) -> MotorFlyCtx: - """Move to the beginning of a run-up distance ready for a fly scan.""" + async def prepare_to_fly( + self, name: str, value: FlyMotorInfo, previous_ctx: MotorPrepareCtx | None + ) -> AsyncGenerator[MotorPrepareCtx | WatcherUpdate, None]: + """Check the scan is possible, then move to the beginning of the run-up. + + Progress of the move is reported. + """ # Velocity at which the motor travels from start to end, in motor egu/s. max_speed, egu = await asyncio.gather( self.max_velocity.get_value(), self.motor_egu.get_value() @@ -225,43 +241,50 @@ async def on_prepare(self, value: FlyMotorInfo) -> MotorFlyCtx: ) # Check the run-up and run-down positions are within limits acceleration_time = await self.acceleration_time.get_value() - ramp_up_start_pos = value.ramp_up_start_pos(acceleration_time) - ramp_down_end_pos = value.ramp_down_end_pos(acceleration_time) - await asyncio.gather( - self.check_move(ramp_up_start_pos), self.check_move(ramp_down_end_pos) + start = value.ramp_up_start_pos(acceleration_time) + end = value.ramp_down_end_pos(acceleration_time) + await asyncio.gather(self.check_move(start), self.check_move(end)) + yield MotorPrepareCtx( + start=start, + end=end, + acceleration_time=acceleration_time, + timeout=value.timeout, ) # Move to the run-up start at maximum velocity await self.velocity.set(abs(max_speed)) old_position = await self.readback.get_value() - timeout = await self.calculate_timeout(old_position, ramp_up_start_pos) - await self.move(ramp_up_start_pos, lambda: timeout) + timeout = await self.calculate_timeout(old_position, start) + async with AsyncStatus(self.move(start, lambda: timeout)) as status: + async for update in watch_readback( + self, status, target=start, initial=old_position + ): + yield update # Set the velocity we will use for the fly scan await self.velocity.set(value.speed) - return MotorFlyCtx(fly_info=value) - async def on_kickoff(self, ctx: MotorFlyCtx) -> MotorFlyCtx: - """Begin moving the motor from the prepared position to the final position.""" - acceleration_time = await self.acceleration_time.get_value() - target = ctx.fly_info.ramp_down_end_pos(acceleration_time) - if ctx.fly_info.timeout == CALCULATE_TIMEOUT: + async def fly( + self, ctx: MotorPrepareCtx + ) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + """Move from the run-up start to the run-down end, reporting progress. + + Kicked off once the run-up is done and the motor is at constant velocity. + """ + if ctx.timeout == CALCULATE_TIMEOUT: initial = await self.readback.get_value() - timeout = await self.calculate_timeout(initial, target) + timeout = await self.calculate_timeout(initial, ctx.end) else: - timeout = ctx.fly_info.timeout - ctx.status = AsyncStatus(self.move(target, lambda: timeout)) - # Wait out the run-up so the motor is at constant velocity before - # kickoff() returns. A plan can then kickoff the motor and afterwards - # kickoff internally-triggered detectors, which is more accurate for - # long acceleration times. - await asyncio.sleep(acceleration_time) - return ctx - - async def on_complete(self, ctx: MotorFlyCtx) -> None: - """Block until the motor reaches the fly-scan end position.""" - status = error_if_none( - ctx.status, f"kickoff for motor {self.readback.name} not called." - ) - await status + timeout = ctx.timeout + async with AsyncStatus(self.move(ctx.end, lambda: timeout)) as status: + # Wait out the run-up so the motor is at constant velocity before + # kickoff() returns. A plan can then kickoff the motor and afterwards + # kickoff internally-triggered detectors, which is more accurate for + # long acceleration times. + await asyncio.sleep(ctx.acceleration_time) + yield KICKED_OFF + async for update in watch_readback( + self, status, target=ctx.end, initial=ctx.start + ): + yield update class InstantMotorMock(DeviceMock["Motor"]): @@ -283,6 +306,10 @@ def _instant_move(value): set_mock_value(device.motor_done_move, 1) # Done callback_on_mock_put(device.user_setpoint, _instant_move) + # A STOP finishes any move the test has started + callback_on_mock_put( + device.motor_stop, lambda _: set_mock_value(device.motor_done_move, 1) + ) class VeloAndAcclRespectingMotorMock(DeviceMock["Motor"]): @@ -320,7 +347,9 @@ async def _move(value: float) -> None: @default_mock_class(InstantMotorMock) class Motor( - StandardMovable[float], StandardFlyable[FlyMotorInfo, MotorFlyCtx], StandardReadable + StandardMovable[float], + StandardFlyable[FlyMotorInfo, MotorPrepareCtx], + StandardReadable, ): """Device that moves a motor record.""" @@ -364,8 +393,8 @@ def __init__(self, prefix: str, name="") -> None: super().__init__(name) @cached_property - def _logic(self) -> MotorFlyableMovableLogic: - """The combined move + fly logic, shared by movable_logic and flyable_logic.""" + def logic(self) -> MotorFlyableMovableLogic: + """The combined move + fly logic, satisfying both mix-ins at once.""" return MotorFlyableMovableLogic( readback=self.user_readback, setpoint=self.user_setpoint, @@ -380,11 +409,3 @@ def _logic(self) -> MotorFlyableMovableLogic: max_velocity=self.max_velocity, motor_egu=self.motor_egu, ) - - @cached_property - def movable_logic(self) -> MotorFlyableMovableLogic: - return self._logic - - @cached_property - def flyable_logic(self) -> MotorFlyableMovableLogic: - return self._logic diff --git a/src/ophyd_async/epics/pmac/_pmac_trajectory.py b/src/ophyd_async/epics/pmac/_pmac_trajectory.py index 292367f612..ac9a889109 100644 --- a/src/ophyd_async/epics/pmac/_pmac_trajectory.py +++ b/src/ophyd_async/epics/pmac/_pmac_trajectory.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import AsyncGenerator from dataclasses import dataclass import numpy as np @@ -7,10 +8,9 @@ from ophyd_async.core import ( DEFAULT_TIMEOUT, - AsyncStatus, + KICKED_OFF, FlyableLogic, SignalR, - error_if_none, gather_dict, observe_value, set_and_wait_for_value, @@ -39,13 +39,6 @@ SLICE_SIZE = 4000 -@dataclass -class PmacPrepareContext: - path: Path - motor_info: _PmacMotorInfo - ramp_up_time: float - - @dataclass class PmacScanInfo: spec: Spec[Motor] @@ -54,82 +47,99 @@ class PmacScanInfo: @dataclass -class PmacFlyCtx: - """State threaded through a PMAC fly scan (prepare -> kickoff -> complete).""" +class PmacPrepareCtx: + """What a PMAC fly scan needs after `prepare_to_fly` has worked it out.""" - #: Trajectory + motion info built by `on_prepare`. - prepare: PmacPrepareContext - #: The running trajectory, produced by `on_kickoff`. - trajectory_status: AsyncStatus | None = None + #: The path to fly, consumed a slice at a time as the trajectory runs. + path: Path + #: The first slice of the path, which is built into the trajectory to start with. + first_slice: Slice + #: How many points of the path are left after the first slice. + path_length: int + motor_info: _PmacMotorInfo + ramp_up_pos: dict[Motor, np.float64] + ramp_up_time: float + turnaround_time: float | None + #: The PVT the last built slice exited at, which the next slice starts from. + next_pvt: PVT | None = None + #: The number of points loaded into the PMAC so far. + loaded: int = 0 -class PmacTrajectoryFlyableLogic(FlyableLogic[PmacScanInfo, PmacFlyCtx]): +class PmacTrajectoryFlyableLogic(FlyableLogic[PmacScanInfo, PmacPrepareCtx]): def __init__(self, pmac: PmacIO) -> None: self.pmac = pmac - # Working state mutated by the running trajectory task; it spans prepare - # and the executing trajectory, so it lives on the logic rather than in - # the cross-stage context. - self._next_pvt: PVT | None = None - self._loaded = 0 - self._turnaround_time: float | None = None - - async def on_prepare(self, value: PmacScanInfo) -> PmacFlyCtx: - spec = value.spec - self._turnaround_time = value.turnaround_time - self._next_pvt = None - self._loaded = 0 - path = Path(spec.calculate()) - slice = path.consume(SLICE_SIZE) - path_length = len(path) - motors = slice.axes() - motor_info = await _PmacMotorInfo.from_motors(self.pmac, motors) + + async def prepare_to_fly( + self, name: str, value: PmacScanInfo, previous_ctx: PmacPrepareCtx | None + ) -> AsyncGenerator[PmacPrepareCtx, None]: + path = Path(value.spec.calculate()) + first_slice = path.consume(SLICE_SIZE) + motor_info = await _PmacMotorInfo.from_motors(self.pmac, first_slice.axes()) ramp_up_pos, ramp_up_time = calculate_ramp_position_and_duration( - slice, motor_info, True, value.ramp_time + first_slice, motor_info, True, value.ramp_time + ) + ctx = PmacPrepareCtx( + path=path, + first_slice=first_slice, + path_length=len(path), + motor_info=motor_info, + ramp_up_pos=ramp_up_pos, + ramp_up_time=ramp_up_time, + turnaround_time=value.turnaround_time, ) + yield ctx await asyncio.gather( - self._build_trajectory(motor_info, slice, path_length, ramp_up_time), - self._move_to_start(motor_info, ramp_up_pos), + self._build_trajectory(ctx), + self._move_to_start(ctx.motor_info, ctx.ramp_up_pos), ) - return PmacFlyCtx( - prepare=PmacPrepareContext( - path=path, motor_info=motor_info, ramp_up_time=ramp_up_time + + async def fly(self, ctx: PmacPrepareCtx) -> AsyncGenerator[KICKED_OFF, None]: + """Execute the trajectory, appending the rest of the path as it runs.""" + trajectory = self.pmac.trajectory + async with trajectory.execute_profile.set(True, timeout=None) as execute: + # Wait for the ramp up to happen + await wait_for_value( + trajectory.total_points, + lambda v: v >= 1, + ctx.ramp_up_time + DEFAULT_TIMEOUT, ) + yield KICKED_OFF + # We consume SLICE_SIZE from self.path and parse a trajectory + # containing at least 2 * SLICE_SIZE, as a gapless trajectory + # will contain 2 points per slice frame. If gaps are present, + # additional points are inserted, overfilling the buffer. + min_buffer_size = SLICE_SIZE * 2 + async for current_point in observe_value( + trajectory.total_points, + done_status=execute, + # Limit on PMAC is 4 seconds between points + # Thus this timeout is fine. + timeout=DEFAULT_TIMEOUT, + ): + # Keep a minimum buffer size, if we have more points to append + if len(ctx.path) != 0 and ctx.loaded - current_point < min_buffer_size: + # Less than SLICE_SIZE * 2 points in the buffer, so refill + next_slice = ctx.path.consume(SLICE_SIZE) + await self._append_trajectory(ctx, next_slice, len(ctx.path)) + await self._check_profile_status( + trajectory.execute_status, trajectory.execute_message ) - async def on_kickoff(self, ctx: PmacFlyCtx) -> PmacFlyCtx: - prepare = ctx.prepare - ctx.trajectory_status = self._execute_trajectory( - prepare.path, prepare.motor_info - ) - # Wait for the ramp up to happen - await wait_for_value( - self.pmac.trajectory.total_points, - lambda v: v >= 1, - prepare.ramp_up_time + DEFAULT_TIMEOUT, - ) - return ctx + async def ensure_stopped(self) -> None: + """Abort the trajectory if one is running, and wait until it has stopped.""" + trajectory = self.pmac.trajectory + if await trajectory.execute_state.get_value() == PmacExecuteState.EXECUTING: + await trajectory.abort_profile.trigger() + await wait_for_value( + trajectory.execute_state, + lambda v: v != PmacExecuteState.EXECUTING, + timeout=DEFAULT_TIMEOUT, + ) - async def on_complete(self, ctx: PmacFlyCtx) -> None: - trajectory_status = error_if_none( - ctx.trajectory_status, "Cannot complete. Must call kickoff first." - ) - await trajectory_status - # Reset number of loaded points - self._loaded = 0 - - async def stop(self) -> None: - # Abort current trajectory, if one is running - if ( - await self.pmac.trajectory.execute_state.get_value() - == PmacExecuteState.EXECUTING - ): - await self.pmac.trajectory.abort_profile.trigger() - - async def on_stage(self) -> None: - # unstage()/on_unstage default to stop(), but stage additionally resets - await self.stop() - - # Run an empty fly scan to reset EQU on Panda Brick + async def ensure_ready(self) -> None: + """Stop, then run an empty fly scan to reset EQU on Panda Brick.""" + await self.ensure_stopped() for use_axis in self.pmac.trajectory.use_axis.values(): await use_axis.set(False) @@ -141,33 +151,6 @@ async def on_stage(self) -> None: await self.pmac.trajectory.build_profile.trigger() await self.pmac.trajectory.execute_profile.set(True) - @AsyncStatus.wrap - async def _execute_trajectory(self, path: Path, motor_info: _PmacMotorInfo): - execute_status = self.pmac.trajectory.execute_profile.set(True, timeout=None) - # We consume SLICE_SIZE from self.path and parse a trajectory - # containing at least 2 * SLICE_SIZE, as a gapless trajectory - # will contain 2 points per slice frame. If gaps are present, - # additional points are inserted, overfilling the buffer. - min_buffer_size = SLICE_SIZE * 2 - async for current_point in observe_value( - self.pmac.trajectory.total_points, - done_status=execute_status, - # Limit on PMAC is 4 seconds between points - # Thus this timeout is fine. - timeout=DEFAULT_TIMEOUT, - ): - # Ensure we maintain a minimum buffer size, if we have more points to append - if len(path) != 0 and self._loaded - current_point < min_buffer_size: - # We have less than SLICE_SIZE * 2 points in the buffer, so refill - next_slice = path.consume(SLICE_SIZE) - path_length = len(path) - await self._append_trajectory(next_slice, path_length, motor_info) - - await self._check_profile_status( - self.pmac.trajectory.execute_status, - self.pmac.trajectory.execute_message, - ) - async def _check_profile_status( self, status_signal: SignalR, @@ -193,10 +176,10 @@ async def _check_profile_status( ) from exc async def _append_trajectory( - self, slice: Slice, path_length: int, motor_info: _PmacMotorInfo + self, ctx: PmacPrepareCtx, slice: Slice, path_length: int ): - trajectory = await self._parse_trajectory(slice, path_length, motor_info, None) - await self._set_trajectory_arrays(trajectory, motor_info) + trajectory = await self._parse_trajectory(ctx, slice, path_length, None) + await self._set_trajectory_arrays(trajectory, ctx.motor_info) await self.pmac.trajectory.append_profile.trigger() await self._check_profile_status( @@ -204,15 +187,10 @@ async def _append_trajectory( self.pmac.trajectory.append_message, ) - async def _build_trajectory( - self, - motor_info: _PmacMotorInfo, - slice: Slice, - path_length: int, - ramp_up_time: float, - ): + async def _build_trajectory(self, ctx: PmacPrepareCtx): + motor_info = ctx.motor_info trajectory = await self._parse_trajectory( - slice, path_length, motor_info, ramp_up_time + ctx, ctx.first_slice, ctx.path_length, ctx.ramp_up_time ) use_axis = { i: (i in motor_info.motor_cs_index.values()) for i in CS_INDEX.values() @@ -236,17 +214,18 @@ async def _build_trajectory( async def _parse_trajectory( self, + ctx: PmacPrepareCtx, slice: Slice, path_length: int, - motor_info: _PmacMotorInfo, ramp_up_time: float | None = None, ) -> Trajectory: + motor_info = ctx.motor_info trajectory, exit_pvt = Trajectory.from_slice( slice, motor_info, - None if ramp_up_time else self._next_pvt, + None if ramp_up_time else ctx.next_pvt, ramp_up_time=ramp_up_time, - turnaround_time=self._turnaround_time, + turnaround_time=ctx.turnaround_time, ) if path_length == 0: @@ -256,8 +235,8 @@ async def _parse_trajectory( trajectory = trajectory.with_ramp_down( exit_pvt, ramp_down_pos, ramp_down_time, 0 ) - self._next_pvt = exit_pvt - self._loaded += len(trajectory) + ctx.next_pvt = exit_pvt + ctx.loaded += len(trajectory) return trajectory diff --git a/src/ophyd_async/fastcs/panda/_fly_logic.py b/src/ophyd_async/fastcs/panda/_fly_logic.py index a97dc18d41..0206855aae 100644 --- a/src/ophyd_async/fastcs/panda/_fly_logic.py +++ b/src/ophyd_async/fastcs/panda/_fly_logic.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncGenerator from dataclasses import dataclass, field from typing import cast @@ -10,6 +11,7 @@ from scanspec.specs import Spec from ophyd_async.core import ( + KICKED_OFF, ConfinedModel, FlyableLogic, SignalRW, @@ -44,12 +46,15 @@ class ScanSpecInfo(ConfinedModel): @dataclass -class StaticSeqTableFlyableLogic(FlyableLogic[SeqTableInfo, None]): +class StaticSeqTableFlyableLogic(FlyableLogic[SeqTableInfo]): """For controlling the PandA `SeqTable` when fly scanning.""" seq: SeqBlock - async def on_prepare(self, value: SeqTableInfo): + async def prepare_to_fly( + self, name: str, value: SeqTableInfo, previous_ctx: None + ) -> AsyncGenerator[None, None]: + yield None await asyncio.gather( self.seq.prescale_units.set(PandaTimeUnits.US), self.seq.enable.set(PandaBitMux.ZERO), @@ -60,14 +65,13 @@ async def on_prepare(self, value: SeqTableInfo): self.seq.table.set(value.sequence_table), ) - async def on_kickoff(self, ctx: None) -> None: + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF, None]: await self.seq.enable.set(PandaBitMux.ONE) await wait_for_value(self.seq.active, True, timeout=1) - - async def on_complete(self, ctx: None) -> None: + yield KICKED_OFF await wait_for_value(self.seq.active, False, timeout=None) - async def stop(self): + async def ensure_stopped(self) -> None: await self.seq.enable.set(PandaBitMux.ZERO) await wait_for_value(self.seq.active, False, timeout=1) @@ -89,11 +93,14 @@ def from_inenc(cls, panda: CommonPandaBlocks, number: int) -> PosOutScaleOffset: @dataclass -class ScanSpecSeqTableFlyableLogic(FlyableLogic[ScanSpecInfo, None]): +class ScanSpecSeqTableFlyableLogic(FlyableLogic[ScanSpecInfo]): seq: SeqBlock motor_pos_outs: dict[Motor, PosOutScaleOffset] = field(default_factory=dict) - async def on_prepare(self, value: ScanSpecInfo): + async def prepare_to_fly( + self, name: str, value: ScanSpecInfo, previous_ctx: None + ) -> AsyncGenerator[None, None]: + yield None await self.seq.enable.set(PandaBitMux.ZERO) slice = Path(value.spec.calculate()).consume() slice_duration = error_if_none(slice.duration, "Slice must have duration") @@ -155,14 +162,13 @@ async def on_prepare(self, value: ScanSpecInfo): self.seq.table.set(rows), ) - async def on_kickoff(self, ctx: None) -> None: + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF, None]: await self.seq.enable.set(PandaBitMux.ONE) await wait_for_value(self.seq.active, True, timeout=1) - - async def on_complete(self, ctx: None) -> None: + yield KICKED_OFF await wait_for_value(self.seq.active, False, timeout=None) - async def stop(self): + async def ensure_stopped(self) -> None: await self.seq.enable.set(PandaBitMux.ZERO) await wait_for_value(self.seq.active, False, timeout=1) @@ -192,12 +198,15 @@ class PcompInfo(ConfinedModel): @dataclass -class StaticPcompFlyableLogic(FlyableLogic[PcompInfo, None]): +class StaticPcompFlyableLogic(FlyableLogic[PcompInfo]): """For controlling the PandA `PcompBlock` when fly scanning.""" pcomp: PcompBlock - async def on_prepare(self, value: PcompInfo): + async def prepare_to_fly( + self, name: str, value: PcompInfo, previous_ctx: None + ) -> AsyncGenerator[None, None]: + yield None await self.pcomp.enable.set(PandaBitMux.ZERO) await asyncio.gather( self.pcomp.start.set(value.start_position), @@ -207,13 +216,12 @@ async def on_prepare(self, value: PcompInfo): self.pcomp.dir.set(value.direction), ) - async def on_kickoff(self, ctx: None) -> None: + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF, None]: await self.pcomp.enable.set(PandaBitMux.ONE) await wait_for_value(self.pcomp.active, True, timeout=1) - - async def on_complete(self, ctx: None) -> None: + yield KICKED_OFF await wait_for_value(self.pcomp.active, False, timeout=None) - async def stop(self): + async def ensure_stopped(self) -> None: await self.pcomp.enable.set(PandaBitMux.ZERO) await wait_for_value(self.pcomp.active, False, timeout=1) diff --git a/src/ophyd_async/sim/_blob_acquire_logic.py b/src/ophyd_async/sim/_blob_acquire_logic.py index 2fa0faf74a..5a2e5c0c39 100644 --- a/src/ophyd_async/sim/_blob_acquire_logic.py +++ b/src/ophyd_async/sim/_blob_acquire_logic.py @@ -1,7 +1,6 @@ import asyncio -from contextlib import suppress -from ophyd_async.core import DetectorAcquireLogic, TriggerInfo +from ophyd_async.core import DetectorAcquireLogic, TriggerInfo, cancel_and_wait from ._pattern_generator import PatternGenerator @@ -24,7 +23,5 @@ async def wait_for_idle(self): async def ensure_stopped(self): # Stop the background task and wait for it to finish if self.task: - self.task.cancel() - with suppress(asyncio.CancelledError): - await self.task - self.task = None + task, self.task = self.task, None + await cancel_and_wait(task) diff --git a/src/ophyd_async/sim/_motor.py b/src/ophyd_async/sim/_motor.py index 41697fd2e1..71940be103 100644 --- a/src/ophyd_async/sim/_motor.py +++ b/src/ophyd_async/sim/_motor.py @@ -1,55 +1,134 @@ import asyncio -from collections.abc import Callable -from dataclasses import dataclass +from collections.abc import AsyncGenerator, Callable +from dataclasses import dataclass, field from functools import cached_property from ophyd_async.core import ( + KICKED_OFF, AsyncStatus, DeviceMock, + FlyableLogic, FlyMotorInfo, MovableLogic, SignalRW, + StandardFlyable, StandardMovable, StandardReadable, TimeoutCalculator, - WatchableAsyncStatus, + WatcherUpdate, + cancel_and_wait, default_mock_class, - error_if_none, simulate_move, soft_signal_r_and_setter, soft_signal_rw, + watch_readback, ) from ophyd_async.core import StandardReadableFormat as Format +@dataclass(frozen=True) +class SimMotorPrepareCtx: + """What a simulated motor fly scan needs after `prepare_to_fly`.""" + + #: Where to start from: the start position with the run-up distance added. + start: float + #: Where to end: the end position with the run-down distance added. + end: float + #: How long the motor takes to get up to speed. + acceleration_time: float + + @dataclass -class SimMotorMoveLogic(MovableLogic[float]): +class SimMotorLogic( + MovableLogic[float], FlyableLogic[FlyMotorInfo, SimMotorPrepareCtx] +): + """Combined move + fly logic for a simulated motor.""" + readback_set: Callable[[float], None] velocity: SignalRW[float] acceleration_time: SignalRW[float] + _motion: asyncio.Task | None = field(default=None, init=False, repr=False) - async def stop(self) -> None: - """Stop the motion.""" + async def ensure_stopped(self) -> None: + """Stop the simulated motion, a fly included, and wait for it to finish. + + Then set the setpoint to where it stopped. + """ + if (motion := self._motion) is not None: + await cancel_and_wait(motion) await self.setpoint.set(await self.readback.get_value()) async def move(self, new_position: float, timeout: TimeoutCalculator) -> None: + """Simulate the move in the calling task, which `ensure_stopped` cancels.""" old_position, velocity, acceleration_time = await asyncio.gather( self.readback.get_value(), self.velocity.get_value(), self.acceleration_time.get_value(), ) - await self.setpoint.set(new_position) - async for position in simulate_move( - old_position, new_position, velocity, acceleration_time - ): - self.readback_set(position) + motion = self._motion = asyncio.current_task() + try: + moves = simulate_move( + old_position, new_position, velocity, acceleration_time + ) + async for position in moves: + self.readback_set(position) + finally: + # A new move may have replaced us already + if self._motion is motion: + self._motion = None + + async def prepare_to_fly( + self, name: str, value: FlyMotorInfo, previous_ctx: SimMotorPrepareCtx | None + ) -> AsyncGenerator[SimMotorPrepareCtx | WatcherUpdate, None]: + """Work out the run-up and run-down positions, then move to the run-up start. + + Progress of the move is reported. + """ + acceleration_time = await self.acceleration_time.get_value() + ctx = SimMotorPrepareCtx( + start=value.ramp_up_start_pos(acceleration_time), + end=value.ramp_down_end_pos(acceleration_time), + acceleration_time=acceleration_time, + ) + yield ctx + # Move to start as fast as we can + await self.velocity.set(0) + async with AsyncStatus(self.move(ctx.start, lambda: None)) as status: + async for update in watch_readback(self, status, target=ctx.start): + yield update + # Set the velocity for the actual move + await self.velocity.set(value.speed) + + async def fly( + self, ctx: SimMotorPrepareCtx + ) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + """Move from the run-up start to the run-down end, reporting progress. + + Kicked off once the run-up is done and the motor is at constant velocity. + """ + async with AsyncStatus(self.move(ctx.end, lambda: None)) as status: + await asyncio.sleep(ctx.acceleration_time) + yield KICKED_OFF + async for update in watch_readback( + self, status, target=ctx.end, initial=ctx.start + ): + yield update # Remove InstantMovableMock as SimMotor owns this logic, depends if instant=True/False @default_mock_class(DeviceMock) -class SimMotor(StandardReadable, StandardMovable[float]): - """For usage when simulating a motor.""" +class SimMotor( + StandardReadable, + StandardMovable[float], + StandardFlyable[FlyMotorInfo, SimMotorPrepareCtx], +): + """For usage when simulating a motor. + + `stage()` calls the logic's `ensure_ready`, which calls `ensure_stopped` and so + stops the simulated motion and sets the setpoint to the current readback. + `unstage()` cancels a fly in progress, then calls `ensure_stopped` too. + """ def __init__( self, @@ -75,47 +154,14 @@ def __init__( self.acceleration_time = soft_signal_rw(float, 0.5) self.user_setpoint = soft_signal_rw(float, initial_value, units=units) - # Stored in prepare - self._fly_info: FlyMotorInfo | None = None - # Set on kickoff(), complete when motor reaches end position - self._fly_status: WatchableAsyncStatus | None = None - super().__init__(name=name) @cached_property - def movable_logic(self): - return SimMotorMoveLogic( + def logic(self) -> SimMotorLogic: + return SimMotorLogic( readback=self.user_readback, readback_set=self._user_readback_set, setpoint=self.user_setpoint, velocity=self.velocity, acceleration_time=self.acceleration_time, ) - - @AsyncStatus.wrap - async def prepare(self, value: FlyMotorInfo): - """Calculate run-up and move there, setting fly velocity when there.""" - self._fly_info = value - # Move to start as fast as we can - await self.velocity.set(0) - await self.set( - value.ramp_up_start_pos(await self.acceleration_time.get_value()) - ) - # Set the velocity for the actual move - await self.velocity.set(value.speed) - - @AsyncStatus.wrap - async def kickoff(self): - """Begin moving motor from prepared position to final position.""" - fly_info = error_if_none( - self._fly_info, "Motor must be prepared before attempting to kickoff" - ) - acceleration_time = await self.acceleration_time.get_value() - self._fly_status = self.set(fly_info.ramp_down_end_pos(acceleration_time)) - # Wait for the acceleration time to ensure we are at velocity - await asyncio.sleep(acceleration_time) - - def complete(self) -> WatchableAsyncStatus: - """Mark as complete once motor reaches completed position.""" - fly_status = error_if_none(self._fly_status, "kickoff not called") - return fly_status diff --git a/src/ophyd_async/tango/demo/_motor.py b/src/ophyd_async/tango/demo/_motor.py index f1194763f0..c7df2676ed 100644 --- a/src/ophyd_async/tango/demo/_motor.py +++ b/src/ophyd_async/tango/demo/_motor.py @@ -23,8 +23,10 @@ class DemoMotorMoveLogic(MovableLogic[float]): stop_: TriggerableCommand state: SignalR[DevStateEnum] - async def stop(self): + async def ensure_stopped(self) -> None: + """Stop the motor and wait for its state to return to ON.""" await self.stop_.trigger() + await wait_for_value(self.state, DevStateEnum.ON, timeout=DEFAULT_TIMEOUT) async def calculate_timeout( self, old_position: float, new_position: float @@ -52,7 +54,7 @@ class DemoMotor(TangoDevice, StandardReadable, StandardMovable[float]): stop_: TriggerableCommand @cached_property - def movable_logic(self) -> MovableLogic: + def logic(self) -> MovableLogic: return DemoMotorMoveLogic( readback=self.readback, setpoint=self.setpoint, diff --git a/tests/unit_tests/core/test_cancel_swallow.py b/tests/unit_tests/core/test_cancel_swallow.py new file mode 100644 index 0000000000..866617e4da --- /dev/null +++ b/tests/unit_tests/core/test_cancel_swallow.py @@ -0,0 +1,71 @@ +import asyncio +from collections.abc import Callable, Coroutine +from typing import Any +from unittest.mock import Mock + +import pytest + +from ophyd_async.core import AsyncStatus, soft_signal_rw + +# The waits these cleanups cancel are internal, with no public way to slow them down +from ophyd_async.core import _signal as signal_module # noqa: PLC2701 +from ophyd_async.sim._blob_acquire_logic import BlobAcquireLogic # noqa: PLC2701 + + +def _slow_to_cancel(cancelling: asyncio.Event): + """A coroutine that takes a moment to finish cancelling, like async cleanup.""" + + async def coro(): + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + cancelling.set() + await asyncio.sleep(0.2) + raise + + return coro() + + +async def _status_aexit(cancelling: asyncio.Event, _: pytest.MonkeyPatch): + async with AsyncStatus(_slow_to_cancel(cancelling)): + await asyncio.sleep(0) # let the status task start + + +async def _blob_ensure_stopped(cancelling: asyncio.Event, _: pytest.MonkeyPatch): + pattern_generator = Mock() + pattern_generator.write_images_to_file = lambda: _slow_to_cancel(cancelling) + logic = BlobAcquireLogic(pattern_generator) + await logic.start_acquiring() + await asyncio.sleep(0) # let the task start + await logic.ensure_stopped() + + +async def _set_and_wait(cancelling: asyncio.Event, monkeypatch: pytest.MonkeyPatch): + async def wait_for_value(self, signal, timeout): + await _slow_to_cancel(cancelling) + + monkeypatch.setattr(signal_module._ValueChecker, "wait_for_value", wait_for_value) + sig = soft_signal_rw(int) + await sig.connect() + # No initial value ever arrives, so this times out and the error handler + # cancels the slow to cancel wait task + with pytest.raises(TimeoutError): + await signal_module.set_and_wait_for_value(sig, 1, timeout=0.1) + + +@pytest.mark.parametrize( + "cleanup", [_status_aexit, _blob_ensure_stopped, _set_and_wait] +) +async def test_cancelling_caller_during_cleanup_waits_for_it_then_raises( + cleanup: Callable[[asyncio.Event, pytest.MonkeyPatch], Coroutine[Any, Any, None]], + monkeypatch: pytest.MonkeyPatch, +): + cancelling = asyncio.Event() + outer = asyncio.create_task(cleanup(cancelling, monkeypatch)) + # Cancel the caller while the cleanup is waiting for the inner task to finish + await cancelling.wait() + outer.cancel() + await asyncio.wait([outer]) + assert outer.cancelled() + # The caller ended only once the slow cleanup had finished + assert asyncio.all_tasks() == {asyncio.current_task()} diff --git a/tests/unit_tests/core/test_flyer.py b/tests/unit_tests/core/test_flyer.py index 70faeb4ce5..d5dd3e6ad7 100644 --- a/tests/unit_tests/core/test_flyer.py +++ b/tests/unit_tests/core/test_flyer.py @@ -1,116 +1,193 @@ import asyncio +import sys +from collections.abc import AsyncGenerator, AsyncIterator from dataclasses import dataclass, field from functools import cached_property +from typing import Any, cast from unittest.mock import Mock import pytest from ophyd_async.core import ( + KICKED_OFF, + AsyncStatus, + Device, FlyableLogic, MovableLogic, StandardFlyable, StandardMovable, StandardReadable, StandardReadableFormat, + TransitionError, WatchableAsyncStatus, + WatcherUpdate, init_devices, set_mock_value, soft_signal_rw, + watch_readback, ) +from ophyd_async.testing import wait_for_pending_wakeups @dataclass -class RecordingFlyableLogic(FlyableLogic[int, None]): - """A `FlyableLogic` that records its calls. Not movable, so complete() reports - no progress to watchers.""" +class RecordingFlyableLogic(FlyableLogic[int]): + """A `FlyableLogic` that records its calls. Not movable, so prepare() and + complete() report no progress to watchers.""" calls: list[str] = field(default_factory=list) prepared_with: list[int] = field(default_factory=list) - async def on_prepare(self, value: int) -> None: - self.calls.append("on_prepare") + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[None, None]: + self.calls.append("prepare_to_fly") self.prepared_with.append(value) + yield None - async def on_kickoff(self, ctx: None) -> None: - self.calls.append("on_kickoff") + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + self.calls.append("kicking off") + yield KICKED_OFF + self.calls.append("completing") - async def on_complete(self, ctx: None) -> None: - self.calls.append("on_complete") + async def ensure_stopped(self) -> None: + self.calls.append("ensure_stopped") - async def stop(self) -> None: - self.calls.append("stop") + +@dataclass(frozen=True) +class Ctx: + value: int + # Anything that changes during the scan is an object the ctx holds + log: list[str] = field(default_factory=list) @dataclass -class CtxFlyableLogic(FlyableLogic[int, dict]): - """A `FlyableLogic` that threads a context dict between its stages.""" +class CtxFlyableLogic(FlyableLogic[int, Ctx]): + """A `FlyableLogic` that passes a ctx from prepare_to_fly to fly.""" - seen: list[tuple[str, dict]] = field(default_factory=list) + seen: list[Ctx] = field(default_factory=list) + made: list[Ctx] = field(default_factory=list) - async def on_prepare(self, value: int) -> dict: - return {"prepared": value} + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[Ctx, None]: + self.made.append(Ctx(value)) + yield self.made[0] - async def on_kickoff(self, ctx: dict) -> dict: - self.seen.append(("kickoff", dict(ctx))) - ctx["kicked"] = True - return ctx + async def fly(self, ctx: Ctx) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + self.seen.append(ctx) + yield KICKED_OFF - async def on_complete(self, ctx: dict) -> None: - self.seen.append(("complete", dict(ctx))) + async def ensure_stopped(self) -> None: ... @dataclass -class MovableFlyableLogic(MovableLogic[float], FlyableLogic[float, None]): +class MovableFlyableLogic(MovableLogic[float], FlyableLogic[float, float]): """Both movable and flyable, mirroring how `Motor` combines the two. - `prepare()` and `complete()` should report progress by observing the readback, - reusing the same watcher-update stream as `StandardMovable.set`. + `prepare()` and `complete()` should report progress by watching the readback + with `watch_readback`, reusing the same watcher updates as + `StandardMovable.set`. """ - async def on_prepare(self, value: float) -> None: - # Move to the fly start, writing the setpoint part-way in like Motor does - await asyncio.sleep(0) - set_mock_value(self.setpoint, value) - for position in (value / 2, value): - set_mock_value(self.readback, position) - await asyncio.sleep(0) - - async def on_kickoff(self, ctx: None) -> None: - pass - - async def on_complete(self, ctx: None) -> None: + async def _step_to(self, target: float) -> None: # Step the readback towards the target so watchers see progress. - target = await self.setpoint.get_value() for position in (target / 2, target): set_mock_value(self.readback, position) await asyncio.sleep(0) + async def prepare_to_fly( + self, name: str, value: float, previous_ctx: None + ) -> AsyncGenerator[float | WatcherUpdate, None]: + yield value + async with AsyncStatus(self._step_to(value)) as status: + async for update in watch_readback(self, status, target=value): + yield update + + async def fly(self, ctx: float) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + async with AsyncStatus(self._step_to(2 * ctx)) as status: + yield KICKED_OFF + async for update in watch_readback(self, status, 2 * ctx, initial=ctx): + yield update + + async def ensure_stopped(self) -> None: ... + -class MovableFlyer(StandardMovable[float], StandardFlyable[float, None]): +class MovableFlyer(StandardMovable[float], StandardFlyable[float, float]): def __init__(self, name: str = ""): self.readback = soft_signal_rw(float) self.setpoint = soft_signal_rw(float) super().__init__(name=name) @cached_property - def _logic(self) -> MovableFlyableLogic: + def logic(self) -> MovableFlyableLogic: return MovableFlyableLogic(setpoint=self.setpoint, readback=self.readback) - @cached_property - def movable_logic(self) -> MovableFlyableLogic: - return self._logic - @cached_property - def flyable_logic(self) -> MovableFlyableLogic: - return self._logic +@dataclass +class BlockingFlyableLogic(RecordingFlyableLogic): + """Blocks forever in the named verb's part of a hook, so it is left in flight.""" + + block: str = "" + + async def _maybe_block(self, verb: str) -> None: + self.calls.append(verb) + if self.block == verb: + await asyncio.Event().wait() + + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[None, None]: + await self._maybe_block("prepare") + yield None + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + await self._maybe_block("kickoff") + yield KICKED_OFF + await self._maybe_block("complete") + + +@dataclass +class FlyingLogic(FlyableLogic[int]): + """Kicks off a status that runs until cancelled; `fail` raises after kickoff + instead of waiting for it.""" + + fail: bool = False + fly_statuses: list[AsyncStatus] = field(default_factory=list) + running_when_stopped: list[bool] = field(default_factory=list) + + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[None, None]: + yield None + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + async with AsyncStatus(asyncio.Event().wait()) as status: + self.fly_statuses.append(status) + yield KICKED_OFF + await asyncio.sleep(0) + if self.fail: + raise ValueError("bad fly") + await status + + async def ensure_stopped(self) -> None: + self.running_when_stopped.append( + any(not status.done for status in self.fly_statuses) + ) + + +async def make_staged_flyer(logic: FlyableLogic[Any, Any]): + async with init_devices(mock=True): + flyer = logic.with_device(name="flyer") + await flyer.stage() + return flyer @pytest.fixture async def recording_flyer(): - """An ephemeral `StandardFlyable` wrapping a fresh `RecordingFlyableLogic`.""" + """A staged ephemeral `StandardFlyable` wrapping a `RecordingFlyableLogic`.""" logic = RecordingFlyableLogic() - async with init_devices(mock=True): - flyer = logic.with_device(name="flyer") + flyer = await make_staged_flyer(logic) + logic.calls.clear() # forget the ensure_stopped() from stage return logic, flyer @@ -120,148 +197,318 @@ async def test_ephemeral_flyable_drives_logic_hooks(recording_flyer): assert flyer.name == "flyer" assert isinstance(flyer, StandardFlyable) # The ephemeral device exposes the same logic instance - assert flyer.flyable_logic is logic + assert flyer.logic is logic - await flyer.prepare(5) - await flyer.kickoff() - await flyer.complete() + # Every point of a scan reuses the staged flyer: FLOWN can prepare again + for value in (5, 6): + await flyer.prepare(value) + await flyer.kickoff() + await flyer.complete() - assert logic.calls == ["on_prepare", "on_kickoff", "on_complete"] - assert logic.prepared_with == [5] + assert logic.calls == ["prepare_to_fly", "kicking off", "completing"] * 2 + assert logic.prepared_with == [5, 6] + await flyer.unstage() + assert logic.calls[-1] == "ensure_stopped" -async def test_flyable_threads_context_between_stages(): +async def test_flyable_passes_ctx_to_fly(): logic = CtxFlyableLogic() - async with init_devices(mock=True): - flyer = logic.with_device(name="flyer") + flyer = await make_staged_flyer(logic) await flyer.prepare(5) await flyer.kickoff() await flyer.complete() - # kickoff sees exactly what prepare produced; complete sees what kickoff added - assert logic.seen == [ - ("kickoff", {"prepared": 5}), - ("complete", {"prepared": 5, "kicked": True}), - ] - - -async def test_flyable_kickoff_before_prepare_raises(recording_flyer): - logic, flyer = recording_flyer + # fly sees the one ctx prepare_to_fly yielded + assert logic.seen == [logic.made[0]] + assert logic.seen[0] is logic.made[0] + + +@pytest.mark.parametrize( + "setup, verb, current", + [ + # Not staged + ([], "prepare", "UNSTAGED"), + # Out of order + (["stage"], "kickoff", "STAGED"), + (["stage"], "complete", "STAGED"), + (["stage", "prepare"], "complete", "PREPARED"), + # Prepared: kickoff before preparing again + (["stage", "prepare"], "prepare", "PREPARED"), + (["stage", "prepare", "kickoff"], "kickoff", "FLYING"), + # A running fly cannot be prepared over + (["stage", "prepare", "kickoff"], "prepare", "FLYING"), + # Flown: only a new prepare (or complete, stage, unstage) is allowed + (["stage", "prepare", "kickoff", "complete"], "kickoff", "FLOWN"), + # A new prepare forgets the last fly + ( + ["stage", "prepare", "kickoff", "complete", "prepare"], + "complete", + "PREPARED", + ), + ], +) +async def test_flyable_verb_in_the_wrong_state_raises(setup, verb, current): + # A fly that is kicked off keeps flying until completed + logic = BlockingFlyableLogic(block="complete" if setup[-1:] == ["kickoff"] else "") + async with init_devices(mock=True): + flyer = logic.with_device(name="flyer") + for step in setup: + await (getattr(flyer, step)(1) if step == "prepare" else getattr(flyer, step)()) + logic.calls.clear() - with pytest.raises(RuntimeError, match="prepare.* before kickoff"): - await flyer.kickoff() + with pytest.raises(TransitionError, match=rf"flyer.*from {current}.*allowed"): + await (getattr(flyer, verb)(1) if verb == "prepare" else getattr(flyer, verb)()) assert logic.calls == [] + await flyer.unstage() -async def test_flyable_complete_before_kickoff_raises(recording_flyer): - logic, flyer = recording_flyer +@pytest.mark.parametrize("verb", ["prepare", "kickoff", "complete"]) +async def test_stage_during_a_verb_raises_and_unstage_cancels_it(verb): + logic = BlockingFlyableLogic(block=verb) + flyer = await make_staged_flyer(logic) + verbs = { + "prepare": lambda: flyer.prepare(1), + "kickoff": flyer.kickoff, + "complete": flyer.complete, + } + running = None + for name, start in verbs.items(): + running = start() + if name == verb: + break + await running + await wait_for_pending_wakeups() + assert logic.calls[-1] == verb + + # Only unstage can end a verb that is in flight + with pytest.raises(TransitionError, match="cannot go to STAGED"): + await flyer.stage() + await flyer.unstage() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(cast(AsyncStatus, running), timeout=1) - with pytest.raises(RuntimeError, match="kickoff.* before complete"): - await flyer.complete() - assert logic.calls == [] + # The cancelled verb did not touch the state unstage left behind + logic.block = "" + await flyer.stage() + await flyer.prepare(2) -async def test_flyable_stage_resets_lifecycle(recording_flyer): - # After prepare+kickoff, staging should reset so a later complete errors again. +async def test_flyable_unstage_resets_lifecycle(recording_flyer): + # After prepare+kickoff, unstaging should reset so a later complete errors. logic, flyer = recording_flyer await flyer.prepare(5) await flyer.kickoff() + await flyer.unstage() await flyer.stage() - with pytest.raises(RuntimeError, match="kickoff.* before complete"): + with pytest.raises(TransitionError, match="from STAGED"): await flyer.complete() -async def test_flyable_complete_without_movable_logic_yields_no_updates( - recording_flyer, -): - # A flyer whose logic is not a MovableLogic has no readback to report, so - # complete() is still a WatchableAsyncStatus but calls no watchers. +async def test_flyable_stage_twice_raises(recording_flyer): logic, flyer = recording_flyer - await flyer.prepare(5) + with pytest.raises(TransitionError, match="cannot go to STAGED from STAGED"): + await flyer.stage() + # Nothing ran, and it is still staged + assert logic.calls == [] + await flyer.prepare(1) + + +async def test_flyable_without_progress_yields_no_updates(recording_flyer): + # A logic whose hooks yield no updates reports no progress, but prepare() + # and complete() are still watchable + logic, flyer = recording_flyer + + prepare_watcher, complete_watcher = Mock(), Mock() + status = flyer.prepare(5) + assert isinstance(status, WatchableAsyncStatus) + status.watch(prepare_watcher) + await status await flyer.kickoff() status = flyer.complete() assert isinstance(status, WatchableAsyncStatus) - updates: list[dict] = [] - status.watch(lambda **kwargs: updates.append(kwargs)) + status.watch(complete_watcher) await status - assert logic.calls == ["on_prepare", "on_kickoff", "on_complete"] - assert updates == [] + prepare_watcher.assert_not_called() + complete_watcher.assert_not_called() + assert logic.calls == ["prepare_to_fly", "kicking off", "completing"] -async def test_flyable_complete_watches_movable_logic(): +async def test_flyable_reports_progress_from_logic(): async with init_devices(mock=True): flyer = MovableFlyer(name="flyer") + await flyer.stage() - # Prepare at 0, then set the fly target on the setpoint without triggering the - # instant-move mock (so the readback stays at 0 until on_complete steps it). - await flyer.prepare(0.0) - set_mock_value(flyer.setpoint, 10.0) + prepare_watcher, complete_watcher = Mock(), Mock() + status = flyer.prepare(10.0) + assert isinstance(status, WatchableAsyncStatus) + status.watch(prepare_watcher) + await status await flyer.kickoff() status = flyer.complete() assert isinstance(status, WatchableAsyncStatus) - updates: list[dict] = [] - status.watch(lambda **kwargs: updates.append(kwargs)) + status.watch(complete_watcher) await status - # Progress was reported by observing the readback up to the target, with the - # target read from the setpoint (as StandardMovable.set reports a move). - assert updates, "expected at least one watcher update" - assert updates[-1]["current"] == 10.0 + # Progress is stamped with the flyer's name, and prepare() moves to its + # start before kickoff is allowed + updates = [c.kwargs for c in prepare_watcher.call_args_list] + assert [u["current"] for u in updates] == [0.0, 5.0, 10.0] assert {u["initial"] for u in updates} == {0.0} assert {u["target"] for u in updates} == {10.0} assert {u["name"] for u in updates} == {"flyer"} + updates = [c.kwargs for c in complete_watcher.call_args_list] + assert updates[-1]["current"] == 20.0 + assert {u["initial"] for u in updates} == {10.0} + assert {u["target"] for u in updates} == {20.0} + assert {u["name"] for u in updates} == {"flyer"} -async def test_flyable_prepare_watches_movable_logic(): +async def test_watch_readback_reraises_the_status_error(): async with init_devices(mock=True): flyer = MovableFlyer(name="flyer") - status = flyer.prepare(10.0) - assert isinstance(status, WatchableAsyncStatus) - watcher = Mock() - status.watch(watcher) - await status + async def fail(): + await asyncio.sleep(0) + raise ValueError("boom") - updates = [c.kwargs for c in watcher.call_args_list] - assert [u["current"] for u in updates] == [0.0, 5.0, 10.0] - assert {u["initial"] for u in updates} == {0.0} - # No updates before on_prepare writes the setpoint, so the target is never stale - assert {u["target"] for u in updates} == {10.0} - assert {u["name"] for u in updates} == {"flyer"} - # prepare() still completes the lifecycle stage, so kickoff is allowed - await flyer.kickoff() + updates = [] + with pytest.raises(ValueError, match="boom"): + async with AsyncStatus(fail()) as status: + async for update in watch_readback(flyer.logic, status, target=1.0): + updates.append(update) + # Initial defaults to the readback now + assert {u.initial for u in updates} == {0.0} -async def test_flyable_stage_unstage_default_to_stop(recording_flyer): +@pytest.mark.parametrize("prepare", [False, True]) +async def test_unstage_always_stops(recording_flyer, prepare): logic, flyer = recording_flyer + if prepare: + await flyer.prepare(1) + await flyer.unstage() + assert "ensure_stopped" in logic.calls + + # ensure_ready defaults to ensure_stopped + logic.calls.clear() await flyer.stage() + assert logic.calls == ["ensure_stopped"] await flyer.unstage() - # on_stage/on_unstage both default to stop() - assert logic.calls == ["stop", "stop"] + assert logic.calls == ["ensure_stopped", "ensure_stopped"] + + +@dataclass +class FailingFlyableLogic(RecordingFlyableLogic): + """Raises from the hook named by `failing`.""" + + failing: str = "" + + def _maybe_fail(self, hook: str) -> None: + if self.failing == hook: + raise ValueError(f"bad {hook}") + + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[None, None]: + self._maybe_fail("prepare_to_fly") + yield None + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + self._maybe_fail("kickoff") + yield KICKED_OFF + self._maybe_fail("complete") + + async def ensure_ready(self) -> None: + self._maybe_fail("ensure_ready") + + async def ensure_stopped(self) -> None: + self._maybe_fail("ensure_stopped") + + +@pytest.mark.parametrize( + "failing, verbs", + [ + ("ensure_ready", ["stage"]), + ("prepare_to_fly", ["stage", "prepare"]), + ("kickoff", ["stage", "prepare", "kickoff"]), + ("complete", ["stage", "prepare", "kickoff", "complete"]), + ("ensure_stopped", ["stage", "prepare", "unstage"]), + ], +) +async def test_an_error_from_a_hook_is_raised(failing, verbs): + logic = FailingFlyableLogic(failing=failing) + async with init_devices(mock=True): + flyer = logic.with_device(name="flyer") + + def call(verb: str): + return flyer.prepare(1) if verb == "prepare" else getattr(flyer, verb)() + + *before, last = verbs + for verb in before: + await call(verb) + with pytest.raises(ValueError, match=f"bad {failing}"): + await call(last) async def test_flyable_stage_unstage_can_differ(): # A flyer (e.g. PmacTrajectoryFlyableLogic) may need distinct stage/unstage. @dataclass class StageUnstageLogic(RecordingFlyableLogic): - async def on_stage(self) -> None: - self.calls.append("on_stage") - - async def on_unstage(self) -> None: - self.calls.append("on_unstage") + async def ensure_ready(self) -> None: + self.calls.append("ensure_ready") logic = StageUnstageLogic() async with init_devices(mock=True): flyer = logic.with_device(name="flyer") await flyer.stage() + assert logic.calls == ["ensure_ready"] + await flyer.prepare(1) + # prepare() does not get it ready again + assert logic.calls == ["ensure_ready", "prepare_to_fly"] await flyer.unstage() - assert logic.calls == ["on_stage", "on_unstage"] + assert logic.calls == ["ensure_ready", "prepare_to_fly", "ensure_stopped"] + + +@pytest.mark.parametrize( + "stale, replacement", + [ + ("on_prepare", "prepare_to_fly"), + ("on_kickoff", "fly"), + ("on_complete", "fly"), + ("on_stage", "ensure_ready"), + ("on_unstage", "ensure_stopped"), + ("stop", "ensure_stopped"), + ], +) +def test_flyable_logic_defining_a_stale_hook_raises(stale: str, replacement: str): + with pytest.raises( + TypeError, + match=rf"StaleLogic defines {stale}, which .*{stale} -> {replacement}", + ): + type("StaleLogic", (RecordingFlyableLogic,), {stale: lambda self: None}) + + +@pytest.mark.parametrize( + "logic_class, kwargs, missing", + [ + (MovableLogic, {"setpoint": None, "readback": None}, "ensure_stopped"), + (FlyableLogic, {}, "ensure_stopped, fly, prepare_to_fly"), + ], +) +def test_logic_missing_an_abstract_hook_raises_at_construction( + logic_class: type, kwargs: dict, missing: str +): + incomplete = type("Incomplete", (logic_class,), {}) + with pytest.raises(TypeError, match=r"Incomplete.*abstract") as exc_info: + incomplete(**kwargs) + # The order of the missing methods in the message varies between Pythons + for name in missing.split(", "): + assert name in str(exc_info.value) async def test_flyable_composes_with_readable_staging(): @@ -276,7 +523,7 @@ def __init__(self, name: str = ""): super().__init__(name=name) @cached_property - def flyable_logic(self) -> FlyableLogic[int, None]: + def logic(self) -> FlyableLogic[int, None]: return logic async with init_devices(mock=True): @@ -287,10 +534,647 @@ def flyable_logic(self) -> FlyableLogic[int, None]: # "not being monitored" if StandardFlyable.stage had shadowed it). await device.sig.read(cached=True) # StandardFlyable contribution ran too. - assert logic.calls == ["stop"] + assert logic.calls == ["ensure_stopped"] await device.unstage() - assert logic.calls == ["stop", "stop"] + # Unstage stops even though it was not prepared + assert logic.calls == ["ensure_stopped", "ensure_stopped"] # The cache is torn down on unstage. with pytest.raises(RuntimeError, match="not being monitored"): await device.sig.read(cached=True) + + +# Bug: fly statuses were never cancelled +async def test_unstage_cancels_the_running_fly_before_stopping_the_hardware(): + logic = FlyingLogic() + flyer = await make_staged_flyer(logic) + await flyer.prepare(1) + await flyer.kickoff() + await asyncio.sleep(0) + assert not logic.fly_statuses[0].done + + await flyer.unstage() + + assert logic.fly_statuses[0].task.cancelled() + # Neither stage's nor unstage's ensure_stopped saw the fly running + assert logic.running_when_stopped == [False, False] + + +# Bug: an unstage() cancelled before it awaited the fly left it running unowned +async def test_cancelled_unstage_still_cancels_the_fly(): + logic = FlyingLogic() + flyer = await make_staged_flyer(logic) + await flyer.prepare(1) + await flyer.kickoff() + + unstaging = flyer.unstage() + await asyncio.sleep(0) + unstaging.task.cancel() + with pytest.raises(asyncio.CancelledError): + await unstaging + await wait_for_pending_wakeups() + + assert logic.fly_statuses[0].task.cancelled() + + +# Bug: an unstage() that failed or was cancelled while stopping left its ctx +# behind as the next prepare's previous_ctx +@pytest.mark.parametrize("stop", ["fail", "hang"]) +async def test_restage_after_an_interrupted_unstage_starts_afresh(stop: str): + @dataclass + class StopLogic(FlyableLogic[int, int]): + stopping: str = "" + previous_ctxs: list[int | None] = field(default_factory=list) + + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: int | None + ) -> AsyncGenerator[int, None]: + self.previous_ctxs.append(previous_ctx) + yield value + + async def fly(self, ctx: int) -> AsyncGenerator[KICKED_OFF, None]: + yield KICKED_OFF + + async def ensure_stopped(self) -> None: + if self.stopping == "fail": + raise ValueError("stop failed") + elif self.stopping == "hang": + await asyncio.Event().wait() + + logic = StopLogic() + flyer = await make_staged_flyer(logic) + await flyer.prepare(1) + logic.stopping = stop + unstaging = flyer.unstage() + await asyncio.sleep(0) + # RE.stop/abort/halt cancel a hanging unstage + unstaging.task.cancel() + with pytest.raises((ValueError, asyncio.CancelledError)): + await unstaging + + logic.stopping = "" + await flyer.stage() + await flyer.prepare(2) + await flyer.kickoff() + await flyer.complete() + await flyer.unstage() + assert logic.previous_ctxs == [None, None] + + +# Bug: a failed re-prepare left the flyer prepared with the failed prepare's ctx +async def test_failed_reprepare_leaves_nothing_to_kick_off(recording_flyer): + logic, flyer = recording_flyer + await flyer.prepare(1) + await flyer.kickoff() + await flyer.complete() + logic.calls.clear() + + async def fail(name: str, value: int, previous_ctx: None) -> AsyncIterator[None]: + raise ValueError("bad prepare") + yield None + + logic.prepare_to_fly = fail # type: ignore + with pytest.raises(ValueError, match="bad prepare"): + await flyer.prepare(2) + + with pytest.raises(TransitionError, match="cannot go to FLYING"): + await flyer.kickoff() + assert logic.calls == [] + + +@pytest.mark.parametrize("failing", ["first", "later"]) +async def test_a_retried_prepare_is_a_first_prepare(recording_flyer, failing: str): + logic, flyer = recording_flyer + seen: list[Ctx | None] = [] + fail = True + + async def prepare_to_fly(name: str, value: int, previous_ctx: Any): + seen.append(previous_ctx) + yield Ctx(value) + if fail: + raise ValueError("bad start") + + logic.prepare_to_fly = prepare_to_fly # type: ignore + if failing == "later": + fail = False + await flyer.prepare(1) + await flyer.kickoff() + await flyer.complete() + fail = True + seen.clear() + with pytest.raises(ValueError, match="bad start"): + await flyer.prepare(1) + fail = False + await flyer.prepare(1) + + # The retry is given no context, not that of a prepare before the failed one + assert seen[-1] is None + + +async def test_complete_returns_the_fly_status_after_unstage(recording_flyer): + _, flyer = recording_flyer + await flyer.prepare(1) + await flyer.kickoff() + status = flyer.complete() + await status + + await flyer.unstage() + + assert flyer.complete() is status + + +async def _raise_before_kicked_off(ctx: None) -> AsyncIterator[Any]: + raise ValueError("bad kickoff") + yield KICKED_OFF + + +async def _yield_an_update_first(ctx: None) -> AsyncIterator[Any]: + yield WatcherUpdate(current=0, initial=0, target=1) + + +async def _finish_without_kicking_off(ctx: None) -> AsyncIterator[Any]: + return + yield KICKED_OFF + + +async def _block_until_cancelled(ctx: None) -> AsyncIterator[Any]: + await asyncio.Event().wait() + yield KICKED_OFF + + +# A kickoff that does not reach KICKED_OFF leaves the flyer staged, so it has to +# be prepared again, and leaves nothing running +@pytest.mark.parametrize( + "fly, error, match", + [ + (_raise_before_kicked_off, ValueError, "bad kickoff"), + (_yield_an_update_first, TypeError, "must yield KICKED_OFF first"), + (_finish_without_kicking_off, RuntimeError, "finished without yielding"), + (_block_until_cancelled, asyncio.CancelledError, None), + ], +) +async def test_kickoff_that_does_not_kick_off_goes_back_to_staged( + recording_flyer, fly, error, match +): + logic, flyer = recording_flyer + await flyer.prepare(1) + + logic.fly = fly + kickoff = flyer.kickoff() + if error is asyncio.CancelledError: + # Cancelled directly, not by unstage + await asyncio.sleep(0) + kickoff.task.cancel() + with pytest.raises(error, match=match): + await kickoff + + with pytest.raises(TransitionError, match="from STAGED"): + await flyer.kickoff() + await flyer.prepare(2) + + +async def test_kickoff_cancelled_just_after_kicked_off_ends_the_fly(recording_flyer): + logic, flyer = recording_flyer + await flyer.prepare(1) + kickoff = flyer.kickoff() + ended = asyncio.Event() + + async def fly(ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + try: + yield KICKED_OFF + # An outside cancel (e.g. a timeout) that lands before kickoff resumes + kickoff.task.cancel() + await asyncio.Event().wait() + finally: + ended.set() + + logic.fly = fly + with pytest.raises(asyncio.CancelledError): + await kickoff + assert ended.is_set() + with pytest.raises(TransitionError, match="from STAGED"): + await flyer.kickoff() + await flyer.prepare(2) + + +@pytest.mark.parametrize("bad_yield", [42, KICKED_OFF]) +async def test_a_non_update_yielded_after_kicked_off_fails_complete( + recording_flyer, bad_yield +): + logic, flyer = recording_flyer + + async def fly(ctx: None) -> AsyncIterator[Any]: + yield KICKED_OFF + yield bad_yield + + logic.fly = fly + await flyer.prepare(1) + await flyer.kickoff() + with pytest.raises(TypeError, match=r"\.fly\(\) must yield only WatcherUpdates"): + await flyer.complete() + + +# A complete() that does not succeed ends the fly and leaves the flyer staged, +# whether it is cancelled directly (not via unstage) or fly raises after kickoff +@pytest.mark.parametrize("fail", [False, True]) +async def test_unsuccessful_complete_cancels_the_fly_and_goes_to_staged(fail): + logic = FlyingLogic(fail=fail) + flyer = await make_staged_flyer(logic) + await flyer.prepare(1) + await flyer.kickoff() + completing = flyer.complete() + + if fail: + with pytest.raises(ValueError, match="bad fly"): + await completing + else: + await asyncio.sleep(0) + assert not logic.fly_statuses[0].done + completing.task.cancel() + with pytest.raises(asyncio.CancelledError): + await completing + + assert logic.fly_statuses[0].task.cancelled() + # complete() still returns the failed fly, but it cannot be kicked off again + assert flyer.complete() is completing + with pytest.raises(TransitionError, match="cannot go to FLYING from STAGED"): + await flyer.kickoff() + # Prepared again, it can fly again + logic.fail = False + await flyer.prepare(2) + await flyer.kickoff() + await flyer.unstage() + + +def test_logic_that_is_also_movable_may_define_stop_with_a_deprecation(): + with pytest.deprecated_call(match=r"CombinedLogic.stop\(\) is deprecated"): + + class CombinedLogic(MovableLogic[float], FlyableLogic[None]): + async def stop(self) -> None: ... + + assert CombinedLogic.ensure_stopped is CombinedLogic.stop + + +@dataclass +class FailingCleanupLogic(RecordingFlyableLogic): + """Its fly raises while it is being cancelled, `block_kickoff` or not.""" + + block_kickoff: bool = False + #: Fail on its own just after kicking off, rather than when cancelled + raise_after_kickoff: bool = False + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + try: + if self.block_kickoff: + await asyncio.Event().wait() + yield KICKED_OFF + if not self.raise_after_kickoff: + await asyncio.Event().wait() + finally: + raise ValueError("cleanup failed") # noqa: B012 + + +# A cancelled kickoff() reports its cancel; complete() is the fly, so reports +# the failure of the fly's cleanup +@pytest.mark.parametrize( + "phase, error", + [ + ("kickoff", asyncio.CancelledError), + ("complete", ValueError), + ], +) +async def test_failing_fly_cleanup_when_kickoff_or_complete_is_cancelled( + phase: str, error: type[BaseException] +): + logic = FailingCleanupLogic(block_kickoff=phase == "kickoff") + flyer = await make_staged_flyer(logic) + await flyer.prepare(1) + status = flyer.kickoff() + if phase == "complete": + await status + status = flyer.complete() + await asyncio.sleep(0) + status.task.cancel() + + with pytest.raises(error): + await status + # Staged again, ready to be prepared + await flyer.prepare(2) + + +async def test_unstage_stops_the_hardware_then_raises_the_fly_failure(): + # The fly fails while being cancelled + logic = FailingCleanupLogic() + flyer = await make_staged_flyer(logic) + await flyer.prepare(1) + await flyer.kickoff() + await asyncio.sleep(0) + logic.calls.clear() + + with pytest.raises(ValueError, match="cleanup failed"): + await flyer.unstage() + + assert logic.calls == ["ensure_stopped"] + # Unstaged all the same + await flyer.stage() + + +class FailsToUnstage(Device): + @AsyncStatus.wrap + async def stage(self) -> None: ... + + @AsyncStatus.wrap + async def unstage(self) -> None: + raise ValueError("unstage failed") + + +@dataclass +class HangingStopLogic(RecordingFlyableLogic): + """Once `hang` is set, its `ensure_stopped` runs until cancelled.""" + + hang: bool = False + stopping: bool = False + + async def ensure_stopped(self) -> None: + self.stopping = self.hang + try: + if self.hang: + await asyncio.Event().wait() + finally: + self.stopping = False + + +async def test_unstage_that_fails_leaves_no_other_hook_running(): + logic = HangingStopLogic() + + class ReadableFlyer(StandardReadable, StandardFlyable[int, None]): + def __init__(self, name: str = ""): + with self.add_children_as_readables(): + self.child = FailsToUnstage() + super().__init__(name=name) + + @cached_property + def logic(self) -> FlyableLogic[int, None]: + return logic + + async with init_devices(mock=True): + device = ReadableFlyer(name="rf") + await device.stage() + logic.hang = True + + with pytest.raises(ValueError, match="unstage failed"): + await device.unstage() + + assert not logic.stopping + + +class FailsToStage(Device): + @AsyncStatus.wrap + async def stage(self) -> None: + raise ValueError("stage failed") + + @AsyncStatus.wrap + async def unstage(self) -> None: ... + + +@dataclass +class HangingReadyLogic(RecordingFlyableLogic): + """Its `ensure_ready` runs until cancelled.""" + + readying: bool = False + + async def ensure_ready(self) -> None: + self.readying = True + try: + await asyncio.Event().wait() + finally: + self.readying = False + + +async def test_stage_that_fails_leaves_no_other_hook_running(): + logic = HangingReadyLogic() + + class ReadableFlyer(StandardReadable, StandardFlyable[int, None]): + def __init__(self, name: str = ""): + with self.add_children_as_readables(): + self.child = FailsToStage() + super().__init__(name=name) + + @cached_property + def logic(self) -> FlyableLogic[int, None]: + return logic + + async with init_devices(mock=True): + device = ReadableFlyer(name="rf") + + with pytest.raises(ValueError, match="stage failed"): + await device.stage() + + assert not logic.readying + + +@dataclass +class SlowToCancelLogic(RecordingFlyableLogic): + """Its verbs take a moment to finish being cancelled.""" + + block: str = "" + running: set[str] = field(default_factory=set) + running_when_stopped: list[set[str]] = field(default_factory=list) + + async def _block(self, verb: str) -> None: + if self.block != verb: + return + self.running.add(verb) + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(0.05) + self.running.discard(verb) + + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[None, None]: + await self._block("prepare") + yield None + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + await self._block("kickoff") + yield KICKED_OFF + + async def ensure_stopped(self) -> None: + self.running_when_stopped.append(set(self.running)) + + +@pytest.mark.parametrize("verb", ["prepare", "kickoff"]) +async def test_unstage_waits_for_the_cancelled_verb_before_stopping(verb: str): + logic = SlowToCancelLogic(block=verb) + flyer = await make_staged_flyer(logic) + logic.running_when_stopped.clear() + if verb == "prepare": + running = flyer.prepare(1) + else: + await flyer.prepare(1) + running = flyer.kickoff() + await wait_for_pending_wakeups() + assert logic.running == {verb} + + await flyer.unstage() + + assert logic.running == set() + assert logic.running_when_stopped == [set()] + with pytest.raises(asyncio.CancelledError): + await running + + +@pytest.mark.parametrize( + "eager", + [ + False, + pytest.param( + True, + marks=pytest.mark.skipif( + sys.version_info < (3, 12), reason="eager_task_factory is 3.12+" + ), + ), + ], +) +async def test_complete_starts_with_the_latest_update_fly_yielded_before_it( + eager: bool, +): + if eager: + loop = asyncio.get_running_loop() + loop.set_task_factory(asyncio.eager_task_factory) # pyright: ignore + release = asyncio.Event() + + class UpdatingLogic(RecordingFlyableLogic): + async def fly( + self, ctx: None + ) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + yield KICKED_OFF + yield WatcherUpdate(current=1, initial=0, target=3) + yield WatcherUpdate(current=2, initial=0, target=3) + await release.wait() + yield WatcherUpdate(current=3, initial=0, target=3) + + flyer = await make_staged_flyer(UpdatingLogic()) + await flyer.prepare(1) + await flyer.kickoff() + await asyncio.sleep(0.01) + + watcher = Mock() + status = flyer.complete() + status.watch(watcher) + await asyncio.sleep(0.01) + release.set() + await status + + assert [c.kwargs["current"] for c in watcher.call_args_list] == [2, 3] + # complete() is the fly's own status, however often it is called + assert flyer.complete() is status + + +@dataclass +class RunUpLogic(RecordingFlyableLogic): + """Its prepare_to_fly runs up until cancelled, or fails part way if `fail`.""" + + fail: bool = False + run_ups: list[AsyncStatus] = field(default_factory=list) + closed: bool = False + + async def prepare_to_fly( + self, name: str, value: int, previous_ctx: None + ) -> AsyncGenerator[Any, None]: + yield value + try: + async with AsyncStatus(asyncio.Event().wait()) as run_up: + self.run_ups.append(run_up) + yield WatcherUpdate(current=0, initial=0, target=1) + if self.fail: + raise ValueError("bad run-up") + await run_up + finally: + self.closed = True + + +@pytest.mark.parametrize("fail", [False, True]) +async def test_prepare_that_fails_or_is_cancelled_stops_the_run_up(fail: bool): + logic = RunUpLogic(fail=fail) + flyer = await make_staged_flyer(logic) + preparing = flyer.prepare(1) + await wait_for_pending_wakeups() + + if fail: + with pytest.raises(ValueError, match="bad run-up"): + await preparing + else: + assert not logic.closed + preparing.task.cancel() + with pytest.raises(asyncio.CancelledError): + await preparing + + # The generator was closed and the run-up stopped before prepare() finished + assert logic.closed + assert logic.run_ups[0].task.cancelled() + # The half-failed prepare left no ctx to fly with + with pytest.raises(TransitionError, match="cannot go to FLYING from STAGED"): + await flyer.kickoff() + + +async def _update_first( + name: str, value: int, previous_ctx: None +) -> AsyncIterator[Any]: + yield WatcherUpdate(current=0, initial=0, target=1) + + +async def _no_ctx(name: str, value: int, previous_ctx: None) -> AsyncIterator[Any]: + return + yield + + +async def _ctx_then_not_an_update( + name: str, value: int, previous_ctx: None +) -> AsyncIterator[Any]: + yield value + yield 42 + + +@pytest.mark.parametrize( + "prepare_to_fly, error, match", + [ + ( + _update_first, + TypeError, + r"\.prepare_to_fly\(\) must yield its context first", + ), + (_no_ctx, RuntimeError, "finished without yielding its context"), + ( + _ctx_then_not_an_update, + TypeError, + r"\.prepare_to_fly\(\) must yield only WatcherUpdates", + ), + ], +) +async def test_prepare_to_fly_that_does_not_yield_its_context_first_raises( + recording_flyer, prepare_to_fly, error, match +): + logic, flyer = recording_flyer + logic.prepare_to_fly = prepare_to_fly + + with pytest.raises(error, match=match): + await flyer.prepare(1) + with pytest.raises(TransitionError, match="cannot go to FLYING from STAGED"): + await flyer.kickoff() + + +async def test_a_fly_that_fails_and_is_completed_is_not_logged( + caplog: pytest.LogCaptureFixture, +): + flyer = await make_staged_flyer(FlyingLogic(fail=True)) + await flyer.prepare(1) + await flyer.kickoff() + with pytest.raises(ValueError, match="bad fly"): + await flyer.complete() + + await flyer.unstage() + + assert not caplog.records diff --git a/tests/unit_tests/core/test_movable.py b/tests/unit_tests/core/test_movable.py index c3df8fd096..1e9c843b14 100644 --- a/tests/unit_tests/core/test_movable.py +++ b/tests/unit_tests/core/test_movable.py @@ -1,12 +1,20 @@ import asyncio +from collections.abc import AsyncGenerator +from dataclasses import dataclass, field from functools import cached_property from unittest.mock import AsyncMock, Mock, patch import pytest from ophyd_async.core import ( + KICKED_OFF, + FlyableLogic, MovableLogic, + StandardFlyable, StandardMovable, + StandardReadable, + StoppedError, + WatcherUpdate, callback_on_mock_put, init_devices, mock_puts_blocked, @@ -43,6 +51,11 @@ async def test_move_timeout_with_timeout_none_returns_none(): assert move_timeout() is None +class InstantStopMovableLogic(MovableLogic[float]): + async def ensure_stopped(self) -> None: + pass + + class StandardMovableImpl(StandardMovable[float]): def __init__(self, name: str = ""): self.readback, _ = soft_signal_r_and_setter(float) @@ -50,8 +63,8 @@ def __init__(self, name: str = ""): super().__init__(name=name) @cached_property - def movable_logic(self) -> MovableLogic[float]: - return MovableLogic(setpoint=self.setpoint, readback=self.readback) + def logic(self) -> MovableLogic[float]: + return InstantStopMovableLogic(setpoint=self.setpoint, readback=self.readback) @pytest.fixture @@ -61,17 +74,17 @@ async def movable() -> StandardMovableImpl: return movable -def test_movable_logic_is_cached(movable: StandardMovableImpl): - logic = movable.movable_logic - logic2 = movable.movable_logic +def test_logic_is_cached(movable: StandardMovableImpl): + logic = movable.logic + logic2 = movable.logic assert logic == logic2 async def test_movable_check_value(movable: StandardMovableImpl): - movable.movable_logic.check_move = AsyncMock() + movable.logic.check_move = AsyncMock() await movable.check_value(5) - movable.movable_logic.check_move.assert_awaited_once_with(5) + movable.logic.check_move.assert_awaited_once_with(5) async def test_locatable(movable: StandardMovableImpl) -> None: @@ -130,15 +143,15 @@ async def test_movable_moving_stopped(movable: StandardMovableImpl): await move_status -async def test_movable_set_calls_movable_logic_check_move_and_calculate_timeout( +async def test_movable_set_calls_logic_check_move_and_calculate_timeout( movable: StandardMovableImpl, ): - mock_check_move = movable.movable_logic.check_move = AsyncMock() + mock_check_move = movable.logic.check_move = AsyncMock() timeout = 5 - mock_calculate_timeout = movable.movable_logic.calculate_timeout = AsyncMock( + mock_calculate_timeout = movable.logic.calculate_timeout = AsyncMock( return_value=timeout ) - mock_move = movable.movable_logic.move = AsyncMock() + mock_move = movable.logic.move = AsyncMock() with patch("ophyd_async.core._movable.MoveTimeout") as move_timeout: pos = 10 @@ -173,3 +186,240 @@ async def test_motor_set_with_instant_mock( await status assert status.success assert await movable.readback.get_value() == -50.0 + + +async def test_movable_logic_stop_is_deprecated_but_still_called(): + with pytest.deprecated_call(match=r"OldStopLogic.stop\(\) is deprecated"): + + class OldStopLogic(MovableLogic[float]): + async def stop(self) -> None: + stopped.append(True) + + stopped: list[bool] = [] + + class OldStopMovable(StandardMovable[float]): + def __init__(self, name: str = ""): + self.readback, _ = soft_signal_r_and_setter(float) + self.setpoint = soft_signal_rw(float) + super().__init__(name=name) + + @cached_property + def logic(self) -> MovableLogic[float]: + # The deprecated stop() is turned into ensure_stopped() when the class + # is made, which pyright can not see + return OldStopLogic( # pyright: ignore[reportAbstractUsage] + setpoint=self.setpoint, readback=self.readback + ) + + async with init_devices(mock=True): + movable = OldStopMovable() + await movable.stop() + assert stopped == [True] + + +@dataclass +class HeldMoveLogic(MovableLogic[float], FlyableLogic[None, None]): + """Moves until released or cancelled; records when it was stopped.""" + + moving: asyncio.Event = field(default_factory=asyncio.Event) + release: asyncio.Event = field(default_factory=asyncio.Event) + calls: list[str] = field(default_factory=list) + + async def move(self, new_position: float, timeout) -> None: + self.calls.append(f"move {new_position}") + self.moving.set() + await self.release.wait() + set_mock_value(self.setpoint, new_position) + + async def ensure_stopped(self) -> None: + self.calls.append("ensure_stopped") + + async def prepare_to_fly( + self, name: str, value: None, previous_ctx: None + ) -> AsyncGenerator[None, None]: + yield None + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + yield KICKED_OFF + + +class StageableMovable(StandardReadable, StandardMovable[float]): + def __init__(self, name: str = ""): + self.readback, _ = soft_signal_r_and_setter(float) + self.setpoint = soft_signal_rw(float) + super().__init__(name=name) + + @cached_property + def logic(self) -> HeldMoveLogic: + return HeldMoveLogic(setpoint=self.setpoint, readback=self.readback) + + +class FlyableMovable(StageableMovable, StandardFlyable[None, None]): + pass + + +@pytest.fixture +async def held() -> FlyableMovable: + async with init_devices(mock=True): + held = FlyableMovable(name="movable") + return held + + +@pytest.mark.parametrize("success", [False, True]) +async def test_stop_ends_a_set_with_or_without_error( + held: FlyableMovable, success: bool +): + move_status = held.set(1.0) + await held.logic.moving.wait() + + await held.stop(success=success) + + if success: + await move_status + else: + with pytest.raises(StoppedError, match="Device movable was stopped"): + await move_status + assert held.logic.calls == ["move 1.0", "ensure_stopped"] + + +async def test_stop_before_any_set_and_twice_is_harmless(held: FlyableMovable): + await held.stop() + move_status = held.set(1.0) + await held.logic.moving.wait() + await held.stop() + await held.stop() + with pytest.raises(StoppedError, match="was stopped"): + await move_status + # A later set is not affected by the earlier stops + held.logic.release.set() + await held.set(2.0) + assert held.logic.calls == [ + "ensure_stopped", + "move 1.0", + "ensure_stopped", + "ensure_stopped", + "move 2.0", + ] + + +async def test_stop_before_the_set_starts_leaves_the_hardware_alone( + held: FlyableMovable, +): + move_status = held.set(1.0) + await held.stop() + with pytest.raises(StoppedError, match="was stopped"): + await move_status + assert held.logic.calls == ["ensure_stopped"] + + +async def test_second_set_supersedes_the_first(held: FlyableMovable): + first = held.set(1.0) + await held.logic.moving.wait() + second = held.set(2.0) + + with pytest.raises(StoppedError, match=r"movable: superseded by a new set\(\)"): + await first + held.logic.release.set() + await second + assert await held.setpoint.get_value() == 2.0 + # Neither set stopped the hardware: the second move retargets it + assert held.logic.calls == ["move 1.0", "move 2.0"] + + +# A timeout of the logic's own is not mistaken for a stop +async def test_a_timeout_inside_move_raises_timeout_error( + movable: StandardMovableImpl, +): + async def move(new_position: float, timeout) -> None: + async with asyncio.timeout(0.01): + await asyncio.Event().wait() + + movable.logic.move = move + with pytest.raises(TimeoutError): + await movable.set(1.0) + + +@pytest.mark.parametrize("flies", [False, True]) +async def test_only_a_flyer_stops_a_set_on_stage_and_unstage(flies: bool): + async with init_devices(mock=True): + movable = FlyableMovable() if flies else StageableMovable() + move_status = movable.set(1.0) + await movable.logic.moving.wait() + + await movable.stage() + await movable.unstage() + + # A pure movable stages nothing; a flyer gets ready once on stage and stops on + # unstage, which also halt a set() in progress without ending it + assert movable.logic.calls == ( + ["move 1.0", "ensure_stopped", "ensure_stopped"] if flies else ["move 1.0"] + ) + assert not move_status.done + movable.logic.release.set() + await move_status + + +async def test_unstage_stops_the_hardware_of_a_set_in_progress(held: FlyableMovable): + await held.stage() + move_status = held.set(1.0) + await held.logic.moving.wait() + + await held.unstage() + + # Never prepared, but stopped anyway; only stop() ends the set() itself + assert not move_status.done + held.logic.release.set() + await move_status + assert held.logic.calls == ["ensure_stopped", "move 1.0", "ensure_stopped"] + + +@dataclass +class HeldFlyLogic(HeldMoveLogic): + """Its prepare and fly run until released.""" + + async def prepare_to_fly( + self, name: str, value: None, previous_ctx: None + ) -> AsyncGenerator[None, None]: + yield None + await self.release.wait() + + async def fly(self, ctx: None) -> AsyncGenerator[KICKED_OFF | WatcherUpdate, None]: + yield KICKED_OFF + await self.release.wait() + + +class HeldFlyer(FlyableMovable): + @cached_property + def logic(self) -> HeldFlyLogic: + return HeldFlyLogic(setpoint=self.setpoint, readback=self.readback) + + +@pytest.mark.parametrize("success", [False, True]) +async def test_stop_halts_the_hardware_but_does_not_end_a_prepare_or_fly( + success: bool, +): + async with init_devices(mock=True): + flyer = HeldFlyer(name="flyer") + await flyer.stage() + preparing = flyer.prepare(None) + await wait_for_pending_wakeups() + + await flyer.stop(success=success) + + # Only the hardware was halted: unstage() and a new prepare() end a prepare() + assert not preparing.done + flyer.logic.release.set() + await preparing + flyer.logic.release.clear() + await flyer.kickoff() + flying = flyer.complete() + await wait_for_pending_wakeups() + + await flyer.stop(success=success) + + assert not flying.done + flyer.logic.release.set() + await flying + # stage() got it ready once, then the two stop()s + assert flyer.logic.calls.count("ensure_stopped") == 3 + await flyer.unstage() diff --git a/tests/unit_tests/core/test_multi_derived_signal.py b/tests/unit_tests/core/test_multi_derived_signal.py index dd252af6fa..3f8b038301 100644 --- a/tests/unit_tests/core/test_multi_derived_signal.py +++ b/tests/unit_tests/core/test_multi_derived_signal.py @@ -298,8 +298,8 @@ async def _put(val: float) -> None: derived_sig = derived_signal_rw(_get, _put, m1=m1, m2=m2) await derived_sig.set(6) - assert await m1.movable_logic.readback.get_value() == 3 - assert await m2.movable_logic.readback.get_value() == 3 + assert await m1.logic.readback.get_value() == 3 + assert await m2.logic.readback.get_value() == 3 assert await derived_sig.get_value() == 6 diff --git a/tests/unit_tests/core/test_readable.py b/tests/unit_tests/core/test_readable.py index c0d682bc1d..5071553b69 100644 --- a/tests/unit_tests/core/test_readable.py +++ b/tests/unit_tests/core/test_readable.py @@ -15,6 +15,7 @@ MockSignalBackend, SignalR, StandardReadable, + TransitionError, soft_signal_r_and_setter, soft_signal_rw, ) @@ -152,10 +153,7 @@ def test_standard_readable_add_children_context_manager(): sr.b = MagicMock(spec=SignalR) sr.c = MagicMock(spec=SignalR) - # Can't use assert_called_once_with() as the order of items returned from - # internal dict comprehension is not guaranteed - mock.assert_called_once() - assert set(mock.call_args.args[0]) == {sr.a, sr.b, sr.c} + mock.assert_called_once_with([sr.a, sr.b, sr.c], Format.CHILD) @pytest.mark.parametrize( @@ -177,10 +175,7 @@ def test_standard_readable_add_children_cm_device_with_mappings(device_type, key with sr.add_children_as_readables(): sr.a = device - # Can't use assert_called_once_with() as the order of items returned from - # internal dict comprehension is not guaranteed - mock.assert_called_once() - assert set(mock.call_args.args[0]) == set(devices) + mock.assert_called_once_with(devices, Format.CHILD) def test_standard_readable_add_children_cm_filters_non_devices(): @@ -195,10 +190,7 @@ def test_standard_readable_add_children_cm_filters_non_devices(): sr.d = "abc" sr.e = MagicMock(spec=MockSignalBackend) - # Can't use assert_called_once_with() as the order of items returned from - # internal dict comprehension is not guaranteed - mock.assert_called_once() - assert set(mock.call_args.args[0]) == {sr.a, sr.b} + mock.assert_called_once_with([sr.a, sr.b], Format.CHILD) def assert_sr_has_attrs(sr: StandardReadable, expected_attrs: dict[str, tuple]): @@ -348,3 +340,13 @@ def __init__(self, name): with pytest.raises(KeyError): DummyDerivedDevice("test_duplicates") + + +async def test_standard_readable_stage_twice_raises(): + sr = StandardReadable(name="sr") + await sr.stage() + # Like ophyd's RedundantStaging + with pytest.raises(TransitionError, match="cannot go to STAGED from STAGED"): + await sr.stage() + await sr.unstage() + await sr.stage() diff --git a/tests/unit_tests/core/test_signal.py b/tests/unit_tests/core/test_signal.py index e9a7b7a973..7792315f3c 100644 --- a/tests/unit_tests/core/test_signal.py +++ b/tests/unit_tests/core/test_signal.py @@ -64,6 +64,7 @@ assert_reading, assert_value, partial_reading, + wait_for_pending_wakeups, ) _array_vals = { @@ -1344,3 +1345,36 @@ async def test_soft_signal_r_and_setter_with_poll_period(): async def test_soft_signal_r_and_setter_poll_period_without_getter_raises(): with pytest.raises(ValueError, match="poll_period requires a getter"): soft_signal_r_and_setter(float, poll_period=0.1) + + +async def test_set_and_wait_for_other_value_cancelled_before_first_value(): + set_signal = epics_signal_rw(int, "pva://signal", name="s") + match_signal = epics_signal_rw(int, "pva://match_signal", name="m") + await set_signal.connect(mock=True) + await match_signal.connect(mock=True) + + async def never_yields(signal): + await asyncio.sleep(1000) + yield # pragma: no cover - unreachable, just makes this a generator + + with patch("ophyd_async.core._signal.observe_value", never_yields): + task = asyncio.create_task( + set_and_wait_for_other_value(set_signal, 1, match_signal, 20, timeout=None) + ) + await wait_for_pending_wakeups() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + # The match wait it started must not outlive it + assert asyncio.all_tasks() == {asyncio.current_task()} + + +async def test_wait_for_value_unsubscribes_before_it_returns(): + signal = soft_signal_rw(int, initial_value=1) + await signal.connect() + with patch.object( + SignalR, "clear_sub", autospec=True, side_effect=SignalR.clear_sub + ) as clear_sub: + await wait_for_value(signal, 1, timeout=1) + # Not left for the abandoned generators to do when they are finalised + clear_sub.assert_called_once() diff --git a/tests/unit_tests/core/test_standard_detector.py b/tests/unit_tests/core/test_standard_detector.py index 0c9b9ab814..96d449feb7 100644 --- a/tests/unit_tests/core/test_standard_detector.py +++ b/tests/unit_tests/core/test_standard_detector.py @@ -720,7 +720,7 @@ async def test_streamable_supports_both_step_and_fly(tmp_path): # Yield so detector can get collections written, then set it so we complete await wait_for_pending_wakeups(raise_if_exceeded=False) await dl.collections_written.set(1) - assert status.done + await status assert status.success docs = [doc async for doc in det.collect_asset_docs()] assert docs == [ diff --git a/tests/unit_tests/core/test_status.py b/tests/unit_tests/core/test_status.py index 767b51ba23..54627656d3 100644 --- a/tests/unit_tests/core/test_status.py +++ b/tests/unit_tests/core/test_status.py @@ -413,3 +413,27 @@ async def test_forward_watcher_updates(moves: dict[str, list], combine: bool): # Updates from different children interleave, but each child's stay in order for name, values in moves.items(): assert [u["current"] for u in updates if u["name"] == name] == values + + +@pytest.mark.parametrize("with_stop", [False, True]) +async def test_a_watcher_that_raises_fails_the_status_and_closes_the_iterator( + with_stop: bool, +): + closed = False + + async def updates(): + nonlocal closed + try: + yield WatcherUpdate(current=0, initial=0, target=2, name="x") + finally: + closed = True + + stop = asyncio.get_running_loop().create_future() if with_stop else None + status = WatchableAsyncStatus(updates(), stop=stop) + status.watch(Mock(side_effect=ValueError("bad watcher"))) + + with pytest.raises(ValueError, match="bad watcher"): + await status + + assert not status.success + assert closed diff --git a/tests/unit_tests/core/test_utils.py b/tests/unit_tests/core/test_utils.py index 378d4449a6..7b7a5e844b 100644 --- a/tests/unit_tests/core/test_utils.py +++ b/tests/unit_tests/core/test_utils.py @@ -1,7 +1,9 @@ import asyncio +import gc import re import sys from asyncio import CancelledError +from contextlib import nullcontext from typing import TypeVar, assert_type import numpy as np @@ -17,12 +19,14 @@ SignalRW, AsyncStatus, callback_on_mock_put, + cancel_and_wait, get_mock_put, merge_gathered_dicts, set_mock_value, simulate_move, soft_signal_rw, ) +from ophyd_async.core._utils import _wait_for from ophyd_async.epics.core import epics_signal_rw from ophyd_async.epics.motor import Motor @@ -482,3 +486,202 @@ async def test_simulate_move_ends_at_target(acceleration_time: float): positions = [p async for p in simulate_move(0, 1.5, 10, acceleration_time)] assert positions[-1] == 1.5 assert positions == sorted(positions) + + +class StuckPutBackend(SoftSignalBackend): + """A put that never finishes, and may fail while being cancelled.""" + + def __init__(self, cleanup_error: Exception | None): + self.cleanup_error = cleanup_error + super().__init__(datatype=int) + + async def put(self, value: int | None) -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + if self.cleanup_error is not None: + raise self.cleanup_error from None + raise + + +@pytest.mark.parametrize( + "cleanup_error, expected", + [(None, TimeoutError), (ValueError("cleanup failed"), ValueError)], +) +async def test_a_timed_out_put_raises_the_error_its_cancellation_raised( + cleanup_error: Exception | None, expected: type[Exception] +): + loop = asyncio.get_running_loop() + unhandled = [] + loop.set_exception_handler(lambda _, context: unhandled.append(context)) + signal = SignalRW(StuckPutBackend(cleanup_error), name="sig") + await signal.connect() + + with pytest.raises(expected): + await signal.set(1, timeout=0.01) + + # The error was retrieved, not left for the loop to log + gc.collect() + await asyncio.sleep(0) + loop.set_exception_handler(None) + assert unhandled == [] + + +async def test_a_cancelled_wait_for_leaves_no_unretrieved_error(): + loop = asyncio.get_running_loop() + unhandled = [] + loop.set_exception_handler(lambda _, context: unhandled.append(context)) + + async def fails_when_cancelled(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + raise ValueError("cleanup failed") from None + + outer = asyncio.create_task(_wait_for(fails_when_cancelled(), 10, "source")) + await asyncio.sleep(0.01) + outer.cancel() + # The error raised while ending the cancelled work shows through + with pytest.raises(ValueError, match="cleanup failed"): + await outer + + gc.collect() + await asyncio.sleep(0) + loop.set_exception_handler(None) + assert unhandled == [] + + +async def test_wait_for_returns_the_result_with_and_without_a_timeout(): + async def value(): + await asyncio.sleep(0) + return 42 + + assert await _wait_for(value(), 5, "source") == 42 + assert await _wait_for(value(), None, "source") == 42 + + +async def test_wait_for_expiry_raises_timeout_error_naming_the_source(): + with pytest.raises(TimeoutError, match="^source$"): + await _wait_for(asyncio.sleep(1), 0.01, "source") + + +async def test_wait_for_relabels_a_timeout_error_from_the_coroutine(): + async def inner(): + raise TimeoutError("inner") + + with pytest.raises(TimeoutError, match="^source$") as exc_info: + await _wait_for(inner(), 5, "source") + assert str(exc_info.value.__cause__) == "inner" + + +async def test_wait_for_does_not_swallow_a_cancel_arriving_as_the_work_finishes(): + finished = asyncio.get_running_loop().create_future() + + async def inner(): + return await finished + + outer = asyncio.create_task(_wait_for(inner(), 5, "source")) + await asyncio.sleep(0) + finished.set_result(1) + outer.cancel() + with pytest.raises(asyncio.CancelledError): + await outer + + +async def test_wait_for_shows_an_error_raised_while_being_cancelled(): + async def inner(): + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + raise ValueError("cleanup") from None + + with pytest.raises(ValueError, match="cleanup"): + await _wait_for(inner(), 0.01, "source") + + +async def test_wait_for_expiry_and_an_outside_cancel_is_a_cancel(): + outer = asyncio.create_task(_wait_for(asyncio.sleep(1), 0.05, "source")) + await asyncio.sleep(0.05) # About when the timeout expires + outer.cancel() + with pytest.raises(asyncio.CancelledError): + await outer + + +async def test_wait_for_expiry_after_a_swallowed_cancel_is_still_a_timeout_error(): + async def body(): + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + pass # Swallowed without uncancel() + await _wait_for(asyncio.sleep(1), 0.01, "source") + + task = asyncio.create_task(body()) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(TimeoutError, match="^source$"): + await task + + +async def _raise_on_cancel( + error: Exception | None, cleanup_time: float, cancelled: asyncio.Event +) -> None: + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + cancelled.set() + await asyncio.sleep(cleanup_time) + if error is not None: + raise error from None + raise + + +@pytest.mark.parametrize( + "errors, raised", + [ + ([None, None], None), + ([None, ValueError()], ValueError), + # The first task in the order given wins, though it ends last + ([ValueError(), RuntimeError()], ValueError), + ], +) +async def test_cancel_and_wait_waits_for_every_task_then_raises_the_first_error( + errors: list[Exception | None], raised: type[Exception] | None +): + cleanup_times = [0.02, 0.0] + tasks = [ + asyncio.create_task(_raise_on_cancel(error, cleanup_time, asyncio.Event())) + for error, cleanup_time in zip(errors, cleanup_times, strict=True) + ] + await asyncio.sleep(0) # let the tasks start + expectation = pytest.raises(raised) if raised else nullcontext() + with expectation: + await cancel_and_wait(*tasks) + assert all(task.done() for task in tasks) + + +@pytest.mark.parametrize("caller_cancels", [1, 2]) +async def test_cancel_and_wait_cancelled_caller_waits_then_raises_its_cancel( + caller_cancels: int, +): + loop = asyncio.get_running_loop() + unhandled: list[str] = [] + old_handler = loop.get_exception_handler() + loop.set_exception_handler(lambda _, context: unhandled.append(context["message"])) + try: + cancelled = asyncio.Event() + task = asyncio.create_task(_raise_on_cancel(ValueError(), 0.02, cancelled)) + await asyncio.sleep(0) # let the task start + caller = asyncio.create_task(cancel_and_wait(task)) + await cancelled.wait() + for _ in range(caller_cancels): + caller.cancel() + await asyncio.sleep(0) + await asyncio.wait([caller]) + assert caller.cancelled() + assert task.done() + # A task whose error was never retrieved logs it when collected + del task, caller + gc.collect() + finally: + loop.set_exception_handler(old_handler) + assert unhandled == [] diff --git a/tests/unit_tests/core/test_watchable_async_status.py b/tests/unit_tests/core/test_watchable_async_status.py index cabdc18b44..114bb4367a 100644 --- a/tests/unit_tests/core/test_watchable_async_status.py +++ b/tests/unit_tests/core/test_watchable_async_status.py @@ -216,3 +216,196 @@ async def set(self, value) -> AsyncIterator: # undecorated methods will not print the device name status: AsyncStatus = ctx.value.args[0] assert f"device: {device_name}" in repr(status) + + +class StopError(Exception): + pass + + +class Ticks: + """Yields an update then waits forever, or until `finish` is resolved.""" + + def __init__(self, finish: asyncio.Future | None = None): + self.finish = finish + self.started = False + + def __aiter__(self): + return self + + async def __anext__(self) -> WatcherUpdate[int]: + if self.started: + await (self.finish or asyncio.sleep(10)) + raise StopAsyncIteration + self.started = True + return WatcherUpdate(current=1, initial=0, target=2) + + +@pytest.mark.parametrize("error", [None, StopError("stopped")]) +@pytest.mark.parametrize("before_start", [True, False]) +async def test_stop_future_ends_a_watchable_status(error, before_start): + ticks = Ticks() + stop = asyncio.get_running_loop().create_future() + updates: list[dict] = [] + if before_start: + stop.set_result(error) + status = WatchableAsyncStatus(ticks, stop=stop) + status.watch(lambda **kw: updates.append(kw)) + if not before_start: + await asyncio.sleep(0.01) + assert not status.done + stop.set_result(error) + if error is None: + await status + assert status.success + else: + with pytest.raises(StopError): + await status + # Stopped before it started, so the iterator body never ran + assert ticks.started is not before_start + assert len(updates) == (0 if before_start else 1) + + +async def test_no_stop_future_adds_nothing(): + tasks = asyncio.all_tasks() + ticks = Ticks(finish=asyncio.get_running_loop().create_future()) + status = WatchableAsyncStatus(ticks) + assert asyncio.all_tasks() - tasks == {status.task} + await asyncio.sleep(0.01) + assert asyncio.all_tasks() - tasks == {status.task} + assert ticks.finish + ticks.finish.set_result(None) + await status + assert status.success + + +@pytest.mark.parametrize("stop_first", [True, False]) +async def test_outside_cancel_with_stop_in_the_same_iteration_cancels(stop_first): + stop = asyncio.get_running_loop().create_future() + status = WatchableAsyncStatus(Ticks(), stop=stop) + await asyncio.sleep(0.01) + if stop_first: + stop.set_result(None) + status.task.cancel() + else: + status.task.cancel() + stop.set_result(None) + with pytest.raises(asyncio.CancelledError): + await status + assert status.task.cancelled() + + +@pytest.mark.parametrize("yields", range(5)) +async def test_outside_cancel_while_the_iterator_is_finishing_is_not_lost(yields): + loop = asyncio.get_running_loop() + finish, stop = loop.create_future(), loop.create_future() + ticks = Ticks(finish) + status = WatchableAsyncStatus(ticks, stop=stop) + await asyncio.sleep(0.01) + finish.set_result(None) + for _ in range(yields): + await asyncio.sleep(0) + if status.task.cancel(): + with pytest.raises(asyncio.CancelledError): + await status + assert status.task.cancelled() + else: + await status + + +async def test_timeout_inside_a_stoppable_watchable_status_raises_timeout_error(): + async def iterator(): + async with asyncio.timeout(0.01): + yield WatcherUpdate(current=1, initial=0, target=2) + await asyncio.sleep(10) + + stop = asyncio.get_running_loop().create_future() + status = WatchableAsyncStatus(iterator(), stop=stop) + with pytest.raises(TimeoutError): + await status + assert not stop.done() + + +async def test_stop_future_decides_even_if_the_iterator_absorbs_the_cancel(): + async def iterator(): + yield WatcherUpdate(current=1, initial=0, target=2) + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + pass + yield WatcherUpdate(current=2, initial=0, target=2) + + stop = asyncio.get_running_loop().create_future() + status = WatchableAsyncStatus(iterator(), stop=stop) + await asyncio.sleep(0.01) + stop.set_result(StopError("stopped")) + with pytest.raises(StopError): + await status + + +class MoveError(Exception): + pass + + +@pytest.mark.parametrize("error", [None, MoveError("failed")]) +async def test_an_iterator_ending_in_the_same_iteration_as_the_stop_wins(error): + loop = asyncio.get_running_loop() + finish, stop = loop.create_future(), loop.create_future() + + async def iterator(): + yield WatcherUpdate(current=1, initial=0, target=2) + await finish + if error: + raise error + + status = WatchableAsyncStatus(iterator(), stop=stop) + await asyncio.sleep(0.01) + finish.set_result(None) + # The iterator resumes before the stop resolves, so both are done by the + # time the race looks + await asyncio.sleep(0) + stop.set_result(StopError("stopped")) + if error: + with pytest.raises(MoveError): + await status + else: + await status + assert status.success + + +async def test_an_error_raised_while_stopping_shows_through(): + async def iterator(): + yield WatcherUpdate(current=1, initial=0, target=2) + try: + await asyncio.sleep(10) + except asyncio.CancelledError as e: + raise MoveError("failed to stop") from e + + stop = asyncio.get_running_loop().create_future() + status = WatchableAsyncStatus(iterator(), stop=stop) + await asyncio.sleep(0.01) + stop.set_result(StopError("stopped")) + with pytest.raises(MoveError): + await status + + +async def test_a_second_outside_cancel_waits_for_the_iterator_cleanup(): + cleaned_up = asyncio.Event() + + async def iterator(): + try: + yield WatcherUpdate(current=1, initial=0, target=2) + await asyncio.sleep(10) + finally: + # Slow cleanup, which a second cancel must not abandon + await asyncio.sleep(0.05) + cleaned_up.set() + + stop = asyncio.get_running_loop().create_future() + status = WatchableAsyncStatus(iterator(), stop=stop) + await asyncio.sleep(0.01) + status.task.cancel() + await asyncio.sleep(0.01) + status.task.cancel() + with pytest.raises(asyncio.CancelledError): + await status + assert cleaned_up.is_set() diff --git a/tests/unit_tests/epics/adcore/test_use_cases.py b/tests/unit_tests/epics/adcore/test_use_cases.py index d3cb2047f0..e215cbb46a 100644 --- a/tests/unit_tests/epics/adcore/test_use_cases.py +++ b/tests/unit_tests/epics/adcore/test_use_cases.py @@ -373,7 +373,7 @@ async def test_flyscan_aravis_detector(static_path_provider: StaticPathProvider) }, ), ] - assert status.done + await status # Check that another kickoff without prepare is not allowed with pytest.raises( RuntimeError, @@ -399,6 +399,7 @@ async def test_flyscan_aravis_detector(static_path_provider: StaticPathProvider) }, ), ] + await status async def test_2_rois_with_hdf(tmp_path): diff --git a/tests/unit_tests/epics/pmac/test_pmac_trajectory.py b/tests/unit_tests/epics/pmac/test_pmac_trajectory.py index f2ee9bf8dc..c5ed3cb17f 100644 --- a/tests/unit_tests/epics/pmac/test_pmac_trajectory.py +++ b/tests/unit_tests/epics/pmac/test_pmac_trajectory.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock, call, patch import numpy as np @@ -8,6 +9,7 @@ callback_on_mock_execute, get_mock, get_mock_execute, + mock_puts_blocked, set_and_wait_for_value, set_callback_filter, set_mock_value, @@ -39,6 +41,7 @@ async def test_pmac_prepare(sim_motors: tuple[PmacIO, Motor, Motor]): spec = Fly(2.0 @ Line(sim_x_motor, 1, 5, 2)) value = PmacScanInfo(spec=spec, ramp_time=None, turnaround_time=None) flyer = PmacTrajectoryFlyableLogic(pmac_io).with_device() + await flyer.stage() await flyer.prepare(value) assert await pmac_io.coord[1].cs_axis_setpoint[7].get_value() == -1.2 @@ -65,6 +68,7 @@ async def test_pmac_prepare_with_configured_ramp( spec = Fly(2.0 @ Line(sim_x_motor, 1, 5, 2)) value = PmacScanInfo(spec=spec, ramp_time=2, turnaround_time=None) flyer = PmacTrajectoryFlyableLogic(pmac_io).with_device() + await flyer.stage() await flyer.prepare(value) assert await pmac_io.coord[1].cs_axis_setpoint[7].get_value() == -3.0 @@ -91,6 +95,7 @@ async def test_pmac_prepare_with_configured_ramp_and_turnaround( spec = Fly(2.0 @ (2 * ~Line(sim_x_motor, 1, 5, 2))) value = PmacScanInfo(spec=spec, ramp_time=2, turnaround_time=3) flyer = PmacTrajectoryFlyableLogic(pmac_io).with_device() + await flyer.stage() await flyer.prepare(value) assert await pmac_io.coord[1].cs_axis_setpoint[7].get_value() == -3.0 @@ -224,6 +229,7 @@ async def test_pmac_trajectory_kickoff( value = PmacScanInfo(spec=spec, ramp_time=None, turnaround_time=None) with patch("ophyd_async.epics.pmac._pmac_trajectory.SLICE_SIZE", 2): # This will prepare the buffer with 2 frames of info + await flyer.stage() await flyer.prepare(value) # This will consume another 2 frames set_mock_value( @@ -332,21 +338,38 @@ async def test_pmac_trajectory_kickoff( ) -async def test_pmac_trajectory_kickoff_trajectory_raises_exception_if_no_prepare( - sim_motors: tuple[PmacIO, Motor, Motor], +@pytest.mark.parametrize("run_up_ends_by", ["timeout", "unstage"]) +async def test_pmac_kickoff_that_fails_during_run_up_stops_the_trajectory( + sim_motors: tuple[PmacIO, Motor, Motor], run_up_ends_by: str ): - # Ordering is enforced by StandardFlyable, so drive through the device. - pmac_io, _, _ = sim_motors - flyer = PmacTrajectoryFlyableLogic(pmac_io).with_device() - with pytest.raises(RuntimeError, match="prepare.* before kickoff"): - await flyer.kickoff() - - -async def test_pmac_trajectory_complete(sim_motors: tuple[PmacIO, Motor, Motor]): - pmac_io, _, _ = sim_motors + pmac_io, sim_x_motor, sim_y_motor = sim_motors flyer = PmacTrajectoryFlyableLogic(pmac_io).with_device() - with pytest.raises(RuntimeError, match="kickoff.* before complete"): - await flyer.complete() + spec = Fly(2.0 @ (Line(sim_y_motor, 1, 5, 2) * ~Line(sim_x_motor, 1, 5, 2))) + value = PmacScanInfo(spec=spec, ramp_time=None, turnaround_time=None) + positions = get_mock(pmac_io.trajectory.positions[7]) + with patch("ophyd_async.epics.pmac._pmac_trajectory.SLICE_SIZE", 2): + await flyer.stage() + await flyer.prepare(value) + appends = len(positions.mock_calls) + # Keep the trajectory executing so it would go on appending + with mock_puts_blocked(pmac_io.trajectory.execute_profile): + if run_up_ends_by == "timeout": + with patch( + "ophyd_async.epics.pmac._pmac_trajectory.wait_for_value", + AsyncMock(side_effect=TimeoutError), + ): + with pytest.raises(TimeoutError): + await flyer.kickoff() + else: + status = flyer.kickoff() + await asyncio.sleep(0.01) + await flyer.unstage() + with pytest.raises(asyncio.CancelledError): + await status + # A trajectory still running would refill the buffer as it empties + set_mock_value(pmac_io.trajectory.total_points, 2) + await asyncio.sleep(0.01) + assert len(positions.mock_calls) == appends async def test_pmac_trajectory_stage(sim_motors: tuple[PmacIO, Motor, Motor]): @@ -377,19 +400,25 @@ async def test_pmac_trajectory_stage(sim_motors: tuple[PmacIO, Motor, Motor]): async def test_pmac_trajectory_unstage(sim_motors: tuple[PmacIO, Motor, Motor]): pmac_io, _, _ = sim_motors logic = PmacTrajectoryFlyableLogic(pmac_io) - logic.stop = AsyncMock() + logic.ensure_stopped = AsyncMock() flyer = logic.with_device() await flyer.unstage() - logic.stop.assert_called_once() + # Unstage stops even though it was not prepared since stage + logic.ensure_stopped.assert_awaited_once() async def test_trajectory_stop_if_running(sim_motors: tuple[PmacIO, Motor, Motor]): pmac_io, _, _ = sim_motors pmac_trajectory = PmacTrajectoryFlyableLogic(pmac_io) execute_mock = get_mock_execute(pmac_io.trajectory.abort_profile) + # Aborting stops the trajectory + callback_on_mock_execute( + pmac_io.trajectory.abort_profile, + lambda: set_mock_value(pmac_io.trajectory.execute_state, PmacExecuteState.DONE), + ) # Method not called as no running trajectory - await pmac_trajectory.stop() + await pmac_trajectory.ensure_stopped() execute_mock.assert_not_awaited() # Mocking that trajectory is executing @@ -397,9 +426,10 @@ async def test_trajectory_stop_if_running(sim_motors: tuple[PmacIO, Motor, Motor pmac_trajectory.pmac.trajectory.execute_state, PmacExecuteState.EXECUTING ) - # Method called as there is now a running trajectory - await pmac_trajectory.stop() + # Method called as there is now a running trajectory, and waits for it to stop + await pmac_trajectory.ensure_stopped() execute_mock.assert_awaited_once_with() + assert await pmac_io.trajectory.execute_state.get_value() == PmacExecuteState.DONE async def test_trajectory_raises_if_profile_status_not_in_good_state( diff --git a/tests/unit_tests/epics/test_motor.py b/tests/unit_tests/epics/test_motor.py index eeed03faf0..f8634312ec 100644 --- a/tests/unit_tests/epics/test_motor.py +++ b/tests/unit_tests/epics/test_motor.py @@ -1,14 +1,17 @@ import asyncio -from unittest.mock import Mock +import contextlib +from unittest.mock import Mock, patch import pytest from bluesky.protocols import Reading from ophyd_async.core import ( CALCULATE_TIMEOUT, + DEFAULT_TIMEOUT, Device, DeviceMock, FlyMotorInfo, + StoppedError, callback_on_mock_put, default_mock_class, get_mock_put, @@ -47,6 +50,11 @@ async def motor(): set_mock_value(motor.high_limit_travel, 21) set_mock_value(motor.dial_low_limit_travel, -11) set_mock_value(motor.dial_high_limit_travel, 21) + # Not moving, and a STOP stops it + set_mock_value(motor.motor_done_move, 1) + callback_on_mock_put( + motor.motor_stop, lambda _: set_mock_value(motor.motor_done_move, 1) + ) yield motor @@ -157,6 +165,15 @@ async def test_motor_moving_stopped(motor: Motor): assert move_status.done +async def test_stage_and_unstage_each_send_stop_even_when_not_moving(motor: Motor): + set_mock_value(motor.motor_done_move, 1) + await motor.stage() + get_mock_put(motor.motor_stop).assert_called_once_with(1) + + await motor.unstage() + assert get_mock_put(motor.motor_stop).call_count == 2 + + async def test_read_motor(motor: Motor): await motor.stage() assert (await motor.read())["motor"]["value"] == 0.0 @@ -244,6 +261,7 @@ async def test_set(motor: Motor, setpoint, velocity, timeout) -> None: async def test_prepare_velocity_limit_error(motor: Motor): + await motor.stage() set_mock_value(motor.max_velocity, 10) with pytest.raises(MotorLimitsError): fly_info = FlyMotorInfo(start_position=-10, end_position=0, time_for_move=0.9) @@ -251,6 +269,7 @@ async def test_prepare_velocity_limit_error(motor: Motor): async def test_valid_prepare_velocity(motor: Motor): + await motor.stage() # Widen user limits to accommodate ramp distances (trajectory goes to -10.5) set_mock_value(motor.low_limit_travel, -11) set_mock_value(motor.high_limit_travel, 21) @@ -279,6 +298,7 @@ async def test_prepare_motor_limits_error( upper_limit, lower_limit, ): + await motor.stage() set_mock_value(motor.acceleration_time, acceleration_time) set_mock_value(motor.low_limit_travel, lower_limit) set_mock_value(motor.high_limit_travel, upper_limit) @@ -292,28 +312,33 @@ async def test_prepare_motor_limits_error( await motor.prepare(fly_info) -async def test_prepare_valid_limits(motor: Motor): +@pytest.mark.parametrize( + "start_position, end_position, expected_setpoint, expected_end", + [(0, 10, -5, 15), (12, 2, 17, -3)], +) +async def test_prepare_valid_limits( + motor: Motor, + start_position: float, + end_position: float, + expected_setpoint: float, + expected_end: float, +): + await motor.stage() set_mock_value(motor.acceleration_time, 1) set_mock_value(motor.low_limit_travel, -10.01) set_mock_value(motor.high_limit_travel, 20.01) set_mock_value(motor.max_velocity, 10) fly_info = FlyMotorInfo( - start_position=0, - end_position=10, + start_position=start_position, + end_position=end_position, time_for_move=1, ) await motor.prepare(fly_info) - assert await motor.user_setpoint.get_value() == -5 - - assert fly_info.ramp_down_end_pos(await motor.acceleration_time.get_value()) == 15 - fly_info = FlyMotorInfo( - start_position=12, - end_position=2, - time_for_move=1, + assert await motor.user_setpoint.get_value() == expected_setpoint + assert ( + fly_info.ramp_down_end_pos(await motor.acceleration_time.get_value()) + == expected_end ) - await motor.prepare(fly_info) - assert await motor.user_setpoint.get_value() == 17 - assert fly_info.ramp_down_end_pos(await motor.acceleration_time.get_value()) == -3 @pytest.mark.parametrize( @@ -331,6 +356,7 @@ async def test_fly_scan( expected_end: float, ): """Happy-path prepare -> kickoff -> complete through the public interface.""" + await motor.stage() # A small acceleration time keeps the up-to-speed wait in kickoff() negligible set_mock_value(motor.acceleration_time, 0.2) set_mock_value(motor.low_limit_travel, -15) @@ -354,30 +380,53 @@ async def test_fly_scan( assert await motor.user_setpoint.get_value() == pytest.approx(expected_end) -async def test_kickoff_before_prepare_raises(motor: Motor): - with pytest.raises(RuntimeError, match="prepare.* before kickoff"): - await motor.kickoff() +async def test_stop_during_a_fly_halts_the_motor_but_does_not_end_complete( + velo_and_accl_respecting_motor: Motor, +): + motor = velo_and_accl_respecting_motor + await motor.stage() + await motor.prepare(FlyMotorInfo(start_position=0, end_position=2, time_for_move=1)) + await motor.kickoff() + status = motor.complete() + await asyncio.sleep(0.2) + await motor.stop() -async def test_complete_before_kickoff_raises(motor: Motor): - with pytest.raises(RuntimeError, match="kickoff.* before complete"): - await motor.complete() + # stop() halts the hardware, but only unstage() ends the fly + assert await motor.motor_done_move.get_value() == 1 + await motor.unstage() + with contextlib.suppress(asyncio.CancelledError): + await status -@pytest.mark.parametrize( - "done_move, expect_stop", - [ - (0, True), # moving -> stop is issued - (1, False), # already stopped -> no redundant STOP write - ], -) -async def test_stop_only_when_moving(motor: Motor, done_move: int, expect_stop: bool): +async def test_unstage_during_kickoff_run_up_cancels_the_move(motor: Motor): + await motor.stage() + get_mock_put(motor.motor_stop).reset_mock() + set_mock_value(motor.acceleration_time, 10) + set_mock_value(motor.max_velocity, 10) + await motor.prepare(FlyMotorInfo(start_position=0, end_position=1, time_for_move=1)) + # The fly move never finishes, and the motor is moving + set_mock_put_proceeds(motor.user_setpoint, False) + kickoff = motor.kickoff() + await wait_for_pending_wakeups() + set_mock_value(motor.motor_done_move, 0) + assert not kickoff.done + + await motor.unstage() + with pytest.raises(asyncio.CancelledError): + await kickoff + get_mock_put(motor.motor_stop).assert_called_once_with(1) + assert await motor.motor_done_move.get_value() == 1 + + +@pytest.mark.parametrize("done_move", [0, 1]) +async def test_stop_is_sent_whether_or_not_the_motor_looks_moving( + motor: Motor, done_move: int +): + # DMOV can flicker at standstill on some motors, so it does not decide set_mock_value(motor.motor_done_move, done_move) await motor.stop() - if expect_stop: - get_mock_put(motor.motor_stop).assert_called_once_with(1) - else: - get_mock_put(motor.motor_stop).assert_not_called() + get_mock_put(motor.motor_stop).assert_called_once_with(1) def test_core_notconnected_emits_deprecation_warning(): @@ -617,6 +666,8 @@ async def test_velo_and_accl_respecting_motor_mock_behavior( [ # RunEngine path: Motor.stop() puts to STOP, then cancels the set() status ("motor_stop", RuntimeError), + # RunEngine pause path: stop(success=True) lets the set() finish normally + ("motor_stop_success", None), # Someone else writes to STOP; the motor record finishes the move normally ("stop_signal", None), # set() timeout cancels the simulated move without any stop @@ -629,9 +680,9 @@ async def test_velo_and_accl_respecting_motor_mock_interrupted_move( motor = velo_and_accl_respecting_motor # velocity 10 over 10 units, so the full move takes about 1.1s status = motor.set(10.0, timeout=0.3 if interrupt == "timeout" else 5) - if interrupt == "motor_stop": + if interrupt in ("motor_stop", "motor_stop_success"): await wait_for_value(motor.user_readback, lambda v: v > 0, timeout=1) - await motor.stop() + await motor.stop(success=interrupt == "motor_stop_success") elif interrupt == "stop_signal": await wait_for_value(motor.user_readback, lambda v: v > 0, timeout=1) await motor.motor_stop.set(1) @@ -644,3 +695,80 @@ async def test_velo_and_accl_respecting_motor_mock_interrupted_move( # the target means the move ended early rather than finishing assert await motor.motor_done_move.get_value() == 1 assert await motor.user_readback.get_value() < 10.0 + + +@pytest.mark.parametrize("moving", [False, True], ids=["before_motion", "moving"]) +async def test_a_set_stopped_before_or_during_motion_fails_and_stays_still( + velo_and_accl_respecting_motor: Motor, moving: bool +): + motor = velo_and_accl_respecting_motor + await motor.stage() + status = motor.set(10.0) + if moving: + await wait_for_value(motor.user_readback, lambda v: v > 0, timeout=1) + + await motor.stop() + + # stop() does not wait for the set to finish failing + with pytest.raises(StoppedError, match=f"Device {motor.name} was stopped"): + await status + assert await motor.motor_done_move.get_value() == 1 + stopped_at = await motor.user_readback.get_value() + assert (0 < stopped_at < 10) if moving else stopped_at == 0 + await asyncio.sleep(0.2) + assert await motor.user_readback.get_value() == stopped_at + + +async def test_stage_stops_a_move_in_progress( + velo_and_accl_respecting_motor: Motor, +): + motor = velo_and_accl_respecting_motor + status = motor.set(0.5) + await wait_for_value(motor.user_readback, lambda v: v > 0, timeout=1) + + await motor.stage() + + get_mock_put(motor.motor_stop).assert_called_once_with(1) + stopped_at = await motor.user_readback.get_value() + assert 0 < stopped_at < 0.5 + await status + assert await motor.user_readback.get_value() == stopped_at + await motor.unstage() + + +# Stopping the motor during the run-up does not end prepare, which finishes short +async def test_stop_during_prepare_run_up_halts_the_motor_and_prepare_finishes_short( + velo_and_accl_respecting_motor: Motor, +): + motor = velo_and_accl_respecting_motor + await motor.stage() + preparing = motor.prepare( + FlyMotorInfo(start_position=100, end_position=200, time_for_move=1) + ) + await wait_for_value(motor.user_readback, lambda v: v > 1, timeout=1) + await motor.stop() + await asyncio.wait_for(preparing, timeout=1) + assert await motor.user_readback.get_value() < 100 + await motor.unstage() + + +async def test_default_mock_motor_stop_finishes_a_move_in_progress(): + async with init_devices(mock=True): + motor = Motor("BLxxI-MO-TABLE-01:X", name="motor") + # A test simulating a move that has not finished + set_mock_value(motor.motor_done_move, 0) + + await asyncio.wait_for(motor.stop(), timeout=1) + + get_mock_put(motor.motor_stop).assert_called_once_with(1) + assert await motor.motor_done_move.get_value() == 1 + + +async def test_motor_stop_allows_the_acceleration_time_to_decelerate(motor: Motor): + await motor.acceleration_time.set(30) + set_mock_value(motor.motor_done_move, 0) + + with patch("ophyd_async.epics.motor.wait_for_value") as wait_for_value: + await motor.stop() + + assert wait_for_value.call_args.kwargs["timeout"] == 30 + DEFAULT_TIMEOUT diff --git a/tests/unit_tests/fastcs/panda/test_hdf_panda.py b/tests/unit_tests/fastcs/panda/test_hdf_panda.py index 7b9e87b4e0..6b5e5599b1 100644 --- a/tests/unit_tests/fastcs/panda/test_hdf_panda.py +++ b/tests/unit_tests/fastcs/panda/test_hdf_panda.py @@ -270,7 +270,7 @@ async def test_flyscan_documents(hdf_panda: HDFPanda, tmp_path): }, ), ] - assert status.done + await status # Check that another kickoff without prepare is not allowed with pytest.raises( RuntimeError, @@ -313,7 +313,7 @@ async def test_flyscan_documents(hdf_panda: HDFPanda, tmp_path): }, ), ] - assert status.done + await status async def test_get_indices_written(hdf_panda: HDFPanda): diff --git a/tests/unit_tests/fastcs/panda/test_trigger.py b/tests/unit_tests/fastcs/panda/test_trigger.py index 345a967a30..31732cf01b 100644 --- a/tests/unit_tests/fastcs/panda/test_trigger.py +++ b/tests/unit_tests/fastcs/panda/test_trigger.py @@ -66,6 +66,7 @@ async def set_active(value: bool): set_mock_value(mock_panda.seq[1].active, value) flyer = trigger_logic.with_device() + await flyer.stage() await flyer.prepare(seq_table_info) await asyncio.gather(flyer.kickoff(), set_active(True)) await asyncio.gather(flyer.complete(), set_active(False)) @@ -105,7 +106,8 @@ async def test_seq_scanspec_trigger_logic(mock_panda, sim_x_motor, sim_y_motor) ), # type: ignore }, ) - await trigger_logic.on_prepare(info) + async for _ in trigger_logic.prepare_to_fly("", info, None): + pass out = await trigger_logic.seq.table.get_value() assert out.repeats == pytest.approx([1, 1, 1, 5, 1, 1, 1, 5, 1, 1, 1, 5]) assert out.trigger == [ @@ -146,7 +148,8 @@ async def test_seq_scanspec_trigger_logic_no_gaps( ) }, ) - await trigger_logic.on_prepare(info) + async for _ in trigger_logic.prepare_to_fly("", info, None): + pass out = await trigger_logic.seq.table.get_value() assert out.repeats == pytest.approx([1, 1, 1, 3]) assert out.trigger == [ @@ -181,7 +184,8 @@ async def test_seq_scanspec_trigger_logic_duration_error( }, ) with pytest.raises(RuntimeError, match="Slice must have duration"): - await trigger_logic.on_prepare(info) + async for _ in trigger_logic.prepare_to_fly("", info, None): + pass async def test_seq_scanspec_trigger_logic_motor_not_passed( @@ -199,7 +203,8 @@ async def test_seq_scanspec_trigger_logic_motor_not_passed( ) }, ) - await trigger_logic.on_prepare(info) + async for _ in trigger_logic.prepare_to_fly("", info, None): + pass out = await trigger_logic.seq.table.get_value() assert out.repeats == pytest.approx([1, 1, 3]) assert out.trigger == [ @@ -227,7 +232,8 @@ async def test_seq_scanspec_trigger_logic_equal( ) }, ) - await trigger_logic.on_prepare(info) + async for _ in trigger_logic.prepare_to_fly("", info, None): + pass out = await trigger_logic.seq.table.get_value() assert out.repeats == pytest.approx([1, 1, 1, 1, 1, 1, 1, 1, 1]) assert out.trigger == [ @@ -261,6 +267,7 @@ async def set_active(value: bool): set_mock_value(mock_panda.pcomp[1].active, value) flyer = trigger_logic.with_device() + await flyer.stage() await flyer.prepare(pcomp_info) await asyncio.gather(flyer.kickoff(), set_active(True)) await asyncio.gather(flyer.complete(), set_active(False)) diff --git a/tests/unit_tests/sim/test_sim_motor.py b/tests/unit_tests/sim/test_sim_motor.py index 6d2d9f061f..9a8baa780c 100644 --- a/tests/unit_tests/sim/test_sim_motor.py +++ b/tests/unit_tests/sim/test_sim_motor.py @@ -1,20 +1,24 @@ import asyncio import sys +import threading from unittest.mock import Mock, call, patch +import bluesky.plan_stubs as bps +import bluesky.preprocessors as bpp import pytest from bluesky.plans import spiral_square from bluesky.protocols import Reading from bluesky.run_engine import RunEngine +from bluesky.utils import FailedStatus, RunEngineInterrupted from ophyd_async.core import ( Device, FlyMotorInfo, + StoppedError, WatchableAsyncStatus, forward_watcher_updates, ) from ophyd_async.sim import SimMotor -from ophyd_async.testing import StatusWatcher class TwoMotorStage(Device): @@ -96,19 +100,86 @@ async def test_short_move_is_exactly_move_time(m2: SimMotor): @pytest.mark.timeout(3) -async def test_stop(m2: SimMotor): +@pytest.mark.parametrize("moving", [False, True], ids=["before_motion", "moving"]) +@pytest.mark.parametrize("success", [False, True]) +async def test_stop(m2: SimMotor, success: bool, moving: bool): # this move should take 10 seconds but we will stop it partway through. # The sim motor updates its readback at 10Hz, so 0.2s is two update ticks # in - enough for the readback to have moved off 0 without waiting longer. move_status = m2.set(10) - await asyncio.sleep(0.2) - await m2.stop(success=False) + if moving: + await asyncio.sleep(0.2) + await m2.stop(success=success) new_pos = await m2.user_readback.get_value() - assert 0 < new_pos < 10 + assert (0 < new_pos < 10) if moving else new_pos == 0 - assert not move_status.success - with pytest.raises(RuntimeError, match=f"Device {m2.name} was stopped"): + if success: await move_status + assert move_status.success + else: + assert not move_status.success + with pytest.raises(StoppedError, match=f"Device {m2.name} was stopped"): + await move_status + # Nothing is left moving + await asyncio.sleep(0.3) + assert await m2.user_readback.get_value() == new_pos + + +@pytest.mark.timeout(3) +async def test_second_set_supersedes_the_first(m2: SimMotor): + first = m2.set(10) + await asyncio.sleep(0.2) + second = m2.set(0.5) + + with pytest.raises(StoppedError, match=r"M2: superseded by a new set\(\)"): + await first + await second + # Only the second move is left moving the motor, and it has finished + assert await m2.user_readback.get_value() == 0.5 + await asyncio.sleep(0.3) + assert await m2.user_readback.get_value() == 0.5 + + +@pytest.mark.timeout(3) +async def test_unstage_stops_a_fly_and_waits_for_it(m2: SimMotor): + await asyncio.gather(m2.velocity.set(10), m2.acceleration_time.set(0.05)) + await m2.stage() + await m2.prepare(FlyMotorInfo(start_position=0, end_position=5, time_for_move=1)) + await m2.kickoff() + status = m2.complete() + await asyncio.sleep(0.2) + + await m2.unstage() + + # Already stopped when unstage returns, not stopping soon after + with pytest.raises(asyncio.CancelledError): + await status + stopped_at = await m2.user_readback.get_value() + assert 0 < stopped_at < 5 + assert await m2.user_setpoint.get_value() == stopped_at + await asyncio.sleep(0.3) + assert await m2.user_readback.get_value() == stopped_at + + +@pytest.mark.timeout(5) +@pytest.mark.parametrize("success", [False, True]) +async def test_stop_during_a_fly_halts_the_motor_and_cancels_complete( + m2: SimMotor, success: bool +): + await asyncio.gather(m2.velocity.set(10), m2.acceleration_time.set(0.05)) + await m2.stage() + await m2.prepare(FlyMotorInfo(start_position=0, end_position=2, time_for_move=1)) + await m2.kickoff() + status = m2.complete() + await asyncio.sleep(0.2) + + await m2.stop(success=success) + + # The sim motion runs in the fly, so stopping it cancels the fly too + with pytest.raises(asyncio.CancelledError): + await status + assert 0 < await m2.user_readback.get_value() < 2 + await m2.unstage() @pytest.mark.skipif("win" in sys.platform, reason="windows CI runners too weedy") @@ -116,26 +187,43 @@ async def test_fly(m1: SimMotor): await m1.acceleration_time.set(0.1) info = FlyMotorInfo(start_position=0, end_position=1, time_for_move=0.2) fly_start, fly_end, velocity = -0.25, 1.25, 5 + await m1.stage() await m1.prepare(info) assert await m1.user_readback.get_value() == fly_start assert await m1.velocity.get_value() == velocity await m1.kickoff() status = m1.complete() - watcher = StatusWatcher(status) - for i, v in enumerate([-0.25, 0, 0.5, 1.0, 1.25]): - await watcher.wait_for_call( - current=pytest.approx(v), - initial=fly_start, - name="M1", - target=fly_end, - time_elapsed=pytest.approx(i * 0.1, abs=0.1), - unit="mm", - ) + watcher = Mock() + status.watch(watcher) await status - watcher.mock.assert_not_called() + # Progress is reported from where the motor is when complete is called + # (it has been running up for the acceleration time) to the end + updates = [c.kwargs for c in watcher.call_args_list] + currents = [u["current"] for u in updates] + assert currents == sorted(currents) + assert currents[-1] == pytest.approx(fly_end) + assert {(u["initial"], u["target"], u["name"], u["unit"]) for u in updates} == { + (fly_start, fly_end, "M1", "mm") + } assert await m1.user_readback.get_value() == fly_end +@pytest.mark.timeout(3) +async def test_unstage_during_kickoff_run_up_leaves_no_move_running(m2: SimMotor): + await m2.acceleration_time.set(0.5) + await m2.stage() + await m2.prepare(FlyMotorInfo(start_position=0, end_position=5, time_for_move=1)) + kickoff_status = m2.kickoff() + await asyncio.sleep(0.2) + await m2.unstage() + with pytest.raises(asyncio.CancelledError): + await kickoff_status + # The readback updates at 10Hz, so a move still running would change it + stopped_at = await m2.user_readback.get_value() + await asyncio.sleep(0.3) + assert await m2.user_readback.get_value() == stopped_at + + async def test_sim_motor_can_be_set_to_its_current_position(m1: SimMotor): await m1.set(0) @@ -221,3 +309,79 @@ def on_readback(value: dict[str, Reading[float]]) -> None: previous <= current for previous, current in zip(readbacks, readbacks[1:], strict=False) ) + + +FLY_INFO = FlyMotorInfo(start_position=0, end_position=1, time_for_move=0.3) +FLY_END = FLY_INFO.ramp_down_end_pos(0.05) + + +async def make_fast_motor() -> SimMotor: + motor = SimMotor("M", instant=False) + await asyncio.gather(motor.velocity.set(20), motor.acceleration_time.set(0.05)) + return motor + + +def move_and_fly_plan( + motor: SimMotor, move_to: float, hang: str | None = None, fail: bool = False +): + """Move then fly the motor in a run, hanging in `hang` ("move" or "fly") if given. + + The hang ends in a failure if `fail`, otherwise it waits to be interrupted. + """ + + def hang_then_fail(): + yield from bps.sleep(0.3) + if fail: + raise ValueError("plan failed") + yield from bps.sleep(5) + + @bpp.run_decorator() + @bpp.stage_decorator([motor]) + def plan(): + if hang == "move": + yield from bps.abs_set(motor, 1000, group="move") + yield from hang_then_fail() + yield from bps.mv(motor, move_to) + yield from bps.prepare(motor, FLY_INFO, wait=True) + yield from bps.kickoff(motor, wait=True) + if hang == "fly": + yield from bps.complete(motor, group="fly") + yield from hang_then_fail() + yield from bps.complete(motor, wait=True) + + return plan() + + +async def test_staged_plans_back_to_back_move_and_fly_the_same_motor(): + RE = RunEngine() + motor = await make_fast_motor() + for move_to in (3, -2): + RE(move_and_fly_plan(motor, move_to)) + assert await motor.user_readback.get_value() == pytest.approx(FLY_END) + assert await motor.user_setpoint.get_value() == pytest.approx(FLY_END) + + +@pytest.mark.timeout(20) +@pytest.mark.parametrize("hang", ["move", "fly"]) +@pytest.mark.parametrize("interrupt", ["fail", "pause_then_abort"]) +async def test_a_plan_interrupted_mid_move_or_fly_leaves_the_motor_ready_for_the_next( + hang: str, interrupt: str +): + RE = RunEngine() + motor = await make_fast_motor() + if interrupt == "fail": + # unstage ends the move or fly, which bluesky reports as a FailedStatus + with pytest.raises((ValueError, FailedStatus)): + RE(move_and_fly_plan(motor, 3, hang=hang, fail=True)) + else: + threading.Timer(0.5, RE.request_pause).start() + with pytest.raises(RunEngineInterrupted): + RE(move_and_fly_plan(motor, 3, hang=hang)) + RE.abort() + # Nothing is left moving the motor + stopped_at = await motor.user_readback.get_value() + await asyncio.sleep(0.3) + assert await motor.user_readback.get_value() == stopped_at + + RE(move_and_fly_plan(motor, 3)) + assert await motor.user_readback.get_value() == pytest.approx(FLY_END) diff --git a/uv.lock b/uv.lock index 5fdfad44f5..e36da88e3b 100644 --- a/uv.lock +++ b/uv.lock @@ -1580,6 +1580,7 @@ dependencies = [ { name = "pyyaml" }, { name = "scanspec" }, { name = "stamina" }, + { name = "typing-extensions" }, { name = "velocity-profile" }, ] @@ -1654,6 +1655,7 @@ requires-dist = [ { name = "pyyaml" }, { name = "scanspec", specifier = ">=1.0.0" }, { name = "stamina", specifier = ">=23.1.0" }, + { name = "typing-extensions", specifier = ">=4.14" }, { name = "velocity-profile" }, ] provides-extras = ["sim", "ca", "pva", "tango", "demo"] @@ -1675,7 +1677,7 @@ dev = [ { name = "pre-commit" }, { name = "pydata-sphinx-theme", specifier = ">=0.12" }, { name = "pyepics", specifier = ">=3.4.2" }, - { name = "pyright" }, + { name = "pyright", specifier = ">=1.1.410" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-cov" }, @@ -2267,15 +2269,15 @@ wheels = [ [[package]] name = "pyright" -version = "1.1.409" +version = "1.1.414" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nodeenv" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/51/4e/3aa27f74211522dba7e9cbc3e74de779c6d4b654c54e50a4840623be8014/pyright-1.1.409.tar.gz", hash = "sha256:986ee05beca9e077c165758ad123667c679e050059a2546aa02473930394bc93", size = 4430434, upload-time = "2026-04-23T11:02:03.799Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/1b/244c7b710031ada80f27e579ec20d28a2285dfc318fed0339866b1047f12/pyright-1.1.414.tar.gz", hash = "sha256:523c0a97c60da6333234955c277730c9cf4f5bd6d5399e7b7d2b0fc5d3599524", size = 4154638, upload-time = "2026-09-10T12:26:53.181Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/16/6b/330d8ebae582b30c2959a1ef4c3bc344ebde48c2ff0c3f113c4710735e11/pyright-1.1.409-py3-none-any.whl", hash = "sha256:aa3ea228cab90c845c7a60d28db7a844c04315356392aa09fafcee98c8c22fb3", size = 6438161, upload-time = "2026-04-23T11:02:01.309Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ba/18b6e682ead424ad24bcc134339ae5d1b931cd9ae260540592a058a91279/pyright-1.1.414-py3-none-any.whl", hash = "sha256:2a6b4b3298c9eec174c5ed83bd338de6eee82df2992f3e1930e6199d381be36f", size = 6225049, upload-time = "2026-09-10T12:26:51.427Z" }, ] [[package]] @@ -2816,23 +2818,23 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version < '3.12'" }, - { name = "babel", marker = "python_full_version < '3.12'" }, - { name = "colorama", marker = "python_full_version < '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version < '3.12'" }, - { name = "imagesize", marker = "python_full_version < '3.12'" }, - { name = "jinja2", marker = "python_full_version < '3.12'" }, - { name = "packaging", marker = "python_full_version < '3.12'" }, - { name = "pygments", marker = "python_full_version < '3.12'" }, - { name = "requests", marker = "python_full_version < '3.12'" }, - { name = "roman-numerals", marker = "python_full_version < '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -2847,23 +2849,23 @@ resolution-markers = [ "python_full_version >= '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version >= '3.12'" }, - { name = "babel", marker = "python_full_version >= '3.12'" }, - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version >= '3.12'" }, - { name = "imagesize", marker = "python_full_version >= '3.12'" }, - { name = "jinja2", marker = "python_full_version >= '3.12'" }, - { name = "packaging", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "requests", marker = "python_full_version >= '3.12'" }, - { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [