devplacepy/quota.py at 5079f40f46dd3291b196c0fbc6441e312edbf438

 # 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