79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
|
# 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
|
||
|
|
from tests.conftest import BASE_URL
|
||
|
|
|
||
|
|
RESET_ALL_URL = f"{BASE_URL}/admin/ai-quota/reset-all"
|
||
|
|
|
||
|
|
_counter = [0]
|
||
|
|
|
||
|
|
|
||
|
|
def _owner():
|
||
|
|
_counter[0] += 1
|
||
|
|
return f"resetall{_counter[0]}-{generate_uid()}"
|
||
|
|
|
||
|
|
|
||
|
|
def _burn(owner_id, cost):
|
||
|
|
refresh_snapshot()
|
||
|
|
get_table(GATEWAY_LEDGER).insert(
|
||
|
|
{
|
||
|
|
"uid": generate_uid(),
|
||
|
|
"owner_kind": "user",
|
||
|
|
"owner_id": owner_id,
|
||
|
|
"app_reference": "typosaurus",
|
||
|
|
"cost_usd": cost,
|
||
|
|
"created_at": datetime.now(timezone.utc).isoformat(),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
refresh_snapshot()
|
||
|
|
|
||
|
|
|
||
|
|
def _spent(owner_id):
|
||
|
|
from devplacepy.services.openai_gateway import quota
|
||
|
|
|
||
|
|
refresh_snapshot()
|
||
|
|
quota._QUOTA_CACHE.clear()
|
||
|
|
return quota.spent_24h("user", owner_id, "typosaurus")
|
||
|
|
|
||
|
|
|
||
|
|
def test_reset_all_clears_the_gateway_spend(seeded_db):
|
||
|
|
owner = _owner()
|
||
|
|
_burn(owner, 3.0)
|
||
|
|
assert _spent(owner) == 3.0
|
||
|
|
response = admin_session(seeded_db).post(RESET_ALL_URL, allow_redirects=False)
|
||
|
|
assert response.status_code == 200
|
||
|
|
assert _spent(owner) == 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def test_reset_all_keeps_the_gateway_usage_history(seeded_db):
|
||
|
|
owner = _owner()
|
||
|
|
_burn(owner, 3.0)
|
||
|
|
before = get_table(GATEWAY_LEDGER).count(owner_id=owner)
|
||
|
|
admin_session(seeded_db).post(RESET_ALL_URL, allow_redirects=False)
|
||
|
|
refresh_snapshot()
|
||
|
|
assert get_table(GATEWAY_LEDGER).count(owner_id=owner) == before
|
||
|
|
|
||
|
|
|
||
|
|
def test_gateway_spend_after_reset_all_counts_again(seeded_db):
|
||
|
|
owner = _owner()
|
||
|
|
_burn(owner, 3.0)
|
||
|
|
admin_session(seeded_db).post(RESET_ALL_URL, allow_redirects=False)
|
||
|
|
_burn(owner, 0.25)
|
||
|
|
assert _spent(owner) == 0.25
|
||
|
|
|
||
|
|
|
||
|
|
def test_reset_all_requires_admin(seeded_db):
|
||
|
|
assert (
|
||
|
|
requests.post(
|
||
|
|
RESET_ALL_URL, headers=JSON_gateway, allow_redirects=False
|
||
|
|
).status_code
|
||
|
|
== 401
|
||
|
|
)
|