"""Cancellation primitives on WorkerPool. These tests stay at the public-ish surface of `WorkerPool` — they exercise the hook registration contextvar that workers use and verify the dispatcher marks cancelled events as failed with the documented marker. They do NOT spin up a real omp subprocess; that's covered by the integration smoke test. """ from __future__ import annotations import asyncio from contextlib import suppress import pytest from robomp.cancellation import ( clear_current_event, register_cancel_hook, set_current_event, unregister_cancel_hook, ) from robomp.config import Settings from robomp.db import Database, EventRow from robomp.queue import WorkerPool from robomp.slot_pool import SlotPool class _StubGitHub: """Sentinel; queue tests don't talk to GitHub.""" class _StubSandbox: """Sentinel; queue tests don't touch the workspace pool.""" natives_cache = None class _StubGitTransport: """Sentinel; queue tests don't push.""" def _make_pool(settings: Settings, db: Database) -> WorkerPool: return WorkerPool( settings=settings, db=db, github=_StubGitHub(), # type: ignore[arg-type] sandbox=_StubSandbox(), # type: ignore[arg-type] git_transport=_StubGitTransport(), # type: ignore[arg-type] slot_pool=SlotPool(), ) def _row(delivery: str = "d1") -> EventRow: return EventRow( delivery_id=delivery, event_type="issues", repo="octo/widget", issue_key="octo/widget#1", payload={"action": "opened"}, received_at="2026-01-01T00:00:00Z", state="running", attempts=1, last_error=None, ) @pytest.mark.asyncio async def test_cancel_fires_hook_armed_by_worker(settings: Settings, db: Database) -> None: """A worker that armed a hook gets it invoked when cancel_event runs.""" pool = _make_pool(settings, db) row = _row() fired = asyncio.Event() async def fake_worker() -> None: # Mimic _run_event entering its contextvar scope: the helpers below are # what worker.py invokes from inside the asyncio.to_thread call. token = set_current_event(pool, row.delivery_id) try: await asyncio.to_thread(register_cancel_hook, fired.set) # Park until somebody fires the hook. await fired.wait() finally: await asyncio.to_thread(unregister_cancel_hook) clear_current_event(token) worker = asyncio.create_task(fake_worker()) try: # Give the worker a tick to register. for _ in range(20): await asyncio.sleep(0) if row.delivery_id in pool._cancel_hooks: # noqa: SLF001 — test inspecting state break assert row.delivery_id in pool._cancel_hooks # noqa: SLF001 assert await pool.cancel_event(row.delivery_id) is True await asyncio.wait_for(worker, timeout=1.0) assert row.delivery_id in pool._cancelled # noqa: SLF001 # Hook is consumed. assert row.delivery_id not in pool._cancel_hooks # noqa: SLF001 finally: if not worker.done(): worker.cancel() with suppress(asyncio.CancelledError): await worker @pytest.mark.asyncio async def test_cancel_before_arm_fires_immediately(settings: Settings, db: Database) -> None: """Cancelling before the worker arms must still terminate it on register.""" pool = _make_pool(settings, db) row = _row("d2") # Cancel is requested before any worker has armed a hook. assert await pool.cancel_event(row.delivery_id) is False assert row.delivery_id in pool._cancelled # noqa: SLF001 # When the worker eventually registers, the hook must fire synchronously. calls: list[int] = [] token = set_current_event(pool, row.delivery_id) try: register_cancel_hook(lambda: calls.append(1)) finally: clear_current_event(token) assert calls == [1] # Late-armed hook is NOT retained; cancel state is one-shot. assert row.delivery_id not in pool._cancel_hooks # noqa: SLF001 @pytest.mark.asyncio async def test_dispatch_marks_cancelled_event_failed_with_marker( settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch ) -> None: """A dispatch that observed cancellation marks the row failed + 'cancelled by operator'.""" pool = _make_pool(settings, db) db.record_event( delivery_id="d3", event_type="issues", repo="octo/widget", issue_key="octo/widget#1", payload={"action": "opened"}, state="running", ) row = _row("d3") async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None: # Simulate cancellation hitting mid-task and the omp subprocess raising. await pool.cancel_event(r.delivery_id) raise RuntimeError("subprocess died") monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch) await pool._run_event(row) # noqa: SLF001 — testing the dispatcher branch directly stored = db.get_event("d3") assert stored is not None assert stored.state == "failed" assert stored.last_error == "cancelled by operator" # State is cleared for future events. assert row.delivery_id not in pool._cancelled # noqa: SLF001 assert row.delivery_id not in pool._cancel_hooks # noqa: SLF001 @pytest.mark.asyncio async def test_non_cancelled_failure_keeps_real_traceback( settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch ) -> None: """A garden-variety dispatch failure still records the traceback path.""" pool = _make_pool(settings, db) monkeypatch.setattr(settings, "event_max_retries", 0) # assert terminal failure, not retry db.record_event( delivery_id="d4", event_type="issues", repo="octo/widget", issue_key="octo/widget#1", payload={"action": "opened"}, state="running", ) row = _row("d4") async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None: raise ValueError("boom 42") monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch) await pool._run_event(row) # noqa: SLF001 stored = db.get_event("d4") assert stored is not None assert stored.state == "failed" assert stored.last_error is not None assert "boom 42" in stored.last_error assert "cancelled by operator" not in stored.last_error @pytest.mark.asyncio async def test_run_event_marks_failed_when_not_shutting_down( settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch ) -> None: """When `_shutting_down` is False, a dispatch failure still marks the row failed.""" pool = _make_pool(settings, db) monkeypatch.setattr(settings, "event_max_retries", 0) # assert terminal failure, not retry assert pool._shutting_down is False # noqa: SLF001 db.record_event( delivery_id="d5", event_type="issues", repo="octo/widget", issue_key="octo/widget#1", payload={"action": "opened"}, state="running", ) row = _row("d5") async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None: raise RuntimeError("regular failure") monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch) await pool._run_event(row) # noqa: SLF001 stored = db.get_event("d5") assert stored is not None assert stored.state == "failed" assert stored.last_error is not None assert "regular failure" in stored.last_error @pytest.mark.asyncio async def test_cancel_unknown_delivery_returns_false(settings: Settings, db: Database) -> None: """Cancelling an unknown delivery is a no-op that returns False.""" pool = _make_pool(settings, db) assert await pool.cancel_event("never-existed") is False # The set still records the request — a later register would fire — but # since no worker is armed, the cancel is harmless. assert "never-existed" in pool._cancelled # noqa: SLF001 @pytest.mark.asyncio async def test_start_reaps_configured_slot_uids( settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch ) -> None: calls: list[int] = [] monkeypatch.setattr("robomp.queue._reap_slot", lambda uid: calls.append(uid)) pool = WorkerPool( settings=settings, db=db, github=_StubGitHub(), # type: ignore[arg-type] sandbox=_StubSandbox(), # type: ignore[arg-type] git_transport=_StubGitTransport(), # type: ignore[arg-type] slot_pool=SlotPool([2001, 2002]), ) await pool.start() try: assert sorted(calls) == [2001, 2002] finally: await pool.stop(drain_timeout=0.01, kill_timeout=0.01) @pytest.mark.asyncio async def test_run_event_reaps_slot_before_release( settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch ) -> None: slot_pool = SlotPool([2001]) pool = WorkerPool( settings=settings, db=db, github=_StubGitHub(), # type: ignore[arg-type] sandbox=_StubSandbox(), # type: ignore[arg-type] git_transport=_StubGitTransport(), # type: ignore[arg-type] slot_pool=slot_pool, ) db.record_event( delivery_id="d-slot", event_type="issues", repo="octo/widget", issue_key="octo/widget#1", payload={"action": "opened"}, state="running", ) order: list[tuple[str, int | None]] = [] monkeypatch.setattr("robomp.queue._reap_slot", lambda uid: order.append(("reap", uid))) release = slot_pool.release def record_release(slot_uid: int | None) -> None: order.append(("release", slot_uid)) release(slot_uid) monkeypatch.setattr(slot_pool, "release", record_release) async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None: assert r.delivery_id == "d-slot" assert slot_uid == 2001 monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch) await pool._run_event(_row("d-slot")) # noqa: SLF001 stored = db.get_event("d-slot") assert stored is not None assert stored.state == "done" assert order == [("reap", 2001), ("release", 2001)]