"""用量记录 + 配额管理 + 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