wehub-resource-sync
302 lines
10 KiB
Python
302 lines
10 KiB
Python
"""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)]
|