feat(v2): complete v2 reimplementation

tests/fault/test_worker_loss.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import asyncio
+import importlib
+from collections.abc import AsyncIterator
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+
+import pytest
+import pytest_asyncio
+from alembic import command
+from backup_tool.config import Settings
+from backup_tool.db.engine import create_engine
+from backup_tool.db.models import Execution, ExecutionEvent, Job, Repository, Source
+from backup_tool.execution import (
+    claim,
+    complete_cancellation,
+    heartbeat,
+    record_event,
+    request_cancellation,
+)
+from backup_tool.worker import Worker
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+
+cli = importlib.import_module("backup_tool.cli")
+
+
+@pytest_asyncio.fixture
+async def database(
+    tmp_path: Path,
+) -> AsyncIterator[tuple[Settings, async_sessionmaker[AsyncSession], AsyncEngine]]:
+    key = tmp_path / "master.key"
+    key.write_bytes(b"m5-fault-test-master-key-material-32-bytes-minimum")
+    key.chmod(0o600)
+    data_dir = tmp_path / "data"
+    repositories = tmp_path / "repositories"
+    sources = tmp_path / "sources"
+    restores = tmp_path / "restores"
+    for directory in (data_dir, repositories, sources, restores):
+        directory.mkdir()
+    settings = Settings(
+        data_dir=data_dir,
+        database_url=f"sqlite+aiosqlite:///{data_dir / 'metadata.db'}",
+        repository_roots=(repositories,),
+        local_source_roots=(sources,),
+        restore_roots=(restores,),
+        master_key_file=key,
+    )
+    command.upgrade(cli.build_alembic_config(settings), "head")
+    engine = create_engine(settings)
+    yield settings, async_sessionmaker(engine, expire_on_commit=False), engine
+    await engine.dispose()
+
+
+async def create_stale_execution(db: AsyncSession, suffix: str, state: str) -> Execution:
+    repository = Repository(
+        name=f"repository-{suffix}",
+        root=f"/repositories/{suffix}",
+        format_version=1,
+        compression="none",
+        encryption="none",
+    )
+    source = Source(
+        name=f"source-{suffix}",
+        kind="local",
+        public_config={"root": f"/sources/{suffix}"},
+        secret_refs=[],
+    )
+    db.add_all([repository, source])
+    await db.flush()
+    job = Job(
+        name=f"job-{suffix}",
+        source_id=source.id,
+        repository_id=repository.id,
+        requested_mode="full",
+        exclusions=[],
+        retention={},
+    )
+    db.add(job)
+    await db.flush()
+    execution = Execution(
+        job_id=job.id,
+        trigger="manual",
+        state=state,
+        lease_owner="lost-worker",
+        lease_expires_at=datetime.now(UTC) - timedelta(seconds=1),
+        heartbeat_at=datetime.now(UTC) - timedelta(seconds=2),
+        progress={},
+        reason_code="cancellation_requested" if state == "cancelling" else None,
+    )
+    db.add(execution)
+    await db.flush()
+    await record_event(db, execution)
+    await db.commit()
+    return execution
+
+
+@pytest.mark.asyncio
+async def test_startup_recovers_each_stale_lease_once_and_fences_lost_owner(
+    database: tuple[Settings, async_sessionmaker[AsyncSession], AsyncEngine],
+) -> None:
+    settings, sessions, _ = database
+    async with sessions() as db:
+        stale = {
+            state: await create_stale_execution(db, state, state)
+            for state in ("preparing", "running", "verifying", "cancelling")
+        }
+
+    worker = Worker(settings, owner="recovery-worker")
+    try:
+        assert await worker.startup() == len(stale)
+    finally:
+        await worker.engine.dispose()
+
+    async with sessions() as db:
+        recovered = {
+            state: await db.get(Execution, execution.id) for state, execution in stale.items()
+        }
+        for state in ("preparing", "running", "verifying"):
+            execution = recovered[state]
+            assert execution is not None
+            assert execution.state == "queued"
+            assert execution.reason_code == "worker_lost"
+            assert execution.lease_owner is None
+            assert execution.lease_expires_at is None
+            assert execution.heartbeat_at is None
+        cancelling = recovered["cancelling"]
+        assert cancelling is not None
+        assert cancelling.state == "cancelled"
+        assert cancelling.completed_at is not None
+        assert cancelling.reason_code == "cancellation_requested"
+        event_counts = {
+            execution.id: await db.scalar(
+                select(func.count()).where(ExecutionEvent.execution_id == execution.id)
+            )
+            for execution in recovered.values()
+            if execution is not None
+        }
+
+    assert event_counts == {execution.id: 2 for execution in stale.values()}
+
+    async with sessions() as db:
+        execution = recovered["running"]
+        assert execution is not None
+        assert await claim(db, execution.id, "replacement-worker") is not None
+        assert not await heartbeat(db, execution.id, "lost-worker")
+        assert await request_cancellation(db, execution.id) is not None
+        assert not await complete_cancellation(db, execution.id, "lost-worker")
+        assert await complete_cancellation(db, execution.id, "replacement-worker")
+
+
+@pytest.mark.asyncio
+async def test_stopping_worker_does_not_claim_new_work(
+    database: tuple[Settings, async_sessionmaker[AsyncSession], AsyncEngine],
+) -> None:
+    settings, _, _ = database
+    worker = Worker(settings, owner="stopping-worker")
+    try:
+        worker.stop()
+        assert not await worker.run_once()
+    finally:
+        await worker.engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_idle_worker_stops_promptly_and_disposes_its_engine(
+    database: tuple[Settings, async_sessionmaker[AsyncSession], AsyncEngine],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    settings, _, _ = database
+    worker = Worker(settings, owner="idle-worker")
+    disposed = asyncio.Event()
+    dispose = AsyncEngine.dispose
+
+    async def track_dispose(engine: AsyncEngine, *, close: bool = True) -> None:
+        disposed.set()
+        await dispose(engine, close=close)
+
+    monkeypatch.setattr(AsyncEngine, "dispose", track_dispose)
+    task = asyncio.create_task(worker.run())
+    await asyncio.sleep(0)
+    worker.stop()
+    await asyncio.wait_for(task, timeout=1)
+    assert disposed.is_set()