291 lines
10 KiB
Python
291 lines
10 KiB
Python
import asyncio
|
|
import inspect
|
|
import os
|
|
import tempfile
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import psycopg
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
from pgqueuer import PsycopgDriver, Queries
|
|
from test_conversations import World
|
|
|
|
from beaver_gateway.conversations.service import parse_at
|
|
from beaver_gateway.jobs.scheduler import Budget, Job, JobRun, Scheduler, next_run
|
|
from beaver_gateway.storage.models import RateLimit
|
|
|
|
DATABASE_URL = os.environ.get("TEST_DATABASE_URL")
|
|
WARSAW = ZoneInfo("Europe/Warsaw")
|
|
|
|
|
|
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:
|
|
if not DATABASE_URL:
|
|
pytest.skip("TEST_DATABASE_URL (postgres) is not set")
|
|
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 inspect.isawaitable(value):
|
|
value = await value
|
|
if value:
|
|
return value
|
|
await asyncio.sleep(0.05)
|
|
msg = "condition never happened"
|
|
raise AssertionError(msg)
|
|
|
|
|
|
def test_cron_is_read_in_the_gateway_tz() -> None:
|
|
summer = datetime(2026, 7, 1, 12, tzinfo=UTC)
|
|
assert next_run("0 4 * * *", WARSAW, summer) == datetime(2026, 7, 2, 2, tzinfo=UTC)
|
|
assert next_run("0 4 * * *", UTC, summer) == datetime(2026, 7, 2, 4, tzinfo=UTC)
|
|
before_fallback = datetime(2026, 10, 24, 12, tzinfo=UTC)
|
|
assert next_run("0 4 * * *", WARSAW, before_fallback) == datetime(
|
|
2026, 10, 25, 3, tzinfo=UTC
|
|
)
|
|
during_fallback = datetime(2026, 10, 25, 1, 30, tzinfo=UTC)
|
|
assert next_run("0 4 * * *", WARSAW, during_fallback) == datetime(
|
|
2026, 10, 25, 3, tzinfo=UTC
|
|
)
|
|
assert next_run("*/15 * * * *", WARSAW, during_fallback) == datetime(
|
|
2026, 10, 25, 1, 45, tzinfo=UTC
|
|
)
|
|
soon = next_run("* * * * *", WARSAW)
|
|
assert soon.tzinfo is UTC
|
|
assert timedelta(0) < soon - datetime.now(UTC) <= timedelta(minutes=1)
|
|
|
|
|
|
def test_parse_at_reads_naive_iso_in_tz() -> None:
|
|
assert parse_at("2026-09-01T10:00", WARSAW) == datetime(
|
|
2026, 9, 1, 8, 0, tzinfo=UTC
|
|
)
|
|
assert parse_at("2026-12-01T10:00", WARSAW) == datetime(
|
|
2026, 12, 1, 9, 0, tzinfo=UTC
|
|
)
|
|
assert parse_at("2026-09-01T10:00:00+02:00", WARSAW) == datetime(
|
|
2026, 9, 1, 8, 0, tzinfo=UTC
|
|
)
|
|
assert parse_at("2026-09-01T10:00") == datetime(2026, 9, 1, 10, 0, tzinfo=UTC)
|
|
|
|
|
|
async def test_runs_are_recorded_without_postgres(world: World) -> None:
|
|
async def ok(run: JobRun) -> None:
|
|
run.payload["seen"] = True
|
|
|
|
async def boom(_run: JobRun) -> None:
|
|
msg = "no network"
|
|
raise ConnectionError(msg)
|
|
|
|
scheduler = Scheduler(
|
|
conversations=world.conversations,
|
|
jobs=[Job("ok", ok, webhook=True), Job("boom", boom)],
|
|
tz="Europe/Warsaw",
|
|
)
|
|
await scheduler.start()
|
|
await scheduler.hook("ok", {"stack": "x"})
|
|
await scheduler.trigger(scheduler.job("boom"))
|
|
await until(lambda: len(scheduler._runs) == 2) # noqa: SLF001
|
|
runs = await scheduler.runs("boom")
|
|
assert [r["status"] for r in runs] == ["failed"]
|
|
assert runs[0]["error"].startswith("ConnectionError: no network\n")
|
|
assert runs[0]["trigger"] == "manual"
|
|
assert runs[0]["started_at"] <= runs[0]["finished_at"]
|
|
assert runs[0]["started_at"].endswith("+02:00")
|
|
ok_runs = await scheduler.runs("ok")
|
|
assert ok_runs[0]["payload"] == {"stack": "x", "seen": True}
|
|
assert ok_runs[0]["trigger"] == "webhook"
|
|
await scheduler.hook("ok", {"stack": "y"})
|
|
await until(lambda: scheduler._runs["ok"]["payload"].get("stack") == "y") # noqa: SLF001
|
|
assert len(await scheduler.runs("ok")) == 2
|
|
assert len(await scheduler.runs("ok", limit=1)) == 1
|
|
snapshot = await scheduler.snapshot()
|
|
assert snapshot["tz"] == "Europe/Warsaw"
|
|
runs_by_name = {j["name"]: j["run"] for j in snapshot["jobs"]}
|
|
assert runs_by_name["ok"]["payload"]["stack"] == "y"
|
|
assert runs_by_name["boom"]["status"] == "failed"
|
|
await scheduler.stop()
|
|
|
|
|
|
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(), tz="Europe/Warsaw"
|
|
)
|
|
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
|
|
assert when.utcoffset() == datetime.now(WARSAW).utcoffset()
|
|
queued = await world.conversations.schedules(conv)
|
|
assert [q["payload"]["text"] for q in queued] == ["push X"]
|
|
assert queued[0]["payload"]["at"] == when.isoformat(timespec="minutes")
|
|
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("[inject: 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"}]
|
|
await until(lambda: scheduler._runs.get("ping")) # noqa: SLF001
|
|
snapshot = await scheduler.snapshot()
|
|
assert [j["name"] for j in snapshot["jobs"]] == ["ping", "quiet"]
|
|
assert snapshot["enabled"] and snapshot["jobs"][0]["webhook"]
|
|
assert snapshot["jobs"][0]["run"]["status"] == "done"
|
|
assert snapshot["jobs"][0]["run"]["trigger"] == "webhook"
|
|
runs = await scheduler.runs("ping")
|
|
assert [r["payload"] for r in runs] == [{"stack": "x"}]
|
|
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),
|
|
tz="Europe/Warsaw",
|
|
)
|
|
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 * * * *"}
|
|
|
|
async def crons_ready() -> dict[str, Any] | None:
|
|
snap = await scheduler.snapshot()
|
|
return snap if all(j["next_run"] for j in snap["jobs"]) else None
|
|
|
|
snapshot = await until(crons_ready)
|
|
expected = datetime.now(WARSAW).utcoffset()
|
|
for job_info in snapshot["jobs"]:
|
|
assert datetime.fromisoformat(job_info["next_run"]).utcoffset() == expected
|
|
await scheduler.stop()
|