71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
import logging
|
|
|
|
from app.services.usage_tracker import UsageTracker
|
|
|
|
|
|
class TestUsageTrackerExceptionHandling:
|
|
"""测试 usage_tracker.py 中的异常处理行为"""
|
|
|
|
def test_usage_tracker_record_handles_exception(self, caplog):
|
|
"""测试 UsageTracker.record 方法对异常的处理"""
|
|
with caplog.at_level(logging.WARNING):
|
|
tracker = UsageTracker()
|
|
|
|
with patch.object(tracker, '_records', side_effect=Exception("Simulated error")):
|
|
try:
|
|
tracker.record(
|
|
user_id="test_user",
|
|
brand_id="test_brand",
|
|
engine_type="chatgpt",
|
|
query="test query",
|
|
input_tokens=100,
|
|
output_tokens=200,
|
|
cost=0.05,
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
def test_usage_tracker_get_summary_handles_empty_records(self):
|
|
"""测试 UsageTracker.get_summary 方法处理空记录"""
|
|
tracker = UsageTracker()
|
|
summary = tracker.get_summary(user_id="test_user")
|
|
|
|
assert summary.total_queries == 0
|
|
assert summary.total_input_tokens == 0
|
|
assert summary.total_output_tokens == 0
|
|
assert summary.total_cost == 0.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_usage_tracker_record_async_requires_session(self):
|
|
"""测试 UsageTracker.record_async 需要有效的 session"""
|
|
tracker = UsageTracker()
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized with AsyncSession"):
|
|
await tracker.record_async(
|
|
user_id="test_user",
|
|
brand_id="test_brand",
|
|
engine_type="chatgpt",
|
|
query="test query",
|
|
input_tokens=100,
|
|
output_tokens=200,
|
|
cost=0.05,
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_usage_tracker_get_summary_async_requires_session(self):
|
|
"""测试 UsageTracker.get_summary_async 需要有效的 session"""
|
|
tracker = UsageTracker()
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized with AsyncSession"):
|
|
await tracker.get_summary_async(user_id="test_user")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_usage_tracker_check_quota_async_requires_session(self):
|
|
"""测试 UsageTracker.check_quota_async 需要有效的 session"""
|
|
tracker = UsageTracker()
|
|
|
|
with pytest.raises(RuntimeError, match="not initialized with AsyncSession"):
|
|
await tracker.check_quota_async(user_id="test_user")
|