Quiiz system

tests/unit/services/devii/tasks/scheduler.py

@@ -0,0 +1,194 @@
+# 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"]