|
# retoor <retoor@molodetz.nl>
|
|
|
|
import asyncio
|
|
|
|
from devplacepy.database import invalidate_admins_cache
|
|
from devplacepy.services.devii.tasks.guards import (
|
|
REASON_EXPIRED,
|
|
REASON_MAX_RUNS,
|
|
REASON_NOT_ADMIN,
|
|
)
|
|
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
|
|
from devplacepy.services.devii.tasks.scheduler import GlobalScheduler
|
|
from devplacepy.services.devii.tasks.store import TaskStore
|
|
from devplacepy.utils import generate_uid
|
|
from tests.conftest import run_async
|
|
|
|
|
|
def _account(local_db, role):
|
|
uid = generate_uid()
|
|
local_db["users"].insert(
|
|
{
|
|
"uid": uid,
|
|
"username": f"sched-{uid[-10:]}",
|
|
"role": role,
|
|
"api_key": "k",
|
|
"deleted_at": None,
|
|
}
|
|
)
|
|
invalidate_admins_cache()
|
|
return uid
|
|
|
|
|
|
def _seed(local_db, owner, **overrides):
|
|
row = {
|
|
"uid": generate_uid(),
|
|
"owner_kind": "user",
|
|
"owner_id": owner,
|
|
"label": "scheduled",
|
|
"prompt": "work",
|
|
"enabled": True,
|
|
"status": "pending",
|
|
"kind": "interval",
|
|
"every_seconds": 900,
|
|
"max_runs": 10,
|
|
"run_count": 0,
|
|
"failure_count": 0,
|
|
"created_at": to_iso(now_utc()),
|
|
"next_run_at": to_iso(now_utc()),
|
|
"expires_at": None,
|
|
"deleted_at": None,
|
|
"deleted_by": None,
|
|
}
|
|
row.update(overrides)
|
|
local_db["devii_tasks"].insert(row)
|
|
return row["uid"]
|
|
|
|
|
|
def _harness(local_db, started, gate, peak=None):
|
|
active: dict[str, int] = {}
|
|
|
|
def resolve(row):
|
|
owner = str(row["owner_id"])
|
|
store = TaskStore(local_db, "user", owner)
|
|
|
|
async def executor(prompt):
|
|
started.append(owner)
|
|
active[owner] = active.get(owner, 0) + 1
|
|
if peak is not None:
|
|
peak.append(max(active.values()))
|
|
try:
|
|
await gate.wait()
|
|
return f"done: {prompt}"
|
|
finally:
|
|
active[owner] -= 1
|
|
|
|
return store, executor, lambda kind, task_row, payload: None
|
|
|
|
return resolve
|
|
|
|
|
|
def test_scheduler_runs_one_task_per_owner_and_completes_it(local_db):
|
|
local_db["devii_tasks"].delete()
|
|
owner = _account(local_db, "Admin")
|
|
first = _seed(local_db, owner)
|
|
second = _seed(local_db, owner)
|
|
|
|
started: list[str] = []
|
|
peak: list[int] = []
|
|
|
|
async def run():
|
|
gate = asyncio.Event()
|
|
scheduler = GlobalScheduler(
|
|
local_db, _harness(local_db, started, gate, peak), tick_seconds=0.05
|
|
)
|
|
scheduler.start()
|
|
await asyncio.sleep(0.4)
|
|
claimed = [r["uid"] for r in local_db["devii_tasks"].find(status="running")]
|
|
held = len(started)
|
|
gate.set()
|
|
await asyncio.sleep(0.5)
|
|
await scheduler.stop()
|
|
return claimed, held
|
|
|
|
claimed, held = run_async(run())
|
|
assert held == 1
|
|
assert max(peak) == 1
|
|
assert len(claimed) == 1
|
|
done = local_db["devii_tasks"].find_one(uid=claimed[0])
|
|
assert int(done["run_count"]) == 1
|
|
assert done["status"] == "pending"
|
|
assert done["last_result"].startswith("done:")
|
|
assert done["next_run_at"] > to_iso(now_utc())
|
|
other = second if claimed[0] == first else first
|
|
assert local_db["devii_tasks"].find_one(uid=other)["status"] == "pending"
|
|
|
|
|
|
def test_scheduler_retires_a_task_whose_owner_may_not_schedule(local_db):
|
|
local_db["devii_tasks"].delete()
|
|
uid = _seed(local_db, _account(local_db, "Member"))
|
|
started: list[str] = []
|
|
|
|
async def run():
|
|
gate = asyncio.Event()
|
|
gate.set()
|
|
scheduler = GlobalScheduler(
|
|
local_db, _harness(local_db, started, gate), tick_seconds=0.05
|
|
)
|
|
scheduler.start()
|
|
await asyncio.sleep(0.3)
|
|
await scheduler.stop()
|
|
|
|
run_async(run())
|
|
row = local_db["devii_tasks"].find_one(uid=uid)
|
|
assert not row["enabled"]
|
|
assert row["status"] == "disabled"
|
|
assert row["last_error"] == REASON_NOT_ADMIN
|
|
assert started == []
|
|
|
|
|
|
def test_scheduler_retires_expired_and_exhausted_tasks_even_while_saturated(local_db):
|
|
local_db["devii_tasks"].delete()
|
|
owner = _account(local_db, "Admin")
|
|
_seed(local_db, owner)
|
|
expired = _seed(
|
|
local_db, owner, expires_at=to_iso(now_utc().replace(year=now_utc().year - 1))
|
|
)
|
|
exhausted = _seed(local_db, owner, max_runs=3, run_count=3)
|
|
started: list[str] = []
|
|
|
|
async def run():
|
|
gate = asyncio.Event()
|
|
scheduler = GlobalScheduler(
|
|
local_db, _harness(local_db, started, gate), tick_seconds=0.05
|
|
)
|
|
scheduler.start()
|
|
await asyncio.sleep(0.4)
|
|
gate.set()
|
|
await asyncio.sleep(0.2)
|
|
await scheduler.stop()
|
|
|
|
run_async(run())
|
|
assert local_db["devii_tasks"].find_one(uid=expired)["last_error"] == REASON_EXPIRED
|
|
assert local_db["devii_tasks"].find_one(uid=exhausted)["last_error"] == REASON_MAX_RUNS
|
|
for uid in (expired, exhausted):
|
|
assert not local_db["devii_tasks"].find_one(uid=uid)["enabled"]
|
|
|
|
|
|
def test_scheduler_disables_a_task_after_repeated_failures(local_db):
|
|
local_db["devii_tasks"].delete()
|
|
owner = _account(local_db, "Admin")
|
|
uid = _seed(local_db, owner, failure_count=2)
|
|
|
|
def resolve(row):
|
|
store = TaskStore(local_db, "user", str(row["owner_id"]))
|
|
|
|
async def executor(prompt):
|
|
raise RuntimeError("upstream is down")
|
|
|
|
return store, executor, lambda kind, task_row, payload: None
|
|
|
|
async def run():
|
|
scheduler = GlobalScheduler(
|
|
local_db, resolve, tick_seconds=0.05, max_failures=3
|
|
)
|
|
scheduler.start()
|
|
await asyncio.sleep(0.4)
|
|
await scheduler.stop()
|
|
|
|
run_async(run())
|
|
row = local_db["devii_tasks"].find_one(uid=uid)
|
|
assert int(row["failure_count"]) == 3
|
|
assert not row["enabled"]
|
|
assert row["status"] == "error"
|
|
assert "upstream is down" in row["last_error"]
|