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

# retoor <retoor@molodetz.nl>

import pytest

from devplacepy.database import invalidate_admins_cache, set_setting
from devplacepy.services.devii.errors import ToolInputError
from devplacepy.services.devii.tasks import limits
from devplacepy.services.devii.tasks.context import task_run_scope
from devplacepy.services.devii.tasks.controller import TaskController
from devplacepy.services.devii.tasks.guards import (
    REASON_CREATE_QUOTA,
    REASON_NESTED,
    REASON_NOT_A_USER,
    AutomationDenied,
)
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
from devplacepy.services.devii.tasks.store import TaskStore, claim, due_rows
from devplacepy.utils import generate_uid

SIGNUP_AFTER_PRIMARY_ADMIN = "2099-01-01T00:00:00"


def _account(local_db, role):
    uid = generate_uid()
    local_db["users"].insert(
        {
            "uid": uid,
            "username": f"store-{uid[-10:]}",
            "terms_version": "1",
            "role": role,
            "deleted_at": None,
            "created_at": SIGNUP_AFTER_PRIMARY_ADMIN,
        }
    )
    invalidate_admins_cache()
    return uid


def _record(**overrides):
    record = {
        "uid": generate_uid(),
        "prompt": "work",
        "enabled": True,
        "status": "pending",
        "created_at": to_iso(now_utc()),
        "next_run_at": to_iso(now_utc()),
        "run_count": 0,
        "failure_count": 0,
        "kind": "interval",
        "every_seconds": 900,
        "max_runs": 10,
    }
    record.update(overrides)
    return record


def _spend_creation_quota(store, count):
    for _ in range(count):
        store.create(_record())


def test_member_can_create_a_task(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    record = _record()
    store.create(record)
    assert store.get(record["uid"]) is not None


def test_member_creation_stops_at_the_daily_limit(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    _spend_creation_quota(store, limits.DEFAULT_MEMBER_CREATE)
    with pytest.raises(AutomationDenied) as denied:
        store.create(_record())
    assert denied.value.reason == REASON_CREATE_QUOTA
    assert denied.value.retry_at is not None


def test_deleting_a_task_does_not_give_the_creation_slot_back(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    records = [_record() for _ in range(limits.DEFAULT_MEMBER_CREATE)]
    for record in records:
        store.create(record)
    for record in records:
        store.delete(record["uid"])
    with pytest.raises(AutomationDenied):
        store.create(_record())


def test_administrator_creation_stops_at_its_own_limit(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Admin"))
    _spend_creation_quota(store, limits.DEFAULT_ADMIN_CREATE)
    with pytest.raises(AutomationDenied) as denied:
        store.create(_record())
    assert denied.value.reason == REASON_CREATE_QUOTA


def test_creation_limit_is_configurable(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    set_setting(limits.FIELD_MEMBER_CREATE, "1")
    try:
        store.create(_record())
        with pytest.raises(AutomationDenied):
            store.create(_record())
    finally:
        set_setting(limits.FIELD_MEMBER_CREATE, str(limits.DEFAULT_MEMBER_CREATE))


def test_guest_cannot_persist_a_task(local_db):
    store = TaskStore(local_db, "guest", "guest-cookie")
    with pytest.raises(AutomationDenied) as denied:
        store.create(_record())
    assert denied.value.reason == REASON_NOT_A_USER


def test_member_cannot_create_a_task_from_inside_a_task_run(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    with task_run_scope():
        with pytest.raises(AutomationDenied) as denied:
            store.create(_record())
    assert denied.value.reason == REASON_NESTED


def test_administrator_can_create_a_task_from_inside_a_task_run(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Admin"))
    record = _record()
    with task_run_scope():
        store.create(record)
    assert store.get(record["uid"]) is not None


def test_member_cannot_enable_a_task_from_inside_a_task_run(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    record = _record()
    store.create(record)
    store.update(record["uid"], {"enabled": False, "status": "disabled"})
    with task_run_scope():
        with pytest.raises(AutomationDenied) as denied:
            store.update(record["uid"], {"enabled": True})
    assert denied.value.reason == REASON_NESTED


def test_administrator_can_enable_a_task_from_inside_a_task_run(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Admin"))
    record = _record()
    store.create(record)
    store.update(record["uid"], {"enabled": False, "status": "disabled"})
    with task_run_scope():
        store.update(record["uid"], {"enabled": True})
    assert store.get(record["uid"])["enabled"]


def test_disabling_from_inside_a_task_run_is_always_allowed(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    record = _record()
    store.create(record)
    with task_run_scope():
        store.update(record["uid"], {"enabled": False, "status": "disabled"})
    assert not store.get(record["uid"])["enabled"]


def test_local_operator_store_bypasses_every_gate(local_db):
    store = TaskStore(local_db, "user", "cli", operator=True)
    with task_run_scope():
        for _ in range(limits.DEFAULT_MEMBER_CREATE + 3):
            store.create(_record())
    assert store.count_active() >= limits.DEFAULT_MEMBER_CREATE


def test_claim_succeeds_once(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Admin"))
    record = _record()
    store.create(record)
    assert claim(local_db, record["uid"]) is True
    assert claim(local_db, record["uid"]) is False


def test_claim_refuses_a_disabled_task(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Admin"))
    record = _record()
    store.create(record)
    store.update(record["uid"], {"enabled": False, "status": "disabled"})
    assert claim(local_db, record["uid"]) is False


def test_due_rows_skips_future_and_running_tasks(local_db):
    admin = _account(local_db, "Admin")
    store = TaskStore(local_db, "user", admin)
    now = now_utc()
    due = _record(next_run_at=to_iso(now.replace(year=now.year - 1)))
    future = _record(next_run_at=to_iso(now.replace(year=now.year + 1)))
    running = _record(
        status="running", next_run_at=to_iso(now.replace(year=now.year - 1))
    )
    for record in (due, future, running):
        store.create(record)
    found = {row["uid"] for row in due_rows(local_db, to_iso(now), 500)}
    assert due["uid"] in found
    assert future["uid"] not in found
    assert running["uid"] not in found


def test_controller_caps_active_tasks_per_owner(local_db, monkeypatch):
    monkeypatch.setattr(
        "devplacepy.services.devii.tasks.controller.max_active_per_owner", lambda: 2
    )
    store = TaskStore(local_db, "user", _account(local_db, "Admin"))
    controller = TaskController(store)
    for _ in range(2):
        controller.create_task(
            {"prompt": "work", "kind": "interval", "every_seconds": 900}
        )
    with pytest.raises(ToolInputError):
        controller.create_task(
            {"prompt": "work", "kind": "interval", "every_seconds": 900}
        )


def test_controller_reports_when_the_next_creation_slot_frees(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    controller = TaskController(store)
    for _ in range(limits.DEFAULT_MEMBER_CREATE):
        controller.create_task(
            {"prompt": "work", "kind": "interval", "every_seconds": 900}
        )
    with pytest.raises(ToolInputError) as failure:
        controller.create_task(
            {"prompt": "work", "kind": "interval", "every_seconds": 900}
        )
    assert REASON_CREATE_QUOTA in str(failure.value)
    assert "frees up at" in str(failure.value)


def test_controller_refuses_nested_creation_for_a_member(local_db):
    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    controller = TaskController(store)
    with task_run_scope():
        with pytest.raises(ToolInputError) as failure:
            controller.create_task(
                {"prompt": "work", "kind": "interval", "every_seconds": 900}
            )
    assert REASON_NESTED in str(failure.value)


def test_controller_refuses_nested_run_now_for_a_member(local_db):
    import json

    store = TaskStore(local_db, "user", _account(local_db, "Member"))
    controller = TaskController(store)
    created = controller.create_task(
        {"prompt": "work", "kind": "interval", "every_seconds": 900}
    )
    uid = json.loads(created)["task"]["uid"]
    with task_run_scope():
        with pytest.raises(ToolInputError) as failure:
            controller.run_task_now({"uid": uid})
    assert REASON_NESTED in str(failure.value)