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

# retoor <retoor@molodetz.nl>

from datetime import timedelta

from devplacepy.database import invalidate_admins_cache, set_setting
from devplacepy.services.devii.tasks import limits
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
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"lim-{uid[-10:]}",
            "terms_version": "1",
            "role": role,
            "deleted_at": None,
            "created_at": SIGNUP_AFTER_PRIMARY_ADMIN,
        }
    )
    invalidate_admins_cache()
    return uid


def _run_at(local_db, owner, moment):
    limits.record_run(local_db, "user", owner, generate_uid(), moment)


def _task_at(local_db, owner, moment, deleted=None):
    local_db["devii_tasks"].insert(
        {
            "uid": generate_uid(),
            "owner_kind": "user",
            "owner_id": owner,
            "prompt": "work",
            "enabled": True,
            "status": "pending",
            "created_at": to_iso(moment),
            "deleted_at": deleted,
        }
    )


def test_default_limits_follow_the_role(local_db):
    assert limits.create_limit(False) == 5
    assert limits.run_limit(False) == 10
    assert limits.create_limit(True) == 5
    assert limits.run_limit(True) == 100


def test_limits_are_configurable(local_db):
    set_setting(limits.FIELD_MEMBER_RUNS, "3")
    try:
        assert limits.run_limit(False) == 3
    finally:
        set_setting(limits.FIELD_MEMBER_RUNS, str(limits.DEFAULT_MEMBER_RUNS))
    assert limits.run_limit(False) == limits.DEFAULT_MEMBER_RUNS


def test_run_quota_counts_only_the_window(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    for offset in (1, 2, 3):
        _run_at(local_db, owner, now - timedelta(hours=offset))
    _run_at(local_db, owner, now - timedelta(hours=30))
    quota = limits.run_quota(local_db, "user", owner, now)
    assert quota.used == 3
    assert quota.limit == 10
    assert quota.exceeded is False
    assert quota.remaining == 7


def test_run_quota_reports_when_the_next_slot_frees(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    oldest = now - timedelta(hours=20)
    _run_at(local_db, owner, oldest)
    for index in range(9):
        _run_at(local_db, owner, now - timedelta(hours=1, minutes=index))
    quota = limits.run_quota(local_db, "user", owner, now)
    assert quota.used == 10
    assert quota.exceeded is True
    assert quota.remaining == 0
    assert quota.free_at == oldest + timedelta(hours=24)


def test_run_quota_is_per_owner(local_db):
    first = _account(local_db, "Member")
    second = _account(local_db, "Member")
    now = now_utc()
    for _ in range(10):
        _run_at(local_db, first, now)
    assert limits.run_quota(local_db, "user", first, now).exceeded is True
    assert limits.run_quota(local_db, "user", second, now).used == 0


def test_administrators_get_the_larger_run_quota(local_db):
    owner = _account(local_db, "Admin")
    now = now_utc()
    for _ in range(50):
        _run_at(local_db, owner, now)
    quota = limits.run_quota(local_db, "user", owner, now)
    assert quota.limit == 100
    assert quota.exceeded is False


def test_create_quota_counts_tasks_in_the_window(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    for offset in range(4):
        _task_at(local_db, owner, now - timedelta(hours=offset))
    _task_at(local_db, owner, now - timedelta(hours=25))
    quota = limits.create_quota(local_db, "user", owner, now)
    assert quota.used == 4
    assert quota.limit == 5
    assert quota.exceeded is False


def test_create_quota_still_counts_a_deleted_task(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    for _ in range(5):
        _task_at(local_db, owner, now, deleted=to_iso(now))
    quota = limits.create_quota(local_db, "user", owner, now)
    assert quota.used == 5
    assert quota.exceeded is True


def test_zero_limit_means_unlimited(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    set_setting(limits.FIELD_MEMBER_RUNS, "0")
    try:
        for _ in range(30):
            _run_at(local_db, owner, now)
        quota = limits.run_quota(local_db, "user", owner, now)
        assert quota.limit == 0
        assert quota.exceeded is False
        assert quota.remaining == -1
    finally:
        set_setting(limits.FIELD_MEMBER_RUNS, str(limits.DEFAULT_MEMBER_RUNS))


def test_reserve_run_grants_exactly_the_quota(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    granted = [
        limits.reserve_run(local_db, "user", owner, generate_uid(), now)
        for _ in range(15)
    ]
    assert granted.count(True) == limits.DEFAULT_MEMBER_RUNS
    assert granted[limits.DEFAULT_MEMBER_RUNS] is False
    assert limits.run_quota(local_db, "user", owner, now).used == limits.DEFAULT_MEMBER_RUNS


def test_reserve_run_grants_an_administrator_more(local_db):
    owner = _account(local_db, "Admin")
    now = now_utc()
    granted = [
        limits.reserve_run(local_db, "user", owner, generate_uid(), now)
        for _ in range(limits.DEFAULT_MEMBER_RUNS + 5)
    ]
    assert all(granted)


def test_reserve_run_is_unbounded_when_the_limit_is_zero(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    set_setting(limits.FIELD_MEMBER_RUNS, "0")
    try:
        granted = [
            limits.reserve_run(local_db, "user", owner, generate_uid(), now)
            for _ in range(25)
        ]
        assert all(granted)
    finally:
        set_setting(limits.FIELD_MEMBER_RUNS, str(limits.DEFAULT_MEMBER_RUNS))


def test_reserve_run_ignores_rows_outside_the_window(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    for index in range(limits.DEFAULT_MEMBER_RUNS):
        _run_at(local_db, owner, now - timedelta(hours=25, minutes=index))
    assert limits.reserve_run(local_db, "user", owner, generate_uid(), now) is True


def test_prune_drops_rows_outside_the_retention_window(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    _run_at(local_db, owner, now)
    _run_at(local_db, owner, now - timedelta(hours=72))
    removed = limits.prune_runs(local_db, now)
    assert removed >= 1
    assert limits.run_quota(local_db, "user", owner, now).used == 1