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,
    _unique_gateway,
)
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


def test_quota_rule_form_pages_require_admin(seeded_db):
    assert (
        requests.get(
            f"{BASE_URL}/admin/gateway/quota-rules/new",
            headers=JSON_gateway,
            allow_redirects=False,
        ).status_code
        == 401
    )
    key = member_key()
    assert (
        requests.get(
            f"{BASE_URL}/admin/gateway/quota-rules/new",
            headers={**JSON_gateway, "X-API-KEY": key},
            allow_redirects=False,
        ).status_code
        == 403
    )


def test_quota_rule_add_edit_delete_reset_via_page(seeded_db):
    admin = admin_session(seeded_db)
    owner = _owner()
    app_ref = _unique_gateway("qrapp").lower()

    new_page = admin.get(f"{BASE_URL}/admin/gateway/quota-rules/new")
    assert new_page.status_code == 200
    assert new_page.json()["is_edit"] is False

    created = admin.post(
        f"{BASE_URL}/admin/gateway/quota-rules/new",
        data={
            "owner_kind": "user",
            "owner_id": owner,
            "app_reference": app_ref,
            "limit_usd": "2.5",
            "is_active": "1",
            "label": "page test",
        },
        allow_redirects=False,
    )
    assert created.status_code == 302
    assert created.headers["location"] == "/admin/gateway?tab=quota"

    rules = admin.get(f"{BASE_URL}/admin/gateway/quota-rules").json()["rules"]
    rule = next(r for r in rules if r["owner_id"] == owner and r["app_reference"] == app_ref)
    uid = rule["uid"]
    assert rule["limit_usd"] == 2.5

    edit_page = admin.get(f"{BASE_URL}/admin/gateway/quota-rules/{uid}/edit")
    assert edit_page.status_code == 200
    assert edit_page.json()["form"]["limit_usd"] == "2.5"

    updated = admin.post(
        f"{BASE_URL}/admin/gateway/quota-rules/{uid}/edit",
        data={
            "owner_kind": "user",
            "owner_id": owner,
            "app_reference": app_ref,
            "limit_usd": "5.0",
            "is_active": "1",
            "label": "page test updated",
        },
        allow_redirects=False,
    )
    assert updated.status_code == 302

    rules = admin.get(f"{BASE_URL}/admin/gateway/quota-rules").json()["rules"]
    rule = next(r for r in rules if r["uid"] == uid)
    assert rule["limit_usd"] == 5.0
    assert rule["label"] == "page test updated"

    _burn(owner, app_ref, 1.0)
    assert _spent(owner, app_ref) == 1.0
    reset_response = admin.post(
        f"{BASE_URL}/admin/gateway/quota-rules/{uid}/reset", allow_redirects=False
    )
    assert reset_response.status_code == 302
    assert _spent(owner, app_ref) == 0.0

    deleted = admin.post(
        f"{BASE_URL}/admin/gateway/quota-rules/{uid}/delete", allow_redirects=False
    )
    assert deleted.status_code == 302
    assert (
        admin.get(f"{BASE_URL}/admin/gateway/quota-rules/{uid}/edit").status_code == 404
    )


def test_quota_rule_page_requires_a_dimension(seeded_db):
    admin = admin_session(seeded_db)
    response = admin.post(
        f"{BASE_URL}/admin/gateway/quota-rules/new",
        data={"limit_usd": "1.0"},
    )
    assert response.status_code == 400
    assert "message" in response.json()["error"]


def test_quota_rule_edit_page_404_for_missing_rule(seeded_db):
    admin = admin_session(seeded_db)
    from devplacepy.utils import generate_uid

    missing = generate_uid()
    assert (
        admin.get(f"{BASE_URL}/admin/gateway/quota-rules/{missing}/edit").status_code
        == 404
    )
    assert (
        admin.post(f"{BASE_URL}/admin/gateway/quota-rules/{missing}/delete").status_code
        == 404
    )
    assert (
        admin.post(f"{BASE_URL}/admin/gateway/quota-rules/{missing}/reset").status_code
        == 404
    )