6033967e6a
Ship Keystone auth, multi-service handlers (Yoga→Dalmatian), Compose/Helm packaging, API contract packs, and pytest/Pulumi coverage labs.
82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
"""Durable task concurrency and recovery tests."""
|
|
|
|
import asyncio
|
|
import os
|
|
import uuid
|
|
|
|
import asyncpg # type: ignore[import-untyped]
|
|
import pytest
|
|
from asyncpg import Pool
|
|
|
|
from app.db.migrations import migrate
|
|
from app.tasks.repository import TaskRepository
|
|
|
|
pytestmark = [
|
|
pytest.mark.integration,
|
|
pytest.mark.skipif(not os.getenv("TEST_DATABASE_URL"), reason="TEST_DATABASE_URL is required"),
|
|
]
|
|
|
|
|
|
async def repository() -> tuple[Pool, TaskRepository]:
|
|
pool = await asyncpg.create_pool(os.environ["TEST_DATABASE_URL"], min_size=1, max_size=4)
|
|
async with pool.acquire() as connection:
|
|
await migrate(connection)
|
|
return pool, TaskRepository(pool)
|
|
|
|
|
|
async def test_two_worker_exclusion_idempotency_and_logs() -> None:
|
|
pool, tasks = await repository()
|
|
key = uuid.uuid4().hex
|
|
try:
|
|
created = await tasks.create(
|
|
upid=f"UPID:pve1:00000001:00000001:00000001:test:{key}:root@pam:",
|
|
task_type="test",
|
|
payload={"value": 1},
|
|
resource_key=f"vm:{key}",
|
|
idempotency_key=key,
|
|
)
|
|
repeated = await tasks.create(
|
|
upid=f"ignored-{key}", task_type="test", payload={}, idempotency_key=key
|
|
)
|
|
assert repeated.id == created.id
|
|
|
|
first, second = await asyncio.gather(
|
|
tasks.claim("worker-a", 30), tasks.claim("worker-b", 30)
|
|
)
|
|
claimed = first or second
|
|
assert claimed is not None
|
|
assert (first is None) != (second is None)
|
|
worker = "worker-a" if first is not None else "worker-b"
|
|
await tasks.append_log(claimed.id, "started")
|
|
await tasks.progress(claimed.id, worker, 50)
|
|
await tasks.finish(claimed.id, worker, status="success", result={"ok": True})
|
|
assert await tasks.logs(claimed.id) == ("started",)
|
|
finished = await tasks.get(claimed.id)
|
|
assert finished is not None
|
|
assert finished.status == "success"
|
|
finally:
|
|
await pool.close()
|
|
|
|
|
|
async def test_expired_lease_is_reclaimed_after_restart() -> None:
|
|
pool, tasks = await repository()
|
|
key = uuid.uuid4().hex
|
|
try:
|
|
created = await tasks.create(
|
|
upid=f"UPID:pve1:00000001:00000001:00000001:test:{key}:root@pam:",
|
|
task_type="test",
|
|
payload={},
|
|
)
|
|
assert await tasks.claim("dead-worker", 0) is not None
|
|
recovered = await tasks.claim("new-worker", 30)
|
|
assert recovered is not None
|
|
assert recovered.id == created.id
|
|
assert recovered.attempt == 2
|
|
await tasks.request_cancel(recovered.id)
|
|
cancelled = await tasks.get(recovered.id)
|
|
assert cancelled is not None
|
|
assert cancelled.cancel_requested
|
|
await tasks.finish(recovered.id, "new-worker", status="cancelled")
|
|
finally:
|
|
await pool.close()
|