196 lines
6.6 KiB
Python
196 lines
6.6 KiB
Python
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, Job, Repository, Source
|
|
from backup_tool.execution import (
|
|
EnqueueError,
|
|
claim,
|
|
complete_cancellation,
|
|
enqueue,
|
|
heartbeat,
|
|
request_cancellation,
|
|
retry,
|
|
)
|
|
from sqlalchemy import 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[async_sessionmaker[AsyncSession], AsyncEngine]]:
|
|
key = tmp_path / "master.key"
|
|
key.write_bytes(b"m5-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 async_sessionmaker(engine, expire_on_commit=False), engine
|
|
await engine.dispose()
|
|
|
|
|
|
async def create_job(
|
|
db: AsyncSession, suffix: str, *, enabled: bool = True, state: str = "active"
|
|
) -> str:
|
|
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={},
|
|
enabled=enabled,
|
|
state=state,
|
|
)
|
|
db.add(job)
|
|
await db.commit()
|
|
return job.id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_enqueue_allows_exactly_one_active_execution(
|
|
database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
|
|
) -> None:
|
|
sessions, _ = database
|
|
async with sessions() as db:
|
|
job_id = await create_job(db, "concurrent")
|
|
|
|
async def start() -> Execution | EnqueueError:
|
|
async with sessions() as db:
|
|
try:
|
|
return await enqueue(db, job_id)
|
|
except EnqueueError as error:
|
|
return error
|
|
|
|
first, second = await asyncio.gather(start(), start())
|
|
results = [first, second]
|
|
successes = [result for result in results if isinstance(result, Execution)]
|
|
failures = [result for result in results if isinstance(result, EnqueueError)]
|
|
|
|
assert len(successes) == 1
|
|
assert len(failures) == 1
|
|
assert failures[0].code == "execution_active"
|
|
assert failures[0].active_execution_id == successes[0].id
|
|
async with sessions() as db:
|
|
executions = list(await db.scalars(select(Execution).where(Execution.job_id == job_id)))
|
|
assert [execution.id for execution in executions] == [successes[0].id]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(("enabled", "state"), [(False, "active"), (True, "archived")])
|
|
async def test_enqueue_rejects_disabled_or_archived_jobs(
|
|
database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
|
|
enabled: bool,
|
|
state: str,
|
|
) -> None:
|
|
sessions, _ = database
|
|
async with sessions() as db:
|
|
job_id = await create_job(
|
|
db, f"unavailable-{enabled}-{state}", enabled=enabled, state=state
|
|
)
|
|
with pytest.raises(EnqueueError) as raised:
|
|
await enqueue(db, job_id)
|
|
|
|
assert raised.value.code == "job_disabled"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reclaimed_lease_fences_the_previous_worker(
|
|
database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
|
|
) -> None:
|
|
sessions, _ = database
|
|
async with sessions() as db:
|
|
execution = await enqueue(db, await create_job(db, "leases"))
|
|
assert await claim(db, execution.id, "worker-a") is not None
|
|
assert await heartbeat(db, execution.id, "worker-a")
|
|
persisted = await db.get(Execution, execution.id)
|
|
assert persisted is not None
|
|
persisted.lease_expires_at = datetime.now(UTC) - timedelta(seconds=1)
|
|
await db.commit()
|
|
|
|
async with sessions() as db:
|
|
assert await claim(db, execution.id, "worker-b") is not None
|
|
assert not await heartbeat(db, execution.id, "worker-a")
|
|
assert await request_cancellation(db, execution.id) is not None
|
|
assert not await complete_cancellation(db, execution.id, "worker-a")
|
|
assert await complete_cancellation(db, execution.id, "worker-b")
|
|
|
|
async with sessions() as db:
|
|
persisted = await db.get(Execution, execution.id)
|
|
assert persisted is not None
|
|
assert persisted.state == "cancelled"
|
|
assert persisted.lease_owner is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_reuses_execution_and_rejects_non_transient_failures(
|
|
database: tuple[async_sessionmaker[AsyncSession], AsyncEngine],
|
|
) -> None:
|
|
sessions, _ = database
|
|
async with sessions() as db:
|
|
execution = await enqueue(db, await create_job(db, "retry"))
|
|
execution.state = "failed"
|
|
execution.reason_code = "transient_io"
|
|
execution.operator_message = "temporary failure"
|
|
execution.lease_owner = "worker-a"
|
|
execution.lease_expires_at = datetime.now(UTC) + timedelta(seconds=60)
|
|
await db.commit()
|
|
|
|
retried = await retry(db, execution.id)
|
|
assert retried is not None
|
|
assert retried.id == execution.id
|
|
assert retried.state == "queued"
|
|
assert retried.attempt == 2
|
|
assert retried.reason_code is None
|
|
assert retried.operator_message is None
|
|
assert retried.lease_owner is None
|
|
assert retried.lease_expires_at is None
|
|
|
|
retried.state = "failed"
|
|
retried.reason_code = "integrity_failure"
|
|
await db.commit()
|
|
with pytest.raises(EnqueueError) as raised:
|
|
await retry(db, execution.id)
|
|
|
|
assert raised.value.code == "retry_not_allowed"
|