198 lines
6.7 KiB
Python
198 lines
6.7 KiB
Python
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()
|