File size: 2,517 Bytes
6d20eab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import pytest
import tempfile
import time
from app.core.usage_tracker import UsageTracker, Tier, UsageRecord
@pytest.fixture
def tracker():
with tempfile.NamedTemporaryFile(suffix=".db") as tmp:
yield UsageTracker(db_path=tmp.name)
def test_get_or_create_api_key(tracker):
assert tracker.get_or_create_api_key("test_key", Tier.FREE) is True
assert tracker.get_tier("test_key") == Tier.FREE
# Second call should return True without error
assert tracker.get_or_create_api_key("test_key") is True
def test_update_api_key_tier(tracker):
tracker.get_or_create_api_key("test_key", Tier.FREE)
assert tracker.update_api_key_tier("test_key", Tier.PRO) is True
assert tracker.get_tier("test_key") == Tier.PRO
# Non-existent key
assert tracker.update_api_key_tier("nonexistent", Tier.PRO) is False
def test_get_remaining_quota_free(tracker):
tracker.get_or_create_api_key("free_key", Tier.FREE)
# Initially 1000 remaining
remaining = tracker.get_remaining_quota("free_key", Tier.FREE)
assert remaining == 1000
# Simulate usage using the atomic method
record = UsageRecord(
api_key="free_key",
tier=Tier.FREE,
timestamp=time.time(),
endpoint="/test"
)
tracker.increment_usage_sync(record)
remaining = tracker.get_remaining_quota("free_key", Tier.FREE)
assert remaining == 999
def test_get_remaining_quota_enterprise(tracker):
tracker.get_or_create_api_key("ent_key", Tier.ENTERPRISE)
remaining = tracker.get_remaining_quota("ent_key", Tier.ENTERPRISE)
assert remaining is None
def test_increment_usage_sync(tracker):
tracker.get_or_create_api_key("test_key", Tier.FREE)
record = UsageRecord(
api_key="test_key",
tier=Tier.FREE,
timestamp=time.time(),
endpoint="/test",
)
result = tracker.increment_usage_sync(record)
assert result is True
# Check quota decreased
remaining = tracker.get_remaining_quota("test_key", Tier.FREE)
assert remaining == 999
def test_get_audit_logs(tracker):
tracker.get_or_create_api_key("test_key", Tier.FREE)
record = UsageRecord(
api_key="test_key",
tier=Tier.FREE,
timestamp=time.time(),
endpoint="/test",
request_body={"foo": "bar"},
response={"status": "ok"},
)
tracker.increment_usage_sync(record)
logs = tracker.get_audit_logs("test_key", limit=10)
assert len(logs) == 1
assert logs[0]["endpoint"] == "/test"
|