461 lines
15 KiB
Python
461 lines
15 KiB
Python
"""
|
|
Tests for UsageTrackingService.
|
|
|
|
These tests verify:
|
|
- Cost calculation for different models
|
|
- Usage logging functionality
|
|
- Daily and monthly usage aggregation
|
|
- Static helper methods
|
|
"""
|
|
|
|
import pytest
|
|
from datetime import datetime, date, timezone, timedelta
|
|
from unittest.mock import Mock, MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
from app.services.usage_tracking_service import (
|
|
UsageTrackingService,
|
|
MODEL_COSTS,
|
|
DEFAULT_COST
|
|
)
|
|
from app.models.ai_usage import (
|
|
AIUsageLog,
|
|
DailyUsageSummary,
|
|
MonthlyUsageSummary,
|
|
TaskType
|
|
)
|
|
|
|
|
|
class TestAIUsageLogModel:
|
|
"""Tests for the AIUsageLog dataclass."""
|
|
|
|
def test_to_dict(self):
|
|
"""Test conversion to dictionary."""
|
|
log = AIUsageLog(
|
|
log_id="log_123",
|
|
user_id="user_456",
|
|
timestamp=datetime(2025, 11, 21, 10, 30, 0, tzinfo=timezone.utc),
|
|
model="anthropic/claude-3.5-sonnet",
|
|
tokens_input=100,
|
|
tokens_output=350,
|
|
tokens_total=450,
|
|
estimated_cost=0.00555,
|
|
task_type=TaskType.STORY_PROGRESSION,
|
|
session_id="sess_789",
|
|
character_id="char_abc",
|
|
request_duration_ms=1500,
|
|
success=True,
|
|
error_message=None
|
|
)
|
|
|
|
result = log.to_dict()
|
|
|
|
assert result["log_id"] == "log_123"
|
|
assert result["user_id"] == "user_456"
|
|
assert result["model"] == "anthropic/claude-3.5-sonnet"
|
|
assert result["tokens_total"] == 450
|
|
assert result["task_type"] == "story_progression"
|
|
assert result["success"] is True
|
|
|
|
def test_from_dict(self):
|
|
"""Test creation from dictionary."""
|
|
data = {
|
|
"log_id": "log_123",
|
|
"user_id": "user_456",
|
|
"timestamp": "2025-11-21T10:30:00+00:00",
|
|
"model": "anthropic/claude-3.5-sonnet",
|
|
"tokens_input": 100,
|
|
"tokens_output": 350,
|
|
"tokens_total": 450,
|
|
"estimated_cost": 0.00555,
|
|
"task_type": "story_progression",
|
|
"session_id": "sess_789",
|
|
"character_id": "char_abc",
|
|
"request_duration_ms": 1500,
|
|
"success": True,
|
|
"error_message": None
|
|
}
|
|
|
|
log = AIUsageLog.from_dict(data)
|
|
|
|
assert log.log_id == "log_123"
|
|
assert log.user_id == "user_456"
|
|
assert log.task_type == TaskType.STORY_PROGRESSION
|
|
assert log.tokens_total == 450
|
|
|
|
def test_from_dict_with_invalid_task_type(self):
|
|
"""Test handling of invalid task type."""
|
|
data = {
|
|
"log_id": "log_123",
|
|
"user_id": "user_456",
|
|
"timestamp": "2025-11-21T10:30:00+00:00",
|
|
"model": "test-model",
|
|
"tokens_input": 100,
|
|
"tokens_output": 200,
|
|
"tokens_total": 300,
|
|
"estimated_cost": 0.001,
|
|
"task_type": "invalid_type"
|
|
}
|
|
|
|
log = AIUsageLog.from_dict(data)
|
|
|
|
# Should default to GENERAL
|
|
assert log.task_type == TaskType.GENERAL
|
|
|
|
|
|
class TestDailyUsageSummary:
|
|
"""Tests for the DailyUsageSummary dataclass."""
|
|
|
|
def test_to_dict(self):
|
|
"""Test conversion to dictionary."""
|
|
summary = DailyUsageSummary(
|
|
date=date(2025, 11, 21),
|
|
user_id="user_123",
|
|
total_requests=15,
|
|
total_tokens=6750,
|
|
total_input_tokens=2000,
|
|
total_output_tokens=4750,
|
|
estimated_cost=0.45,
|
|
requests_by_task={"story_progression": 10, "combat_narration": 5}
|
|
)
|
|
|
|
result = summary.to_dict()
|
|
|
|
assert result["date"] == "2025-11-21"
|
|
assert result["total_requests"] == 15
|
|
assert result["estimated_cost"] == 0.45
|
|
assert result["requests_by_task"]["story_progression"] == 10
|
|
|
|
|
|
class TestCostCalculation:
|
|
"""Tests for cost calculation functionality."""
|
|
|
|
def test_calculate_cost_llama(self):
|
|
"""Test cost calculation for Llama model."""
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="meta/meta-llama-3-8b-instruct",
|
|
tokens_input=1000,
|
|
tokens_output=1000
|
|
)
|
|
|
|
# Llama: $0.0001 per 1K input + $0.0001 per 1K output
|
|
expected = 0.0001 + 0.0001
|
|
assert abs(cost - expected) < 0.000001
|
|
|
|
def test_calculate_cost_haiku(self):
|
|
"""Test cost calculation for Claude Haiku."""
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="anthropic/claude-3.5-haiku",
|
|
tokens_input=1000,
|
|
tokens_output=1000
|
|
)
|
|
|
|
# Haiku: $0.001 per 1K input + $0.005 per 1K output
|
|
expected = 0.001 + 0.005
|
|
assert abs(cost - expected) < 0.000001
|
|
|
|
def test_calculate_cost_sonnet(self):
|
|
"""Test cost calculation for Claude Sonnet."""
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="anthropic/claude-3.5-sonnet",
|
|
tokens_input=1000,
|
|
tokens_output=1000
|
|
)
|
|
|
|
# Sonnet: $0.003 per 1K input + $0.015 per 1K output
|
|
expected = 0.003 + 0.015
|
|
assert abs(cost - expected) < 0.000001
|
|
|
|
def test_calculate_cost_opus(self):
|
|
"""Test cost calculation for Claude Opus."""
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="anthropic/claude-3-opus",
|
|
tokens_input=1000,
|
|
tokens_output=1000
|
|
)
|
|
|
|
# Opus: $0.015 per 1K input + $0.075 per 1K output
|
|
expected = 0.015 + 0.075
|
|
assert abs(cost - expected) < 0.000001
|
|
|
|
def test_calculate_cost_unknown_model(self):
|
|
"""Test cost calculation for unknown model uses default."""
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="unknown/model",
|
|
tokens_input=1000,
|
|
tokens_output=1000
|
|
)
|
|
|
|
# Default: $0.001 per 1K input + $0.005 per 1K output
|
|
expected = DEFAULT_COST["input"] + DEFAULT_COST["output"]
|
|
assert abs(cost - expected) < 0.000001
|
|
|
|
def test_calculate_cost_fractional_tokens(self):
|
|
"""Test cost calculation with fractional token counts."""
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="anthropic/claude-3.5-sonnet",
|
|
tokens_input=500,
|
|
tokens_output=250
|
|
)
|
|
|
|
# Sonnet: (500/1000 * 0.003) + (250/1000 * 0.015)
|
|
expected = 0.0015 + 0.00375
|
|
assert abs(cost - expected) < 0.000001
|
|
|
|
def test_get_model_cost_info(self):
|
|
"""Test getting cost info for a model."""
|
|
cost_info = UsageTrackingService.get_model_cost_info(
|
|
"anthropic/claude-3.5-sonnet"
|
|
)
|
|
|
|
assert cost_info["input"] == 0.003
|
|
assert cost_info["output"] == 0.015
|
|
|
|
def test_get_model_cost_info_unknown(self):
|
|
"""Test getting cost info for unknown model."""
|
|
cost_info = UsageTrackingService.get_model_cost_info("unknown/model")
|
|
|
|
assert cost_info == DEFAULT_COST
|
|
|
|
|
|
class TestUsageTrackingService:
|
|
"""Tests for UsageTrackingService class."""
|
|
|
|
@pytest.fixture
|
|
def mock_env(self):
|
|
"""Set up mock environment variables."""
|
|
with patch.dict('os.environ', {
|
|
'APPWRITE_ENDPOINT': 'https://cloud.appwrite.io/v1',
|
|
'APPWRITE_PROJECT_ID': 'test_project',
|
|
'APPWRITE_API_KEY': 'test_api_key',
|
|
'APPWRITE_DATABASE_ID': 'test_db'
|
|
}):
|
|
yield
|
|
|
|
@pytest.fixture
|
|
def mock_databases(self):
|
|
"""Create mock Databases service."""
|
|
with patch('app.services.usage_tracking_service.Databases') as mock:
|
|
yield mock
|
|
|
|
@pytest.fixture
|
|
def service(self, mock_env, mock_databases):
|
|
"""Create UsageTrackingService instance with mocked dependencies."""
|
|
service = UsageTrackingService()
|
|
return service
|
|
|
|
def test_init_missing_env(self):
|
|
"""Test initialization fails with missing env vars."""
|
|
with patch.dict('os.environ', {}, clear=True):
|
|
with pytest.raises(ValueError, match="Appwrite configuration incomplete"):
|
|
UsageTrackingService()
|
|
|
|
def test_log_usage_success(self, service):
|
|
"""Test logging usage successfully."""
|
|
# Mock the create_document response
|
|
service.databases.create_document = Mock(return_value={
|
|
"$id": "doc_123"
|
|
})
|
|
|
|
result = service.log_usage(
|
|
user_id="user_123",
|
|
model="anthropic/claude-3.5-sonnet",
|
|
tokens_input=100,
|
|
tokens_output=350,
|
|
task_type=TaskType.STORY_PROGRESSION,
|
|
session_id="sess_789"
|
|
)
|
|
|
|
# Verify result
|
|
assert result.user_id == "user_123"
|
|
assert result.model == "anthropic/claude-3.5-sonnet"
|
|
assert result.tokens_input == 100
|
|
assert result.tokens_output == 350
|
|
assert result.tokens_total == 450
|
|
assert result.task_type == TaskType.STORY_PROGRESSION
|
|
assert result.estimated_cost > 0
|
|
|
|
# Verify Appwrite was called
|
|
service.databases.create_document.assert_called_once()
|
|
call_args = service.databases.create_document.call_args
|
|
assert call_args.kwargs["database_id"] == "test_db"
|
|
assert call_args.kwargs["collection_id"] == "ai_usage_logs"
|
|
|
|
def test_log_usage_with_error(self, service):
|
|
"""Test logging usage when request failed."""
|
|
service.databases.create_document = Mock(return_value={
|
|
"$id": "doc_123"
|
|
})
|
|
|
|
result = service.log_usage(
|
|
user_id="user_123",
|
|
model="anthropic/claude-3.5-sonnet",
|
|
tokens_input=100,
|
|
tokens_output=0,
|
|
task_type=TaskType.STORY_PROGRESSION,
|
|
success=False,
|
|
error_message="API timeout"
|
|
)
|
|
|
|
assert result.success is False
|
|
assert result.error_message == "API timeout"
|
|
assert result.tokens_output == 0
|
|
|
|
def test_get_daily_usage(self, service):
|
|
"""Test getting daily usage summary."""
|
|
# Mock list_rows response
|
|
service.tables_db.list_rows = Mock(return_value={
|
|
"rows": [
|
|
{
|
|
"tokens_input": 100,
|
|
"tokens_output": 300,
|
|
"tokens_total": 400,
|
|
"estimated_cost": 0.005,
|
|
"task_type": "story_progression"
|
|
},
|
|
{
|
|
"tokens_input": 150,
|
|
"tokens_output": 350,
|
|
"tokens_total": 500,
|
|
"estimated_cost": 0.006,
|
|
"task_type": "story_progression"
|
|
},
|
|
{
|
|
"tokens_input": 50,
|
|
"tokens_output": 200,
|
|
"tokens_total": 250,
|
|
"estimated_cost": 0.003,
|
|
"task_type": "combat_narration"
|
|
}
|
|
]
|
|
})
|
|
|
|
result = service.get_daily_usage("user_123", date(2025, 11, 21))
|
|
|
|
assert result.user_id == "user_123"
|
|
assert result.date == date(2025, 11, 21)
|
|
assert result.total_requests == 3
|
|
assert result.total_tokens == 1150
|
|
assert result.total_input_tokens == 300
|
|
assert result.total_output_tokens == 850
|
|
assert abs(result.estimated_cost - 0.014) < 0.0001
|
|
assert result.requests_by_task["story_progression"] == 2
|
|
assert result.requests_by_task["combat_narration"] == 1
|
|
|
|
def test_get_daily_usage_empty(self, service):
|
|
"""Test getting daily usage when no usage exists."""
|
|
service.tables_db.list_rows = Mock(return_value={
|
|
"rows": []
|
|
})
|
|
|
|
result = service.get_daily_usage("user_123", date(2025, 11, 21))
|
|
|
|
assert result.total_requests == 0
|
|
assert result.total_tokens == 0
|
|
assert result.estimated_cost == 0.0
|
|
assert result.requests_by_task == {}
|
|
|
|
def test_get_monthly_cost(self, service):
|
|
"""Test getting monthly cost summary."""
|
|
service.tables_db.list_rows = Mock(return_value={
|
|
"rows": [
|
|
{"tokens_total": 1000, "estimated_cost": 0.01},
|
|
{"tokens_total": 2000, "estimated_cost": 0.02},
|
|
{"tokens_total": 1500, "estimated_cost": 0.015}
|
|
]
|
|
})
|
|
|
|
result = service.get_monthly_cost("user_123", 2025, 11)
|
|
|
|
assert result.year == 2025
|
|
assert result.month == 11
|
|
assert result.user_id == "user_123"
|
|
assert result.total_requests == 3
|
|
assert result.total_tokens == 4500
|
|
assert abs(result.estimated_cost - 0.045) < 0.0001
|
|
|
|
def test_get_monthly_cost_invalid_month(self, service):
|
|
"""Test monthly cost with invalid month raises ValueError."""
|
|
with pytest.raises(ValueError, match="Invalid month"):
|
|
service.get_monthly_cost("user_123", 2025, 13)
|
|
|
|
with pytest.raises(ValueError, match="Invalid month"):
|
|
service.get_monthly_cost("user_123", 2025, 0)
|
|
|
|
def test_get_total_daily_cost(self, service):
|
|
"""Test getting total daily cost across all users."""
|
|
service.tables_db.list_rows = Mock(return_value={
|
|
"rows": [
|
|
{"estimated_cost": 0.10},
|
|
{"estimated_cost": 0.25},
|
|
{"estimated_cost": 0.15}
|
|
]
|
|
})
|
|
|
|
result = service.get_total_daily_cost(date(2025, 11, 21))
|
|
|
|
assert abs(result - 0.50) < 0.0001
|
|
|
|
def test_get_user_request_count_today(self, service):
|
|
"""Test getting user request count for today."""
|
|
service.tables_db.list_rows = Mock(return_value={
|
|
"rows": [
|
|
{"tokens_total": 100, "tokens_input": 30, "tokens_output": 70, "estimated_cost": 0.001, "task_type": "story_progression"},
|
|
{"tokens_total": 200, "tokens_input": 50, "tokens_output": 150, "estimated_cost": 0.002, "task_type": "story_progression"}
|
|
]
|
|
})
|
|
|
|
result = service.get_user_request_count_today("user_123")
|
|
|
|
assert result == 2
|
|
|
|
|
|
class TestCostEstimations:
|
|
"""Tests for realistic cost estimation scenarios."""
|
|
|
|
def test_free_tier_daily_cost(self):
|
|
"""Test estimated daily cost for free tier user with Llama."""
|
|
# 20 requests per day, average 500 total tokens each
|
|
total_input = 20 * 200
|
|
total_output = 20 * 300
|
|
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="meta/meta-llama-3-8b-instruct",
|
|
tokens_input=total_input,
|
|
tokens_output=total_output
|
|
)
|
|
|
|
# Should be very cheap (essentially free)
|
|
assert cost < 0.01
|
|
|
|
def test_premium_tier_daily_cost(self):
|
|
"""Test estimated daily cost for premium tier user with Sonnet."""
|
|
# 100 requests per day, average 1000 total tokens each
|
|
total_input = 100 * 300
|
|
total_output = 100 * 700
|
|
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="anthropic/claude-3.5-sonnet",
|
|
tokens_input=total_input,
|
|
tokens_output=total_output
|
|
)
|
|
|
|
# Should be under $2/day for heavy usage
|
|
assert cost < 2.0
|
|
|
|
def test_elite_tier_monthly_cost(self):
|
|
"""Test estimated monthly cost for elite tier user."""
|
|
# 200 requests per day * 30 days = 6000 requests
|
|
# Average 1500 tokens per request
|
|
total_input = 6000 * 500
|
|
total_output = 6000 * 1000
|
|
|
|
cost = UsageTrackingService.estimate_cost_for_model(
|
|
model="anthropic/claude-4.5-sonnet",
|
|
tokens_input=total_input,
|
|
tokens_output=total_output
|
|
)
|
|
|
|
# Elite tier should be under $100/month even with heavy usage
|
|
assert cost < 100.0
|