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.statuses(conv) == [("wake", "done")] 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()