File size: 3,922 Bytes
94f31ec d9aaa76 94f31ec | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Tests for token usage API endpoints."""
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import event, text
from app.core.database import engine
from app.core.models import Base
from app.core.models import (
TokenUsageRecord, # noqa: F401 - registers model with Base
)
from app.main import app
@event.listens_for(engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
"""Enable foreign keys for SQLite."""
if "sqlite" in str(engine.url):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
@pytest.fixture(autouse=True)
def setup_test_db():
"""Ensure tables exist and clean up after each test."""
# For SQLite, we need to manually create the table with correct autoincrement
if "sqlite" in str(engine.url):
with engine.connect() as conn:
conn.execute(text("""
CREATE TABLE IF NOT EXISTS token_usage_records (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER,
project_id INTEGER,
operation VARCHAR(100) NOT NULL,
model VARCHAR(100) NOT NULL,
input_tokens INTEGER NOT NULL,
output_tokens INTEGER NOT NULL,
total_tokens INTEGER NOT NULL,
cost_usd FLOAT,
latency_ms INTEGER,
status VARCHAR(20) NOT NULL DEFAULT 'success',
error_message TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL
)
"""))
conn.commit()
else:
Base.metadata.create_all(bind=engine)
yield
# Clean up token_usage_records table after each test
with engine.connect() as conn:
conn.execute(text("DELETE FROM token_usage_records"))
conn.commit()
def test_create_token_usage():
"""Test creating a token usage record."""
client = TestClient(app)
response = client.post(
"/api/token-usage/",
json={
"operation": "prd_generation",
"model": "gpt-4",
"input_tokens": 1500,
"output_tokens": 3000,
"total_tokens": 4500,
"cost_usd": 0.15,
"latency_ms": 2500,
},
)
assert response.status_code == 201
data = response.json()
assert data["operation"] == "prd_generation"
assert data["total_tokens"] == 4500
assert data["cost_usd"] == 0.15
def test_list_token_usage():
"""Test listing token usage records."""
client = TestClient(app)
# Create some records first
for _ in range(3):
client.post(
"/api/token-usage/",
json={
"operation": "test_op",
"model": "test-model",
"input_tokens": 100,
"output_tokens": 200,
"total_tokens": 300,
},
)
response = client.get("/api/token-usage/")
assert response.status_code == 200
data = response.json()
assert len(data) == 3
def test_get_token_usage_summary():
"""Test getting aggregated token usage summary."""
client = TestClient(app)
# Create records
for _ in range(5):
client.post(
"/api/token-usage/",
json={
"operation": "test_op",
"model": "test-model",
"input_tokens": 100,
"output_tokens": 200,
"total_tokens": 300,
"cost_usd": 0.01,
"latency_ms": 1000,
},
)
response = client.get("/api/token-usage/summary")
assert response.status_code == 200
data = response.json()
assert data["total_operations"] == 5
assert data["total_tokens"] == 1500
assert data["total_cost"] == 0.05
|