feat: auto committed

services/ai/tests/test_usage.py

@@ -0,0 +1,148 @@
+"""用量记录 + 配额管理 + Kafka 生产者测试."""
+
+import pytest
+
+from src.ai.errors import AIQuotaExceededError
+from src.ai.usage.kafka_producer import KafkaProducer, UsageEvent
+from src.ai.usage.quota_enforcer import QuotaEnforcer, QuotaStatus
+from src.ai.usage.usage_recorder import UsageRecord, UsageRecorder
+
+
+class TestUsageRecorder:
+    """UsageRecorder 测试（Redis=None 降级模式）."""
+
+    async def test_record_no_redis_skipped(self) -> None:
+        """Redis=None 时跳过记录."""
+        recorder = UsageRecorder(redis=None)
+        record = UsageRecord(
+            user_id="u-1",
+            school_id="s-1",
+            provider="openai",
+            model="gpt-4o",
+            operation="chat",
+            total_tokens=100,
+        )
+        await recorder.record(record)  # 不抛异常即通过
+
+    async def test_get_user_usage_no_redis(self) -> None:
+        """Redis=None 时返回 0."""
+        recorder = UsageRecorder(redis=None)
+        usage = await recorder.get_user_usage("u-1")
+        assert usage == 0
+
+    async def test_get_school_usage_no_redis(self) -> None:
+        """Redis=None 时返回 0."""
+        recorder = UsageRecorder(redis=None)
+        usage = await recorder.get_school_usage("s-1")
+        assert usage == 0
+
+
+class TestQuotaEnforcer:
+    """QuotaEnforcer 测试."""
+
+    async def test_within_budget_allowed(self) -> None:
+        """预算内允许."""
+        recorder = UsageRecorder(redis=None)
+        enforcer = QuotaEnforcer(
+            usage_recorder=recorder,
+            user_monthly_budget=100_000,
+            school_monthly_budget=1_000_000,
+        )
+        statuses = await enforcer.check("u-1", "s-1")
+        assert len(statuses) == 2
+        assert all(s.allowed for s in statuses)
+
+    async def test_exceeds_budget_raises(self) -> None:
+        """超预算抛 AIQuotaExceededError."""
+        recorder = UsageRecorder(redis=None)
+        enforcer = QuotaEnforcer(
+            usage_recorder=recorder,
+            user_monthly_budget=0,  # budget=0 → used(0) < 0 is False
+        )
+        # used=0, budget=0 → allowed = 0 < 0 = False
+        with pytest.raises(AIQuotaExceededError):
+            await enforcer.check("u-1")
+
+    async def test_no_user_id_no_check(self) -> None:
+        """无 user_id 返回空列表."""
+        recorder = UsageRecorder(redis=None)
+        enforcer = QuotaEnforcer(usage_recorder=recorder)
+        statuses = await enforcer.check("")
+        assert statuses == []
+
+    async def test_only_user(self) -> None:
+        """仅检查 user."""
+        recorder = UsageRecorder(redis=None)
+        enforcer = QuotaEnforcer(usage_recorder=recorder)
+        statuses = await enforcer.check("u-1")
+        assert len(statuses) == 1
+        assert statuses[0].scope == "user"
+
+    def test_quota_status_usage_percentage(self) -> None:
+        """usage_percentage 计算."""
+        status = QuotaStatus(allowed=True, scope="user", used=50, budget=200, remaining=150)
+        assert status.usage_percentage == 25.0
+
+    def test_quota_status_zero_budget(self) -> None:
+        """budget=0 时 usage_percentage=0."""
+        status = QuotaStatus(allowed=True, scope="user", used=10, budget=0, remaining=0)
+        assert status.usage_percentage == 0.0
+
+
+class TestUsageEvent:
+    """UsageEvent dataclass 测试."""
+
+    def test_to_dict_fills_defaults(self) -> None:
+        """to_dict 填充 event_id + occurred_at."""
+        event = UsageEvent(user_id="u-1", operation="chat")
+        data = event.to_dict()
+        assert data["event_id"]  # 自动生成
+        assert data["occurred_at"] > 0  # 自动填充
+        assert data["user_id"] == "u-1"
+
+    def test_from_usage_record(self) -> None:
+        """from_usage_record 工厂方法."""
+        event = UsageEvent.from_usage_record(
+            user_id="u-1",
+            school_id="s-1",
+            request_id="r-1",
+            provider="openai",
+            model="gpt-4o",
+            operation="chat",
+            prompt_tokens=10,
+            completion_tokens=20,
+            total_tokens=30,
+            latency_ms=500,
+        )
+        assert event.event_id
+        assert event.aggregate_id == "r-1"
+        assert event.event_type == "AIUsageRecorded"
+        assert event.total_tokens == 30
+
+
+class TestKafkaProducer:
+    """KafkaProducer 测试（降级模式，不连真实 Kafka）."""
+
+    async def test_start_without_kafka_degrades(self) -> None:
+        """Kafka 不可用时降级."""
+        producer = KafkaProducer(bootstrap_servers="localhost:9999")
+        await producer.start()
+        assert producer.is_started is False
+
+    async def test_publish_not_started_skipped(self) -> None:
+        """未启动时 publish 跳过."""
+        producer = KafkaProducer()
+        event = UsageEvent(user_id="u-1")
+        await producer.publish(event)  # 不抛异常
+
+    async def test_stop_when_not_started(self) -> None:
+        """未启动时 stop 不抛异常."""
+        producer = KafkaProducer()
+        await producer.stop()
+
+    async def test_stop_after_failed_start(self) -> None:
+        """启动失败后 stop 不抛异常."""
+        producer = KafkaProducer(bootstrap_servers="localhost:9999")
+        await producer.start()
+        await producer.stop()
+        assert producer.is_started is False