backup-tool/tests/fault/test_worker_loss.py

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()