first commit

api/tests/test_usage_tracking_service.py

@@ -0,0 +1,460 @@
+"""
+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