112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
"""Controlled scheduler collaborators for tests (fixture side of F6 seams).
|
|
|
|
``DictDeployments``, ``fixture_environment``, and ``ScriptedDispatcher``
|
|
are the test doubles behind the production
|
|
:mod:`wf_scheduling.prepare` / :mod:`wf_scheduling.dispatch` protocols.
|
|
They carry the fixture behavior the old production ``Scheduler`` used to
|
|
hardcode (fixture input, fixture environments, scripted outcomes); the
|
|
production package no longer contains any of it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Sequence
|
|
from datetime import datetime
|
|
from typing import Any, Union
|
|
|
|
from tests.artifacts.test_run_store import artifact as _artifact
|
|
from tests.artifacts.test_run_store import deployment as _deployment
|
|
from wf_artifacts.runs.models import PinnedRunEnvironment, RunAdmission
|
|
from wf_core import RunState, RunStatus
|
|
from wf_scheduling.dispatch import DispatchResult, StillRunning, Stopped
|
|
|
|
OutcomeSpec = Union[str, Callable[[RunAdmission, datetime], DispatchResult]]
|
|
|
|
|
|
class DictDeployments:
|
|
"""Deployment directory over a plain ``{id: {"rev": n, "required": [...]}}`` map."""
|
|
|
|
def __init__(self, deployments: dict[str, dict[str, Any]]) -> None:
|
|
self._deployments = deployments
|
|
|
|
def deployment_revision(self, deployment_id: str) -> int:
|
|
return int(self._deployments[deployment_id]["rev"])
|
|
|
|
def required_inputs(self, deployment_id: str) -> Sequence[str]:
|
|
return list(self._deployments[deployment_id].get("required", []))
|
|
|
|
|
|
def fixture_environment(sched: Any) -> PinnedRunEnvironment:
|
|
"""Build the historic fixture environment for schedule tests."""
|
|
deployment = _deployment().model_copy(
|
|
update={"id": getattr(sched, "deployment_id", "dep-1")}
|
|
)
|
|
return PinnedRunEnvironment(
|
|
deployment=deployment, root_artifact=_artifact(), child_artifacts=[]
|
|
)
|
|
|
|
|
|
def stopped_state(admission: RunAdmission, outcome: str) -> RunState:
|
|
"""Build a genuine stopped RunState echoing the admitted input."""
|
|
status = {
|
|
"complete": RunStatus.COMPLETED,
|
|
"interrupt": RunStatus.INTERRUPTED,
|
|
"fail": RunStatus.FAILED,
|
|
}[outcome]
|
|
return RunState(
|
|
workflow_name="sched",
|
|
status=status,
|
|
workflow_input=dict(admission.resolved_input),
|
|
state={},
|
|
)
|
|
|
|
|
|
class ScriptedDispatcher:
|
|
"""Scripted execution double: ``{run_id | "*": outcome | callable}``.
|
|
|
|
Outcomes are ``complete`` | ``interrupt`` | ``fail`` | ``hang``.
|
|
Callables receive ``(admission, now)`` and may raise or terminate the
|
|
process (crash tests). ``finish`` builds a stopped state for a hanging
|
|
run so tests can settle it through the scheduler's async-completion
|
|
seam without re-dispatching.
|
|
"""
|
|
|
|
def __init__(self, script: dict[str, OutcomeSpec] | None = None) -> None:
|
|
self.script: dict[str, OutcomeSpec] = dict(script or {})
|
|
|
|
def dispatch(self, *, admission: RunAdmission, now: datetime) -> DispatchResult:
|
|
spec = self.script.get(admission.id, self.script.get("*", "complete"))
|
|
if callable(spec):
|
|
return spec(admission, now)
|
|
assert isinstance(spec, str), f"unknown scripted outcome {spec!r}"
|
|
if spec == "hang":
|
|
return StillRunning()
|
|
if spec in ("complete", "interrupt", "fail"):
|
|
return Stopped(result=stopped_state(admission, spec))
|
|
raise ValueError(f"unknown scripted outcome {spec!r} for {admission.id!r}")
|
|
|
|
def finish(self, admission: RunAdmission, outcome: str) -> RunState:
|
|
"""Build the stopped state that settles a hanging run."""
|
|
return stopped_state(admission, outcome)
|
|
|
|
|
|
class WorkflowDispatcher:
|
|
"""Execution double driving a real workflow to its first stopped state.
|
|
|
|
Dispatch runs the workflow through the genuine ``wf_core`` runtime and
|
|
returns the resulting stopped state, so scheduler tests exercise real
|
|
execution (including durable interruptions) without enabling any server.
|
|
"""
|
|
|
|
def __init__(self, workflow: Any, registry: dict[str, Any] | None = None) -> None:
|
|
self.workflow = workflow
|
|
self.registry = registry or {}
|
|
|
|
def dispatch(self, *, admission: RunAdmission, now: datetime) -> DispatchResult:
|
|
from wf_core import execute_workflow
|
|
|
|
state = execute_workflow(
|
|
self.workflow, dict(admission.resolved_input), self.registry
|
|
)
|
|
return Stopped(result=state)
|