svcforge/tests/integration/test_worker.py

"""The worker loop: drains on SIGTERM, retries, stays idempotent.

These run in milliseconds against a FakeProvisioner. That is the payoff for putting helm
behind a Protocol in Module 5: the crash-safety properties are testable without a cluster.
"""

from __future__ import annotations

import asyncio
import time
from pathlib import Path
from uuid import UUID

import pytest

from services.worker.deps import WorkerDeps
from services.worker.main import run_worker
from svcforge_core.adapters.clock import SystemClock
from svcforge_core.domain.catalog import load_catalog
from svcforge_core.domain.models import TaskKind
from svcforge_core.domain.states import InstanceState
from svcforge_core.repo.db import DictPool
from svcforge_core.repo.instances import InstanceRepo
from svcforge_core.repo.tasks import TaskRepo
from svcforge_core.settings import Settings
from tests.fakes import FakeNotifier, FakeProvisioner
from tests.integration.helpers import build_instance

CATALOG = load_catalog(Path(__file__).resolve().parents[2] / "catalog.yaml")


def _settings(**over: object) -> Settings:
    base: dict[str, object] = {
        "pg_dsn": "postgresql://x:x@127.0.0.1:5432/x",
        "worker_id": "w-test",
        "worker_concurrency": 4,
        "poll_interval_s": 0.05,
        "max_attempts": 3,
    }
    base.update(over)
    return Settings(**base)  # type: ignore[arg-type]


def _deps(
    pool: DictPool,
    prov: FakeProvisioner,
    notifier: FakeNotifier | None = None,
    **over: object,
) -> WorkerDeps:
    return WorkerDeps(
        pool=pool,
        instances=InstanceRepo(pool),
        tasks=TaskRepo(pool),
        provisioner=prov,
        notifier=notifier or FakeNotifier(),
        clock=SystemClock(),
        catalog=CATALOG,
        settings=_settings(**over),
    )


async def _seed(pool: DictPool, **kw: object) -> tuple[str, int]:
    inst = build_instance(**kw)  # type: ignore[arg-type]
    async with pool.connection() as conn:
        await InstanceRepo(pool).create(conn, inst)
    tid = await TaskRepo(pool).enqueue_standalone(inst.id, TaskKind.PROVISION)
    return str(inst.id), tid


async def _state_of(pool: DictPool, tid: int) -> str:
    async with pool.connection() as conn, conn.cursor() as cur:
        await cur.execute("select state from tasks where id = %s", (tid,))
        row = await cur.fetchone()
    assert row is not None
    return str(row["state"])


async def test_worker_provisions_and_marks_ready(pool: DictPool) -> None:
    iid, tid = await _seed(pool)
    prov = FakeProvisioner()
    notifier = FakeNotifier()
    stop = asyncio.Event()

    worker = asyncio.create_task(run_worker(_deps(pool, prov, notifier), stop))
    await asyncio.sleep(0.5)
    stop.set()
    await asyncio.wait_for(worker, timeout=5)

    assert await _state_of(pool, tid) == "done"
    inst = await InstanceRepo(pool).get(UUID(iid), team="platform")
    assert inst is not None
    assert inst.state is InstanceState.READY
    assert inst.endpoint is not None
    assert len(prov.installed) == 1

    # Assert the notification fired, and that it happened on the FIRST attempt.
    #
    # Without this the handler could raise after marking the instance ready — the task
    # requeues, the retry hits the idempotency early-return, and everything above still
    # passes while the worker is quietly crashing on every provision. Idempotency is
    # supposed to make crashes survivable, not invisible; asserting attempts==1 is what
    # keeps a masked crash from reading as success.
    assert notifier.events() == ["instance.ready"]

    async with pool.connection() as conn, conn.cursor() as cur:
        await cur.execute("select attempts from tasks where id = %s", (tid,))
        row = await cur.fetchone()
    assert row is not None
    assert row["attempts"] == 1, "task was retried: the handler raised after doing the work"


async def test_sigterm_drains_in_flight(pool: DictPool) -> None:
    """Stop is requested mid-provision: the worker must FINISH the task, then exit.

    Abandoning it would not lose the task — the lease would recover it — but only after
    five minutes of a tenant watching 'provisioning'. Draining costs two seconds.
    """
    _, tid = await _seed(pool)
    prov = FakeProvisioner(delay=2.0)
    stop = asyncio.Event()

    started = time.monotonic()
    worker = asyncio.create_task(run_worker(_deps(pool, prov), stop))
    await asyncio.sleep(0.5)
    stop.set()  # mid-flight: the handler is still inside its 2s install

    await asyncio.wait_for(worker, timeout=5)
    elapsed = time.monotonic() - started

    assert await _state_of(pool, tid) == "done", "worker abandoned an in-flight task"
    assert elapsed >= 2.0, "worker returned before the in-flight task finished"
    assert elapsed < 5.0


async def test_idle_worker_stops_promptly(pool: DictPool) -> None:
    """Nothing queued: stop must wake the poll sleep, not wait it out."""
    stop = asyncio.Event()
    worker = asyncio.create_task(run_worker(_deps(pool, FakeProvisioner(), poll_interval_s=5.0), stop))
    await asyncio.sleep(0.2)
    started = time.monotonic()
    stop.set()
    await asyncio.wait_for(worker, timeout=2)
    assert time.monotonic() - started < 1.0, "stop did not interrupt the poll sleep"


async def test_failed_task_is_requeued_with_backoff(pool: DictPool) -> None:
    _, tid = await _seed(pool)
    prov = FakeProvisioner(fail_on={"platform-elasticsearch"})
    stop = asyncio.Event()

    worker = asyncio.create_task(run_worker(_deps(pool, prov), stop))
    await asyncio.sleep(0.6)
    stop.set()
    await asyncio.wait_for(worker, timeout=5)

    async with pool.connection() as conn, conn.cursor() as cur:
        await cur.execute("select state, attempts, last_error from tasks where id = %s", (tid,))
        row = await cur.fetchone()
    assert row is not None
    assert row["state"] == "queued"  # requeued, not failed — attempts remain
    assert row["attempts"] >= 1
    assert row["last_error"]


async def test_provision_twice_installs_once(pool: DictPool) -> None:
    """The idempotency claim, executed.

    Simulates the crash window: the release is installed and the instance is READY, but
    the task got re-queued (worker died before reporting). Re-running must not re-install.
    """
    from services.worker.handlers import handle_provision

    inst = build_instance(state=InstanceState.REQUESTED)
    async with pool.connection() as conn:
        await InstanceRepo(pool).create(conn, inst)
    tid = await TaskRepo(pool).enqueue_standalone(inst.id, TaskKind.PROVISION)
    task = await TaskRepo(pool).claim("w1")
    assert task is not None

    prov = FakeProvisioner()
    deps = _deps(pool, prov)

    await handle_provision(task, deps)
    await handle_provision(task, deps)  # the redelivery

    assert len(prov.installed) == 1, "second run re-installed: handler is not idempotent"
    assert tid == task.id


@pytest.mark.parametrize("concurrency", [1, 4])
async def test_concurrency_cap_is_respected(pool: DictPool, concurrency: int) -> None:
    """The semaphore is what stops one worker from starting 200 helm processes."""
    for _ in range(6):
        await _seed(pool)

    prov = FakeProvisioner(delay=0.2)
    stop = asyncio.Event()
    worker = asyncio.create_task(run_worker(_deps(pool, prov, worker_concurrency=concurrency), stop))
    await asyncio.sleep(0.5)
    stop.set()
    await asyncio.wait_for(worker, timeout=10)

    # BOTH bounds. `<= concurrency` alone passes on a worker whose semaphore is broken to 1,
    # or that lost concurrency entirely — it only proves the cap is not exceeded, never that
    # concurrency exists. With 6 tasks seeded and a 0.2s handler, a working worker reaches
    # its cap exactly.
    assert prov.max_concurrent <= concurrency, f"ran {prov.max_concurrent} at once, cap was {concurrency}"
    assert prov.max_concurrent == concurrency, (
        f"only reached {prov.max_concurrent} concurrent with a cap of {concurrency}: "
        "the worker is not actually running tasks in parallel"
    )