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

# retoor <retoor@molodetz.nl>

import asyncio
from datetime import timedelta

from devplacepy.database import invalidate_admins_cache
from devplacepy.services.devii.tasks import limits
from devplacepy.services.devii.tasks.context import in_task_run
from devplacepy.services.devii.tasks.guards import (
    REASON_BUDGET,
    REASON_EXPIRED,
    REASON_MAX_RUNS,
    REASON_NOT_A_USER,
    REASON_RUN_QUOTA,
)
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,
            "created_at": to_iso(now_utc()),
        }
    )
    invalidate_admins_cache()
    return uid


def _seed(local_db, owner, owner_kind="user", **overrides):
    row = {
        "uid": generate_uid(),
        "owner_kind": owner_kind,
        "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, flags=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)
            if flags is not None:
                flags.append(in_task_run())
            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 _run_scheduler(local_db, resolve, hold=0.4, release=0.5, **kwargs):
    async def run():
        gate = kwargs.pop("gate", None) or asyncio.Event()
        scheduler = GlobalScheduler(local_db, resolve, tick_seconds=0.05, **kwargs)
        scheduler.start()
        await asyncio.sleep(hold)
        claimed = [r["uid"] for r in local_db["devii_tasks"].find(status="running")]
        held = len(claimed)
        gate.set()
        await asyncio.sleep(release)
        await scheduler.stop()
        return claimed, held

    return run


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] = []
    gate = asyncio.Event()

    claimed, held = run_async(
        _run_scheduler(
            local_db, _harness(local_db, started, gate, peak), gate=gate
        )()
    )
    assert held == 1
    assert max(peak) == 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_runs_a_member_task(local_db):
    local_db["devii_tasks"].delete()
    owner = _account(local_db, "Member")
    uid = _seed(local_db, owner)
    started: list[str] = []
    gate = asyncio.Event()

    run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
    assert started == [owner]
    row = local_db["devii_tasks"].find_one(uid=uid)
    assert row["enabled"]
    assert int(row["run_count"]) == 1


def test_the_executor_runs_inside_the_task_run_context(local_db):
    local_db["devii_tasks"].delete()
    owner = _account(local_db, "Member")
    _seed(local_db, owner)
    started: list[str] = []
    flags: list[bool] = []
    gate = asyncio.Event()

    run_async(
        _run_scheduler(
            local_db, _harness(local_db, started, gate, flags=flags), gate=gate
        )()
    )
    assert flags == [True]
    assert in_task_run() is False


def test_every_run_is_recorded_against_the_owner_quota(local_db):
    local_db["devii_tasks"].delete()
    local_db[limits.RUNS_TABLE].delete()
    owner = _account(local_db, "Member")
    _seed(local_db, owner)
    started: list[str] = []
    gate = asyncio.Event()

    run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
    quota = limits.run_quota(local_db, "user", owner, now_utc())
    assert quota.used == 1
    assert quota.limit == limits.DEFAULT_MEMBER_RUNS


def test_a_member_over_the_run_quota_is_postponed_not_disabled(local_db):
    local_db["devii_tasks"].delete()
    local_db[limits.RUNS_TABLE].delete()
    owner = _account(local_db, "Member")
    uid = _seed(local_db, owner)
    now = now_utc()
    for index in range(limits.DEFAULT_MEMBER_RUNS):
        limits.record_run(
            local_db, "user", owner, generate_uid(), now - timedelta(minutes=index)
        )
    started: list[str] = []
    gate = asyncio.Event()

    run_async(
        _run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)()
    )
    row = local_db["devii_tasks"].find_one(uid=uid)
    assert started == []
    assert row["enabled"]
    assert row["status"] == "pending"
    assert row["last_error"] == REASON_RUN_QUOTA
    assert row["next_run_at"] > to_iso(now)


def test_an_administrator_keeps_running_past_the_member_limit(local_db):
    local_db["devii_tasks"].delete()
    local_db[limits.RUNS_TABLE].delete()
    owner = _account(local_db, "Admin")
    _seed(local_db, owner)
    now = now_utc()
    for index in range(limits.DEFAULT_MEMBER_RUNS + 5):
        limits.record_run(
            local_db, "user", owner, generate_uid(), now - timedelta(minutes=index)
        )
    started: list[str] = []
    gate = asyncio.Event()

    run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
    assert started == [owner]


def test_budget_postpones_the_task(local_db):
    local_db["devii_tasks"].delete()
    local_db[limits.RUNS_TABLE].delete()
    owner = _account(local_db, "Admin")
    uid = _seed(local_db, owner)
    started: list[str] = []
    gate = asyncio.Event()

    run_async(
        _run_scheduler(
            local_db,
            _harness(local_db, started, gate),
            gate=gate,
            budget_exceeded=lambda kind, owner_id: True,
        )()
    )
    row = local_db["devii_tasks"].find_one(uid=uid)
    assert started == []
    assert row["enabled"]
    assert row["last_error"] == REASON_BUDGET


def test_scheduler_retires_a_guest_owned_task(local_db):
    local_db["devii_tasks"].delete()
    uid = _seed(local_db, "guest-cookie", owner_kind="guest")
    started: list[str] = []
    gate = asyncio.Event()
    gate.set()

    run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
    row = local_db["devii_tasks"].find_one(uid=uid)
    assert not row["enabled"]
    assert row["last_error"] == REASON_NOT_A_USER
    assert started == []


def test_scheduler_retires_expired_and_exhausted_tasks_even_while_saturated(local_db):
    local_db["devii_tasks"].delete()
    local_db[limits.RUNS_TABLE].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] = []
    gate = asyncio.Event()

    run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
    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()
    local_db[limits.RUNS_TABLE].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"]


def test_a_failed_run_still_counts_against_the_quota(local_db):
    local_db["devii_tasks"].delete()
    local_db[limits.RUNS_TABLE].delete()
    owner = _account(local_db, "Member")
    _seed(local_db, owner)

    def resolve(row):
        store = TaskStore(local_db, "user", str(row["owner_id"]))

        async def executor(prompt):
            raise RuntimeError("nope")

        return store, executor, lambda kind, task_row, payload: None

    async def run():
        scheduler = GlobalScheduler(local_db, resolve, tick_seconds=0.05)
        scheduler.start()
        await asyncio.sleep(0.3)
        await scheduler.stop()

    run_async(run())
    assert limits.run_quota(local_db, "user", owner, now_utc()).used >= 1