sched: require proven ownership at poll/admin/recovery entries (F5)

tests/scheduling/test_ownership_guard.py

@@ -0,0 +1,191 @@
+"""Ownership guards at scheduler mutation/recovery/dispatch entries (R4/F5).
+
+The held lock must be proven at the entry points, not merely exist as a
+class: ``poll``, schedule administration, and ``recover`` reject when the
+caller cannot prove live ownership, before any write or side effect.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+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, SecondOwnerError
+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 _scheduler(
+    tmp_path: Path,
+    ownership: SchedulerOwnership,
+    *,
+    script: dict | None = None,
+) -> tuple[Scheduler, FileScheduleStore, FileRunStore]:
+    sched_store = FileScheduleStore(tmp_path / "sched")
+    run_store = FileRunStore(tmp_path / "runs")
+    sched = Scheduler(
+        schedule_store=sched_store,
+        run_store=run_store,
+        sources={},
+        capacity=4,
+        preparer=SchedulePreparer(
+            DictDeployments({"dep-1": {"rev": 1, "required": []}}),
+            fixture_environment,
+        ),
+        dispatcher=ScriptedDispatcher(script),
+        ownership=ownership,
+    )
+    return sched, sched_store, run_store
+
+
+def _due_setup(sched: Scheduler, store: FileScheduleStore, intended: datetime) -> None:
+    store.create_schedule(_sched_model("a"))
+    store.save_consumed("a", intended - timedelta(hours=1))
+    sched.sources["a"] = OneShotSource(intended)
+
+
+def test_poll_without_ownership_rejects_before_writes(tmp_path: Path) -> None:
+    calls: list[str] = []
+
+    def spy(admission: Any, now: datetime) -> Any:
+        calls.append(admission.id)
+        raise AssertionError("dispatcher must not run without ownership")
+
+    ownership = SchedulerOwnership(tmp_path / "sched", owner="never-acquired")
+    sched, store, runs = _scheduler(tmp_path, ownership)
+    sched.dispatcher = ScriptedDispatcher({"*": spy})
+    intended = ts(2026, 9, 8, 12, 0)
+    _due_setup(sched, store, intended)
+    with pytest.raises(SecondOwnerError):
+        sched.poll(intended)
+    assert calls == []
+    assert runs.list_runs() == []
+    assert runs.list_admissions() == []
+    assert store.list_occurrences("a", limit=100)["total"] == 0
+    assert store.get_consumed("a") == intended - timedelta(hours=1)
+
+
+def test_poll_with_released_ownership_rejects(tmp_path: Path) -> None:
+    ownership = SchedulerOwnership(tmp_path / "sched", owner="test").acquire()
+    sched, store, runs = _scheduler(tmp_path, ownership)
+    intended = ts(2026, 9, 8, 12, 0)
+    _due_setup(sched, store, intended)
+    ownership.release()
+    assert not ownership.held
+    with pytest.raises(SecondOwnerError):
+        sched.poll(intended)
+    assert runs.list_runs() == []
+    assert store.list_occurrences("a", limit=100)["total"] == 0
+
+
+def test_recover_without_ownership_rejects_before_writes(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 / "sched", owner="never-acquired")
+    with pytest.raises(SecondOwnerError):
+        sched_recovery.recover(
+            schedule_store=sched_store,
+            run_store=run_store,
+            now=ts(2026, 9, 8, 12, 0),
+            ownership=ownership,
+        )
+    # Rejected before any write: the admitted run is untouched.
+    assert run_store.get_run(admission.id).status.value == "admitted"
+
+
+def test_held_ownership_allows_poll_and_recovery(tmp_path: Path) -> None:
+    ownership = SchedulerOwnership(tmp_path / "sched", owner="test").acquire()
+    try:
+        sched, store, runs = _scheduler(tmp_path, ownership, script={"*": "hang"})
+        intended = ts(2026, 9, 8, 12, 0)
+        _due_setup(sched, store, intended)
+        assert sched.poll(intended) == {"a": f"admit:{runs.list_runs()[0].id}"}
+        diags = sched_recovery.recover(
+            schedule_store=store,
+            run_store=runs,
+            now=intended,
+            ownership=ownership,
+        )
+        assert diags != []
+    finally:
+        ownership.release()
+
+
+def test_second_owner_cannot_acquire_for_poll(tmp_path: Path) -> None:
+    first = SchedulerOwnership(tmp_path / "sched", owner="first").acquire()
+    try:
+        second = SchedulerOwnership(tmp_path / "sched", owner="second")
+        with pytest.raises(SecondOwnerError):
+            second.acquire()
+        assert not second.held
+        # A scheduler proving the loser's (unheld) ownership is rejected.
+        sched, store, _ = _scheduler(tmp_path, second)
+        intended = ts(2026, 9, 8, 12, 0)
+        _due_setup(sched, store, intended)
+        with pytest.raises(SecondOwnerError):
+            sched.poll(intended)
+    finally:
+        first.release()
+
+
+def test_admin_mutations_require_ownership(tmp_path: Path) -> None:
+    ownership = SchedulerOwnership(tmp_path / "sched", owner="test")
+    sched, store, _ = _scheduler(tmp_path, ownership)
+    store.create_schedule(_sched_model("a", paused=True))
+    with pytest.raises(SecondOwnerError):
+        sched.resume_schedule("a", ts(2026, 9, 8, 12, 0))
+    with pytest.raises(SecondOwnerError):
+        sched.edit_schedule("a", ts(2026, 9, 8, 12, 0))
+    assert store.get_schedule("a").paused is True
+    assert store.get_schedule("a").revision == 1