Refactor to flat APIKey model with quota, admin UI, .env config, and Berlin timezone
- Remove User/Quota models; quota fields now live directly on APIKey - Admin UI: login, API key management, settings (Ollama URL/model), proxy info display - .env/.env.example: ADMIN_PASSWORD, PROXY_HOST/PORT, DATABASE_URL, APP_TZ - Admin API runs on 127.0.0.1 only; proxy host/port configurable - API keys support optional expires_at; verified against Europe/Berlin timezone - Daily/monthly quota resets use Europe/Berlin midnight boundary - Fix all tests to use new flat model; add expiry tests Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Oliver Hofmann
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
cfa874a4c3
commit
c8235ec274
@@ -6,68 +6,23 @@ os.environ.setdefault("OLLAMA_URL", "http://127.0.0.1:9999")
|
||||
|
||||
def _setup_db():
|
||||
from database import Base, engine, SessionLocal
|
||||
from models import User, Quota
|
||||
from crud import create_api_key, hash_password
|
||||
from crud import create_api_key
|
||||
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
Base.metadata.create_all(bind=engine)
|
||||
|
||||
db = SessionLocal()
|
||||
|
||||
test_user = User(
|
||||
username="testuser",
|
||||
email="test@example.com",
|
||||
hashed_password=hash_password("test123"),
|
||||
is_active=True,
|
||||
)
|
||||
db.add(test_user)
|
||||
db.commit()
|
||||
db.refresh(test_user)
|
||||
|
||||
db.add(Quota(
|
||||
user_id=test_user.id,
|
||||
daily_tokens=1_000_000,
|
||||
monthly_tokens=10_000_000,
|
||||
daily_requests=1000,
|
||||
monthly_requests=10000,
|
||||
))
|
||||
db.commit()
|
||||
|
||||
_, raw_key = create_api_key(db, test_user.id, "test-key")
|
||||
_, raw_key = create_api_key(db, name="test-key", daily_tokens=1_000_000,
|
||||
monthly_tokens=10_000_000, daily_requests=1000,
|
||||
monthly_requests=10000)
|
||||
os.environ["TEST_API_KEY"] = raw_key
|
||||
|
||||
admin_user = User(
|
||||
username="admin",
|
||||
email="admin@example.com",
|
||||
hashed_password=hash_password("admin123"),
|
||||
is_active=True,
|
||||
is_admin=True,
|
||||
)
|
||||
db.add(admin_user)
|
||||
db.commit()
|
||||
db.refresh(admin_user)
|
||||
|
||||
db.add(Quota(
|
||||
user_id=admin_user.id,
|
||||
daily_tokens=10_000_000,
|
||||
monthly_tokens=100_000_000,
|
||||
daily_requests=10000,
|
||||
monthly_requests=100000,
|
||||
))
|
||||
db.commit()
|
||||
|
||||
_, admin_raw_key = create_api_key(db, admin_user.id, "admin-key")
|
||||
os.environ["ADMIN_API_KEY"] = admin_raw_key
|
||||
|
||||
db.close()
|
||||
|
||||
|
||||
def _teardown_db():
|
||||
from database import engine
|
||||
from models import Base
|
||||
from database import engine, Base
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
os.environ.pop("TEST_API_KEY", None)
|
||||
os.environ.pop("ADMIN_API_KEY", None)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
|
||||
+77
-75
@@ -5,33 +5,21 @@ from datetime import datetime, timedelta, timezone
|
||||
os.environ.setdefault("OLLAMA_URL", "http://127.0.0.1:9999")
|
||||
|
||||
from database import Base, engine, SessionLocal
|
||||
from models import User, Quota, Usage
|
||||
from crud import check_and_increment_quota, count_tokens, hash_password
|
||||
from models import APIKey, Usage
|
||||
from crud import check_and_increment_quota, count_tokens, create_api_key, verify_api_key
|
||||
|
||||
|
||||
def make_user_and_quota(db, daily_tokens=None, monthly_tokens=None,
|
||||
daily_requests=None, monthly_requests=None):
|
||||
user = User(
|
||||
username="quotauser",
|
||||
email="quota@example.com",
|
||||
hashed_password=hash_password("pass"),
|
||||
is_active=True,
|
||||
)
|
||||
db.add(user)
|
||||
db.commit()
|
||||
db.refresh(user)
|
||||
|
||||
quota = Quota(
|
||||
user_id=user.id,
|
||||
def make_api_key(db, daily_tokens=None, monthly_tokens=None,
|
||||
daily_requests=None, monthly_requests=None):
|
||||
db_key, _ = create_api_key(
|
||||
db,
|
||||
name="test-key",
|
||||
daily_tokens=daily_tokens,
|
||||
monthly_tokens=monthly_tokens,
|
||||
daily_requests=daily_requests,
|
||||
monthly_requests=monthly_requests,
|
||||
)
|
||||
db.add(quota)
|
||||
db.commit()
|
||||
|
||||
return user.id
|
||||
return db_key.id
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -64,119 +52,133 @@ def test_count_tokens_more_accurate_than_split():
|
||||
# --- check_and_increment_quota ---
|
||||
|
||||
def test_allowed_within_daily_token_limit(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000)
|
||||
assert check_and_increment_quota(db, user_id, tokens=100, requests=1) is True
|
||||
api_key_id = make_api_key(db, daily_tokens=1000)
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=100, requests=1) is True
|
||||
|
||||
def test_denied_when_daily_tokens_exceeded(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=50)
|
||||
assert check_and_increment_quota(db, user_id, tokens=100, requests=1) is False
|
||||
api_key_id = make_api_key(db, daily_tokens=50)
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=100, requests=1) is False
|
||||
|
||||
def test_denied_when_monthly_tokens_exceeded(db):
|
||||
user_id = make_user_and_quota(db, monthly_tokens=50)
|
||||
assert check_and_increment_quota(db, user_id, tokens=100, requests=1) is False
|
||||
api_key_id = make_api_key(db, monthly_tokens=50)
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=100, requests=1) is False
|
||||
|
||||
def test_denied_when_daily_requests_exceeded(db):
|
||||
user_id = make_user_and_quota(db, daily_requests=1)
|
||||
check_and_increment_quota(db, user_id, tokens=0, requests=1)
|
||||
assert check_and_increment_quota(db, user_id, tokens=0, requests=1) is False
|
||||
api_key_id = make_api_key(db, daily_requests=1)
|
||||
check_and_increment_quota(db, api_key_id, tokens=0, requests=1)
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=0, requests=1) is False
|
||||
|
||||
def test_denied_when_monthly_requests_exceeded(db):
|
||||
user_id = make_user_and_quota(db, monthly_requests=1)
|
||||
check_and_increment_quota(db, user_id, tokens=0, requests=1)
|
||||
assert check_and_increment_quota(db, user_id, tokens=0, requests=1) is False
|
||||
api_key_id = make_api_key(db, monthly_requests=1)
|
||||
check_and_increment_quota(db, api_key_id, tokens=0, requests=1)
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=0, requests=1) is False
|
||||
|
||||
def test_increments_both_daily_and_monthly_counters(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000, monthly_tokens=10000,
|
||||
daily_requests=100, monthly_requests=1000)
|
||||
check_and_increment_quota(db, user_id, tokens=50, requests=1)
|
||||
api_key_id = make_api_key(db, daily_tokens=1000, monthly_tokens=10000,
|
||||
daily_requests=100, monthly_requests=1000)
|
||||
check_and_increment_quota(db, api_key_id, tokens=50, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
assert usage.tokens_used_today == 50
|
||||
assert usage.tokens_used_month == 50
|
||||
assert usage.requests_today == 1
|
||||
assert usage.requests_month == 1
|
||||
|
||||
def test_creates_usage_record_on_first_call(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000)
|
||||
assert db.query(Usage).filter(Usage.user_id == user_id).first() is None
|
||||
api_key_id = make_api_key(db, daily_tokens=1000)
|
||||
assert db.query(Usage).filter(Usage.api_key_id == api_key_id).first() is None
|
||||
|
||||
check_and_increment_quota(db, user_id, tokens=10, requests=1)
|
||||
check_and_increment_quota(db, api_key_id, tokens=10, requests=1)
|
||||
|
||||
assert db.query(Usage).filter(Usage.user_id == user_id).first() is not None
|
||||
assert db.query(Usage).filter(Usage.api_key_id == api_key_id).first() is not None
|
||||
|
||||
def test_no_quota_allows_any_request(db):
|
||||
user_id = make_user_and_quota(db) # all limits None
|
||||
assert check_and_increment_quota(db, user_id, tokens=999999, requests=9999) is True
|
||||
api_key_id = make_api_key(db) # all limits None
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=999999, requests=9999) is True
|
||||
|
||||
def test_cumulative_usage_across_calls(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=200)
|
||||
check_and_increment_quota(db, user_id, tokens=100, requests=1)
|
||||
check_and_increment_quota(db, user_id, tokens=99, requests=1)
|
||||
# 199 used, 1 remaining – exactly 1 more token should pass
|
||||
assert check_and_increment_quota(db, user_id, tokens=1, requests=1) is True
|
||||
# Now 200 used – next request must fail
|
||||
assert check_and_increment_quota(db, user_id, tokens=1, requests=1) is False
|
||||
api_key_id = make_api_key(db, daily_tokens=200)
|
||||
check_and_increment_quota(db, api_key_id, tokens=100, requests=1)
|
||||
check_and_increment_quota(db, api_key_id, tokens=99, requests=1)
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=1, requests=1) is True
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=1, requests=1) is False
|
||||
|
||||
|
||||
# --- Reset logic ---
|
||||
|
||||
def test_daily_reset_restores_access(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=100)
|
||||
check_and_increment_quota(db, user_id, tokens=90, requests=1)
|
||||
api_key_id = make_api_key(db, daily_tokens=100)
|
||||
check_and_increment_quota(db, api_key_id, tokens=90, requests=1)
|
||||
|
||||
# Backdate daily_reset_at to yesterday
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
usage.daily_reset_at = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
db.commit()
|
||||
|
||||
# Should pass again after reset
|
||||
assert check_and_increment_quota(db, user_id, tokens=90, requests=1) is True
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=90, requests=1) is True
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
assert usage.tokens_used_today == 90
|
||||
|
||||
def test_daily_reset_does_not_affect_monthly_counter(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000, monthly_tokens=10000)
|
||||
check_and_increment_quota(db, user_id, tokens=50, requests=1)
|
||||
api_key_id = make_api_key(db, daily_tokens=1000, monthly_tokens=10000)
|
||||
check_and_increment_quota(db, api_key_id, tokens=50, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
usage.daily_reset_at = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
db.commit()
|
||||
|
||||
check_and_increment_quota(db, user_id, tokens=50, requests=1)
|
||||
check_and_increment_quota(db, api_key_id, tokens=50, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
assert usage.tokens_used_today == 50
|
||||
assert usage.tokens_used_month == 100 # cumulative across days
|
||||
assert usage.tokens_used_month == 100
|
||||
|
||||
def test_monthly_reset_restores_access(db):
|
||||
user_id = make_user_and_quota(db, monthly_tokens=100)
|
||||
check_and_increment_quota(db, user_id, tokens=90, requests=1)
|
||||
api_key_id = make_api_key(db, monthly_tokens=100)
|
||||
check_and_increment_quota(db, api_key_id, tokens=90, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
usage.monthly_reset_at = datetime.now(timezone.utc) - timedelta(days=32)
|
||||
db.commit()
|
||||
|
||||
assert check_and_increment_quota(db, user_id, tokens=90, requests=1) is True
|
||||
assert check_and_increment_quota(db, api_key_id, tokens=90, requests=1) is True
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
assert usage.tokens_used_month == 90
|
||||
|
||||
def test_failed_quota_check_still_commits_reset(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=100, daily_requests=5)
|
||||
check_and_increment_quota(db, user_id, tokens=80, requests=1)
|
||||
api_key_id = make_api_key(db, daily_tokens=100, daily_requests=5)
|
||||
check_and_increment_quota(db, api_key_id, tokens=80, requests=1)
|
||||
|
||||
# Backdate so a reset fires, but the new request still exceeds the limit
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
usage.daily_reset_at = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
usage.tokens_used_today = 80
|
||||
db.commit()
|
||||
|
||||
# After reset tokens_used_today = 0; 200 tokens exceeds 100 limit
|
||||
result = check_and_increment_quota(db, user_id, tokens=200, requests=1)
|
||||
result = check_and_increment_quota(db, api_key_id, tokens=200, requests=1)
|
||||
assert result is False
|
||||
|
||||
# Reset must still be persisted so the next request sees fresh counters
|
||||
db.expire_all()
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage = db.query(Usage).filter(Usage.api_key_id == api_key_id).first()
|
||||
assert usage.tokens_used_today == 0
|
||||
|
||||
|
||||
# --- verify_api_key expiry ---
|
||||
|
||||
def _create_raw_key(db, expires_at=None):
|
||||
_, raw_key = create_api_key(db, name="expiry-test", expires_at=expires_at)
|
||||
return raw_key
|
||||
|
||||
def test_key_without_expiry_is_valid(db):
|
||||
raw_key = _create_raw_key(db)
|
||||
assert verify_api_key(db, raw_key) is not None
|
||||
|
||||
def test_key_with_future_expiry_is_valid(db):
|
||||
future = datetime.now(timezone.utc) + timedelta(days=30)
|
||||
raw_key = _create_raw_key(db, expires_at=future)
|
||||
assert verify_api_key(db, raw_key) is not None
|
||||
|
||||
def test_key_with_past_expiry_is_rejected(db):
|
||||
past = datetime.now(timezone.utc) - timedelta(seconds=1)
|
||||
raw_key = _create_raw_key(db, expires_at=past)
|
||||
assert verify_api_key(db, raw_key) is None
|
||||
|
||||
Reference in New Issue
Block a user