lda-wf/tests/scheduling/test_recovery.py

"""Startup recovery and reconciliation with real stores (T10)."""

from __future__ import annotations

from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from tests.scheduling.controlled import (
    DictDeployments,
    ScriptedDispatcher,
    fixture_environment,
)
from wf_artifacts.runs.store import FileRunStore
from wf_scheduling import recovery as sched_recovery
from wf_scheduling.calendar import OneShotSource
from wf_scheduling.models import Schedule
from wf_scheduling.ownership import SchedulerOwnership
from wf_scheduling.poll import Scheduler
from wf_scheduling.prepare import SchedulePreparer
from wf_scheduling.store import FileScheduleStore


def ts(y: int, mo: int, d: int, h: int = 0, mi: int = 0) -> datetime:
    return datetime(y, mo, d, h, mi, tzinfo=UTC)


def _sched_model(sid: str, **kw: Any) -> Schedule:
    now = ts(2026, 9, 8, 12, 0)
    base: dict[str, Any] = {
        "id": sid,
        "deployment_id": "dep-1",
        "trigger": {"kind": "cron", "expression": "0 * * * *", "timezone": "UTC"},
        "input_bindings": [],
        "created_at": now.isoformat(),
        "updated_at": now.isoformat(),
    }
    base.update(kw)
    return Schedule.model_validate(base)


def test_recovery_materializes_missing_view_as_pending(tmp_path: Path) -> None:
    from tests.artifacts.test_run_store import artifact as _artifact
    from tests.artifacts.test_run_store import deployment as _deployment
    from wf_api.run_lifecycle import persist_admission
    from wf_artifacts import PinnedRunEnvironment

    sched_store = FileScheduleStore(tmp_path / "sched")
    run_store = FileRunStore(tmp_path / "runs")
    sched_store.create_schedule(_sched_model("a"))
    env = PinnedRunEnvironment(
        deployment=_deployment(), root_artifact=_artifact(), child_artifacts=[]
    )
    admission = persist_admission(
        store=run_store,
        run_id=run_store.allocate_run_id(),
        environment=env,
        resolved_input={},
        max_steps=None,
    )
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        diags = sched_recovery.recover(
            schedule_store=sched_store,
            run_store=run_store,
            now=ts(2026, 9, 8, 12, 0),
            ownership=ownership,
        )
    finally:
        ownership.release()
    assert any("pending-dispatch" in d for d in diags)
    assert run_store.get_run(admission.id).status.value == "admitted"
    assert sched_recovery._is_pending(run_store, admission.id)


def test_recovery_fails_abandoned_admitted_without_replay(tmp_path: Path) -> None:
    from tests.artifacts.test_run_store import artifact as _artifact
    from tests.artifacts.test_run_store import deployment as _deployment
    from wf_api.run_lifecycle import materialize_admitted_view, persist_admission
    from wf_artifacts import PinnedRunEnvironment

    sched_store = FileScheduleStore(tmp_path / "sched")
    run_store = FileRunStore(tmp_path / "runs")
    sched_store.create_schedule(_sched_model("a"))
    env = PinnedRunEnvironment(
        deployment=_deployment(), root_artifact=_artifact(), child_artifacts=[]
    )
    admission = persist_admission(
        store=run_store,
        run_id=run_store.allocate_run_id(),
        environment=env,
        resolved_input={},
        max_steps=None,
        scheduled_at=ts(2026, 9, 8, 12, 0),
        schedule_id="a",
        schedule_revision=1,
    )
    materialize_admitted_view(store=run_store, admission=admission)
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        diags = sched_recovery.recover(
            schedule_store=sched_store,
            run_store=run_store,
            now=ts(2026, 9, 8, 12, 0),
            ownership=ownership,
        )
    finally:
        ownership.release()
    assert any("failed-closed" in d for d in diags)
    assert run_store.get_run(admission.id).status.value == "failed"


def test_recovery_never_executes_pending_until_poll(tmp_path: Path) -> None:

    from tests.artifacts.test_run_store import artifact as _artifact
    from tests.artifacts.test_run_store import deployment as _deployment
    from wf_api.run_lifecycle import persist_admission
    from wf_artifacts import PinnedRunEnvironment

    sched_store = FileScheduleStore(tmp_path / "sched")
    run_store = FileRunStore(tmp_path / "runs")
    sched_store.create_schedule(_sched_model("a"))
    env = PinnedRunEnvironment(
        deployment=_deployment(), root_artifact=_artifact(), child_artifacts=[]
    )
    admission = persist_admission(
        store=run_store,
        run_id=run_store.allocate_run_id(),
        environment=env,
        resolved_input={},
        max_steps=None,
        scheduled_at=ts(2026, 9, 8, 12, 0),
        schedule_id="a",
        schedule_revision=1,
    )
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        diags = sched_recovery.recover(
            schedule_store=sched_store,
            run_store=run_store,
            now=ts(2026, 9, 8, 12, 0),
            ownership=ownership,
        )
    finally:
        ownership.release()
    assert any("pending-dispatch" in d for d in diags)
    # Recovery itself produced no terminal history; the poll sweep dispatches.
    assert sched_store.list_occurrences("a", limit=100)["total"] == 0
    sources = {"a": OneShotSource(ts(2026, 9, 8, 12, 0))}
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        sched = Scheduler(
            schedule_store=sched_store,
            run_store=run_store,
            sources=sources,  # type: ignore[arg-type]
            capacity=4,
            preparer=SchedulePreparer(
                DictDeployments({"dep-1": {"rev": 1, "required": []}}),
                fixture_environment,
            ),
            dispatcher=ScriptedDispatcher({"*": "complete"}),
            ownership=ownership,
        )
        sched_store.save_consumed("a", ts(2026, 9, 8, 12, 0))
        sched.poll(ts(2026, 9, 8, 12, 1))
    finally:
        ownership.release()
    assert run_store.get_run(admission.id).status.value == "completed"
    assert not sched_recovery._is_pending(run_store, admission.id)