devplacepy/tests/api/admin/gateway/quota.py

# retoor <retoor@molodetz.nl>

from datetime import datetime, timezone

import requests

from devplacepy.database import get_table, refresh_snapshot
from devplacepy.services.openai_gateway.usage import GATEWAY_LEDGER
from devplacepy.utils import generate_uid
from tests.api.admin.gateway.index import (
    JSON_gateway,
    admin_session,
    member_key,
)
from tests.conftest import BASE_URL

RESETS_URL = f"{BASE_URL}/admin/gateway/quota-resets"

_counter = [0]


def _owner():
    _counter[0] += 1
    return f"gwquota{_counter[0]}-{generate_uid()}"


def _burn(owner_id, app_reference, cost):
    refresh_snapshot()
    get_table(GATEWAY_LEDGER).insert(
        {
            "uid": generate_uid(),
            "owner_kind": "user",
            "owner_id": owner_id,
            "app_reference": app_reference,
            "cost_usd": cost,
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
    )
    refresh_snapshot()


def _spent(owner_id, app_reference):
    from devplacepy.services.openai_gateway import quota

    refresh_snapshot()
    quota._QUOTA_CACHE.clear()
    return quota.spent_24h("user", owner_id, app_reference)


def test_quota_reset_requires_admin(seeded_db):
    assert (
        requests.post(RESETS_URL, headers=JSON_gateway, allow_redirects=False).status_code
        == 401
    )
    key = member_key()
    assert (
        requests.post(
            RESETS_URL,
            headers={**JSON_gateway, "X-API-KEY": key},
            json={},
            allow_redirects=False,
        ).status_code
        == 403
    )


def test_admin_can_reset_a_scoped_spend(seeded_db):
    owner = _owner()
    _burn(owner, "appa", 2.5)
    assert _spent(owner, "appa") == 2.5
    response = admin_session(seeded_db).post(
        RESETS_URL,
        json={"owner_kind": "user", "owner_id": owner, "app_reference": "appa"},
    )
    assert response.status_code == 200
    assert response.json()["ok"] is True
    assert _spent(owner, "appa") == 0.0


def test_the_reset_response_echoes_the_scope(seeded_db):
    owner = _owner()
    _burn(owner, "appa", 1.0)
    payload = admin_session(seeded_db).post(
        RESETS_URL,
        json={"owner_kind": "user", "owner_id": owner, "app_reference": "appa"},
    ).json()["reset"]
    assert payload["owner_kind"] == "user"
    assert payload["owner_id"] == owner
    assert payload["app_reference"] == "appa"
    assert payload["reset_at"]


def test_a_reset_leaves_another_caller_capped(seeded_db):
    first, second = _owner(), _owner()
    _burn(first, "appa", 2.0)
    _burn(second, "appa", 2.0)
    admin_session(seeded_db).post(
        RESETS_URL,
        json={"owner_kind": "user", "owner_id": first, "app_reference": "appa"},
    )
    assert _spent(first, "appa") == 0.0
    assert _spent(second, "appa") == 2.0


def test_a_reset_keeps_the_usage_history(seeded_db):
    owner = _owner()
    _burn(owner, "appa", 2.0)
    before = get_table(GATEWAY_LEDGER).count(owner_id=owner)
    admin_session(seeded_db).post(
        RESETS_URL,
        json={"owner_kind": "user", "owner_id": owner, "app_reference": "appa"},
    )
    refresh_snapshot()
    assert get_table(GATEWAY_LEDGER).count(owner_id=owner) == before


def test_an_unscoped_reset_clears_every_caller(seeded_db):
    owner = _owner()
    _burn(owner, "appa", 2.0)
    assert admin_session(seeded_db).post(RESETS_URL, json={}).status_code == 200
    assert _spent(owner, "appa") == 0.0


def test_an_unknown_owner_kind_is_rejected(seeded_db):
    response = admin_session(seeded_db).post(RESETS_URL, json={"owner_kind": "wizard"})
    assert response.status_code == 400
    assert response.json()["ok"] is False


def test_a_malformed_app_reference_is_rejected(seeded_db):
    response = admin_session(seeded_db).post(
        RESETS_URL, json={"app_reference": "not a valid app!"}
    )
    assert response.status_code == 400


def test_the_reset_is_audited(seeded_db):
    owner = _owner()
    _burn(owner, "appa", 1.0)
    admin = admin_session(seeded_db)
    admin.post(
        RESETS_URL,
        json={"owner_kind": "user", "owner_id": owner, "app_reference": "appa"},
    )
    refresh_snapshot()
    entries = admin.get(
        f"{BASE_URL}/admin/audit-log",
        headers=JSON_gateway,
        params={"event_key": "gateway.quota.reset"},
    ).json()["entries"]
    assert any(entry.get("target_type") == "gateway_quota" for entry in entries)


def test_spend_recorded_after_a_reset_counts_again(seeded_db):
    owner = _owner()
    _burn(owner, "appa", 2.0)
    admin_session(seeded_db).post(
        RESETS_URL,
        json={"owner_kind": "user", "owner_id": owner, "app_reference": "appa"},
    )
    _burn(owner, "appa", 0.5)
    assert _spent(owner, "appa") == 0.5