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

# retoor <retoor@molodetz.nl>

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 task_run_scope
from devplacepy.services.devii.tasks.guards import (
    REASON_BUDGET,
    REASON_CREATE_QUOTA,
    REASON_EXPIRED,
    REASON_FAILURES,
    REASON_MAX_RUNS,
    REASON_NESTED,
    REASON_NOT_A_USER,
    REASON_RUN_QUOTA,
    automation_allowed,
    budget_deferral,
    creation_denial,
    retire_reason,
    run_deferral,
)
from devplacepy.services.devii.tasks.limits import reserve_run
from devplacepy.services.devii.tasks.schedule import now_utc, to_iso
from devplacepy.utils import generate_uid


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


def _row(owner_uid, **overrides):
    row = {
        "uid": generate_uid(),
        "owner_kind": "user",
        "owner_id": owner_uid,
        "run_count": 0,
        "created_at": to_iso(now_utc()),
    }
    row.update(overrides)
    return row


def _burn_runs(local_db, owner, count, reference):
    for index in range(count):
        limits.record_run(
            local_db,
            "user",
            owner,
            generate_uid(),
            reference - timedelta(minutes=index),
        )


def test_any_signed_in_account_may_automate(local_db):
    member = _account(local_db, "Member")
    admin = _account(local_db, "Admin")
    assert automation_allowed("user", member) is True
    assert automation_allowed("user", admin) is True
    assert automation_allowed("guest", "guest-cookie") is False
    assert automation_allowed("user", "") is False


def test_member_task_is_no_longer_retired_for_its_role(local_db):
    member = _account(local_db, "Member")
    assert retire_reason(_row(member), now_utc()) is None


def test_guest_owned_task_is_retired(local_db):
    row = _row("cookie", owner_kind="guest")
    assert retire_reason(row, now_utc()) == REASON_NOT_A_USER


def test_expired_task_is_retired(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    row = _row(owner, expires_at=to_iso(now - timedelta(minutes=1)))
    assert retire_reason(row, now) == REASON_EXPIRED


def test_exhausted_task_is_retired(local_db):
    owner = _account(local_db, "Member")
    assert retire_reason(_row(owner, max_runs=5, run_count=5), now_utc()) == REASON_MAX_RUNS


def test_repeated_failures_retire_the_task(local_db):
    owner = _account(local_db, "Member")
    row = _row(owner, failure_count=3)
    assert retire_reason(row, now_utc(), max_failures=3) == REASON_FAILURES
    assert retire_reason(row, now_utc(), max_failures=0) is None


def test_reservation_succeeds_until_the_quota_is_spent(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    granted = [
        reserve_run(local_db, "user", owner, generate_uid(), now) for _ in range(12)
    ]
    assert granted.count(True) == 10
    assert granted.count(False) == 2


def test_a_spent_quota_defers_instead_of_retiring(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    row = _row(owner)
    _burn_runs(local_db, owner, 10, now)
    assert reserve_run(local_db, "user", owner, row["uid"], now) is False
    postponement = run_deferral(local_db, row, now)
    assert postponement.reason == REASON_RUN_QUOTA
    assert postponement.retry_at > now
    assert retire_reason(row, now) is None


def test_administrator_run_quota_is_larger(local_db):
    owner = _account(local_db, "Admin")
    now = now_utc()
    _burn_runs(local_db, owner, 10, now)
    assert reserve_run(local_db, "user", owner, generate_uid(), now) is True


def test_budget_defers_the_task(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    postponement = budget_deferral(
        _row(owner), now, budget_exceeded=lambda kind, uid: True
    )
    assert postponement is not None
    assert postponement.reason == REASON_BUDGET
    assert postponement.retry_at > now
    assert budget_deferral(_row(owner), now, lambda kind, uid: False) is None
    assert budget_deferral(_row(owner), now, None) is None


def test_creation_is_allowed_within_the_quota(local_db):
    owner = _account(local_db, "Member")
    assert creation_denial(local_db, "user", owner, now_utc()) is None


def test_creation_is_denied_over_the_quota(local_db):
    owner = _account(local_db, "Member")
    now = now_utc()
    for _ in range(5):
        local_db["devii_tasks"].insert(
            {
                "uid": generate_uid(),
                "owner_kind": "user",
                "owner_id": owner,
                "created_at": to_iso(now),
                "deleted_at": None,
            }
        )
    denial = creation_denial(local_db, "user", owner, now)
    assert denial is not None
    assert denial.reason == REASON_CREATE_QUOTA
    assert denial.retry_at is not None


def test_creation_inside_a_task_run_is_denied_for_a_member(local_db):
    owner = _account(local_db, "Member")
    with task_run_scope():
        denial = creation_denial(local_db, "user", owner, now_utc())
    assert denial is not None
    assert denial.reason == REASON_NESTED


def test_creation_inside_a_task_run_is_allowed_for_an_administrator(local_db):
    owner = _account(local_db, "Admin")
    with task_run_scope():
        assert creation_denial(local_db, "user", owner, now_utc()) is None


def test_guest_creation_is_denied(local_db):
    denial = creation_denial(local_db, "guest", "cookie", now_utc())
    assert denial is not None
    assert denial.reason == REASON_NOT_A_USER