149 lines
5.2 KiB
Python
149 lines
5.2 KiB
Python
"""用量记录 + 配额管理 + 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
|