devplacepy/tests/e2e/devii/quota.py

# retoor <retoor@molodetz.nl>

from datetime import datetime, timezone
from types import SimpleNamespace
from tests.conftest import BASE_URL
from devplacepy.database import get_table, set_setting
from devplacepy.utils import generate_uid
from devplacepy.services.devii.config import (
    effective_daily_limit,
    FIELD_USER_DAILY_USD,
    FIELD_GUEST_DAILY_USD,
    FIELD_ADMIN_DAILY_USD,
)
from devplacepy.services.devii.store import UsageLedger
LEDGER = "devii_usage_ledger"
def _now_iso_devii_quota():
    return datetime.now(timezone.utc).isoformat()
def _seed_ledger(owner_kind, owner_id, cost, n=1):
    table = get_table(LEDGER)
    for _ in range(n):
        table.insert(
            {
                "owner_kind": owner_kind,
                "owner_id": owner_id,
                "created_at": _now_iso_devii_quota(),
                "cost_usd": cost,
            }
        )
def _ensure_devii():
    from devplacepy.services.manager import service_manager
    from devplacepy.services.devii import DeviiService

    if service_manager.get_service("devii") is None:
        service_manager.register(DeviiService())
    return service_manager.get_service("devii")
def _make_user_devii_quota(role="Member"):
    uid = generate_uid()
    username = f"quota_{uid[-12:]}"
    get_table("users").insert(
        {
            "uid": uid,
            "username": username,
            "terms_version": "1",
            "email": f"{username}@t.dev",
            "api_key": generate_uid(),
            "role": role,
            "is_active": True,
        }
    )
    return uid, username


def test_admin_reset_user_quota_endpoint(alice):
    page, _ = alice
    bob = get_table("users").find_one(username="bob_test")
    get_table(LEDGER).insert(
        {
            "owner_kind": "user",
            "owner_id": bob["uid"],
            "created_at": _now_iso_devii_quota(),
            "cost_usd": 0.5,
        }
    )
    resp = page.request.post(f"{BASE_URL}/admin/users/{bob['uid']}/reset-ai-quota")
    assert resp.ok
    assert get_table(LEDGER).count(owner_kind="user", owner_id=bob["uid"]) == 0