"""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) assert prov.max_concurrent <= concurrency, f"ran {prov.max_concurrent} at once, cap was {concurrency}"