svcforge/tests/integration/test_claim.py

"""The claim race. If only one test in this repo survives, it should be this one."""

from __future__ import annotations

import asyncio

from svcforge_core.domain.models import TaskKind
from svcforge_core.repo.db import DictPool
from svcforge_core.repo.tasks import TaskRepo
from tests.integration.helpers import make_instance


async def test_skip_locked_claims_each_task_exactly_once(pool: DictPool) -> None:
    """50 workers, 50 tasks, one claim each — no double-claims, no lost tasks.

    This is the test that fails if you split the claim into select-then-update.
    """
    repo = TaskRepo(pool)
    inst = await make_instance(pool)
    ids = {await repo.enqueue_standalone(inst, TaskKind.PROVISION) for _ in range(50)}

    async with asyncio.TaskGroup() as tg:
        claims = [tg.create_task(repo.claim(f"w{i}")) for i in range(50)]
    got = [c.result() for c in claims]

    assert all(t is not None for t in got)
    assert sorted(t.id for t in got if t is not None) == sorted(ids)  # each exactly once
    assert all(t.attempts == 1 for t in got if t is not None)


async def test_more_workers_than_tasks_get_none_not_a_duplicate(pool: DictPool) -> None:
    """Contention must produce None for the losers, never a second claim on one row."""
    repo = TaskRepo(pool)
    inst = await make_instance(pool)
    await repo.enqueue_standalone(inst, TaskKind.PROVISION)

    async with asyncio.TaskGroup() as tg:
        claims = [tg.create_task(repo.claim(f"w{i}")) for i in range(10)]
    got = [c.result() for c in claims]

    won = [t for t in got if t is not None]
    assert len(won) == 1
    assert len([t for t in got if t is None]) == 9