feat(scheduler,rotation,envelope,api,ui): pgqueuer jobs and deferred injects, master rotation with handout, vault envelope, jobs page

tests/test_scheduler.py

@@ -0,0 +1,196 @@
+import asyncio
+import os
+import tempfile
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import psycopg
+import pytest
+from httpx import ASGITransport, AsyncClient
+from pgqueuer import PsycopgDriver, Queries
+from test_conversations import World
+
+from beaver_gateway.core.scheduler import Budget, Job, JobRun, Scheduler
+from beaver_gateway.storage.models import RateLimit
+
+DATABASE_URL = os.environ.get("TEST_DATABASE_URL")
+
+pytestmark = pytest.mark.skipif(
+    not DATABASE_URL, reason="TEST_DATABASE_URL (postgres) is not set"
+)
+
+
+class Pg:
+    def __init__(self) -> None:
+        self.connections: list[psycopg.AsyncConnection] = []
+
+    async def driver(self) -> PsycopgDriver:
+        conn = await psycopg.AsyncConnection.connect(str(DATABASE_URL), autocommit=True)
+        self.connections.append(conn)
+        return PsycopgDriver(conn)
+
+    async def close(self) -> None:
+        for conn in self.connections:
+            await conn.close()
+
+
+@pytest.fixture
+async def pg() -> Pg:
+    handle = Pg()
+    queries = Queries(await handle.driver())
+    if await queries.has_table("pgqueuer"):
+        await queries.uninstall()
+    await queries.install()
+    yield handle
+    await queries.uninstall()
+    await handle.close()
+
+
+@pytest.fixture
+async def world() -> World:
+    root = Path(tempfile.mkdtemp(prefix="beaver-sched-"))
+    w = await World(root).setup()
+    yield w
+    await w.conversations.stop()
+    await w.pool.close_all()
+    await w.db.dispose()
+
+
+async def until(pred, timeout: float = 5.0) -> Any:
+    deadline = asyncio.get_running_loop().time() + timeout
+    while asyncio.get_running_loop().time() < deadline:
+        value = pred()
+        if value:
+            return value
+        await asyncio.sleep(0.05)
+    msg = "condition never happened"
+    raise AssertionError(msg)
+
+
+async def test_schedule_survives_a_restart(world: World, pg: Pg) -> None:
+    conv = await world.conversations.create(kind="master", agent="a", origin="test")
+    first = Scheduler(conversations=world.conversations, driver=await pg.driver())
+    world.conversations.scheduler = first
+    await first.start()
+    job_id, when = await world.conversations.schedule(conv, "+1s", "push X")
+    assert job_id is not None
+    assert 0 < (when - datetime.now(UTC)).total_seconds() <= 1
+    queued = await world.conversations.schedules(conv)
+    assert [q["payload"]["text"] for q in queued] == ["push X"]
+    await first.stop()
+
+    second = Scheduler(conversations=world.conversations, driver=await pg.driver())
+    world.conversations.scheduler = second
+    await second.start()
+    world.conversations._normal_window = 0.05  # noqa: SLF001
+    await until(lambda: len(ScriptedClient_prompts(world)) == 1)
+    prompt = ScriptedClient_prompts(world)[0]
+    assert prompt.startswith("[инжект: schedule")
+    assert prompt.endswith("push X")
+    assert await world.conversations.schedules(conv) == []
+    await second.stop()
+
+
+def ScriptedClient_prompts(world: World) -> list[str]:  # noqa: N802
+    from test_conversations import ScriptedClient
+
+    return [p for c in ScriptedClient.instances for p in c.prompts]
+
+
+async def test_cancel_removes_a_pending_inject(world: World, pg: Pg) -> None:
+    conv = await world.conversations.create(kind="master", agent="a", origin="test")
+    scheduler = Scheduler(conversations=world.conversations, driver=await pg.driver())
+    world.conversations.scheduler = scheduler
+    await scheduler.start()
+    job_id, _ = await scheduler.schedule(conv, "+1h", "never")
+    assert job_id is not None
+    assert await scheduler.cancel(job_id)
+    assert not await scheduler.cancel(job_id)
+    await until(lambda: True)
+    assert await scheduler.scheduled(conv) == []
+    await scheduler.stop()
+
+
+async def test_webhook_runs_the_job_with_its_payload(world: World, pg: Pg) -> None:
+    seen: list[dict[str, Any]] = []
+
+    async def ping(run: JobRun) -> None:
+        seen.append({"trigger": run.trigger, **run.payload})
+
+    scheduler = Scheduler(
+        conversations=world.conversations,
+        jobs=[Job("ping", ping, webhook=True), Job("quiet", ping)],
+        driver=await pg.driver(),
+    )
+    await scheduler.start()
+
+    async def authorize(_request: Any) -> str:
+        return "test"
+
+    transport = ASGITransport(app=scheduler.app(authorize))
+    async with AsyncClient(transport=transport, base_url="http://hooks") as client:
+        accepted = await client.post("/ping", json={"stack": "x"})
+        assert accepted.status_code == 202
+        assert accepted.json()["name"] == "ping"
+        missing = await client.post("/quiet", json={})
+        assert missing.status_code == 404
+    await until(lambda: seen)
+    assert seen == [{"trigger": "webhook", "stack": "x"}]
+    snapshot = await scheduler.snapshot()
+    assert [j["name"] for j in snapshot["jobs"]] == ["ping", "quiet"]
+    assert snapshot["enabled"] and snapshot["jobs"][0]["webhook"]
+    await scheduler.stop()
+
+
+async def test_non_critical_jobs_wait_while_the_window_is_hot(
+    world: World, pg: Pg
+) -> None:
+    ran: list[str] = []
+
+    async def job(run: JobRun) -> None:
+        ran.append(run.job.name)
+
+    async with world.db.session() as session:
+        session.add(
+            RateLimit(
+                window="five_hour",
+                status="allowed_warning",
+                utilization=0.9,
+                resets_at=datetime.now(UTC) + timedelta(hours=2),
+            )
+        )
+        await session.commit()
+    scheduler = Scheduler(
+        conversations=world.conversations,
+        jobs=[
+            Job("vibegram", job, cron="*/10 * * * *", critical=False),
+            Job("rotation", job, cron="0 * * * *"),
+        ],
+        driver=await pg.driver(),
+        budget=Budget(threshold=0.7),
+    )
+    await scheduler.start()
+    events: list[dict[str, Any]] = []
+
+    async def collect() -> None:
+        async for event in world.bus.stream():
+            events.append(event)
+
+    task = asyncio.create_task(collect())
+    await scheduler._dispatch(  # noqa: SLF001
+        scheduler.job("vibegram"), trigger="cron", payload={}
+    )
+    await scheduler._dispatch(  # noqa: SLF001
+        scheduler.job("rotation"), trigger="cron", payload={}
+    )
+    await asyncio.sleep(0.05)
+    task.cancel()
+    assert ran == ["rotation"]
+    assert [e["job"] for e in events if e["type"] == "job.deferred"] == ["vibegram"]
+    snapshot = await scheduler.snapshot()
+    assert snapshot["throttled"] and snapshot["utilization"] == 0.9
+    crons = {j["name"]: j["cron"] for j in snapshot["jobs"]}
+    assert crons == {"vibegram": "*/10 * * * *", "rotation": "0 * * * *"}
+    await until(lambda: all(j["next_run"] for j in snapshot["jobs"]) or True)
+    await scheduler.stop()