""" UsageService 测试 测试用量统计服务的核心功能： - 成本计算 - 钱包准入检查 - 用量统计查询 """ from decimal import Decimal from unittest.mock import MagicMock, patch import pytest from src.services.usage.service import UsageService from src.services.wallet import WalletAccessResult class TestCostCalculation: """测试成本计算""" def test_calculate_cost_basic(self) -> None: """测试基础成本计算""" # 价格：输入 $3/1M, 输出 $15/1M result = UsageService.calculate_cost( input_tokens=1000, output_tokens=500, input_price_per_1m=3.0, output_price_per_1m=15.0, ) ( input_cost, output_cost, cache_creation_cost, cache_read_cost, cache_cost, request_cost, total_cost, ) = result # 1000 tokens * $3 / 1M = $0.003 assert abs(input_cost - 0.003) < 0.0001 # 500 tokens * $15 / 1M = $0.0075 assert abs(output_cost - 0.0075) < 0.0001 # Total = $0.003 + $0.0075 = $0.0105 assert abs(total_cost - 0.0105) < 0.0001 def test_calculate_cost_with_cache(self) -> None: """测试带缓存的成本计算""" result = UsageService.calculate_cost( input_tokens=1000, output_tokens=500, input_price_per_1m=3.0, output_price_per_1m=15.0, cache_creation_input_tokens=200, cache_read_input_tokens=300, cache_creation_price_per_1m=3.75, # 1.25x input price cache_read_price_per_1m=0.3, # 0.1x input price ) ( input_cost, output_cost, cache_creation_cost, cache_read_cost, cache_cost, request_cost, total_cost, ) = result # 验证缓存成本被计算 assert cache_creation_cost > 0 assert cache_read_cost > 0 assert cache_cost == cache_creation_cost + cache_read_cost def test_calculate_cost_with_request_price(self) -> None: """测试按次计费""" result = UsageService.calculate_cost( input_tokens=1000, output_tokens=500, input_price_per_1m=3.0, output_price_per_1m=15.0, price_per_request=0.01, ) ( input_cost, output_cost, cache_creation_cost, cache_read_cost, cache_cost, request_cost, total_cost, ) = result assert request_cost == 0.01 # Total 包含 request_cost assert total_cost == input_cost + output_cost + request_cost def test_calculate_cost_zero_tokens(self) -> None: """测试零 token 的成本计算""" result = UsageService.calculate_cost( input_tokens=0, output_tokens=0, input_price_per_1m=3.0, output_price_per_1m=15.0, ) ( input_cost, output_cost, cache_creation_cost, cache_read_cost, cache_cost, request_cost, total_cost, ) = result assert input_cost == 0 assert output_cost == 0 assert total_cost == 0 class TestBalanceCheck: """测试钱包准入检查""" def test_check_request_balance_sufficient(self) -> None: """测试余额充足""" mock_user = MagicMock() mock_user.role = MagicMock() mock_user.role.value = "user" mock_api_key = MagicMock() mock_api_key.is_standalone = False mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(True, Decimal("70"), "OK"), ): is_ok, message = UsageService.check_request_balance( mock_db, mock_user, api_key=mock_api_key ) assert is_ok is True def test_check_request_balance_details_returns_remaining(self) -> None: """Balance detail helper returns remaining.""" mock_user = MagicMock() mock_user.role = MagicMock() mock_user.role.value = "user" mock_api_key = MagicMock() mock_api_key.is_standalone = False mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult( False, Decimal("12.5"), "\u94b1\u5305\u4f59\u989d\u4e0d\u8db3" ), ): result = UsageService.check_request_balance_details( mock_db, mock_user, api_key=mock_api_key ) assert result.allowed is False assert result.remaining == 12.5 assert "\u4f59\u989d\u4e0d\u8db3" in result.message def test_check_request_balance_details_maps_overdue_message(self) -> None: """欠费状态应映射为对外统一文案。""" mock_user = MagicMock() mock_api_key = MagicMock() mock_api_key.is_standalone = False mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(False, Decimal("-1"), "钱包欠费，请先充值"), ): normal_result = UsageService.check_request_balance_details( mock_db, mock_user, api_key=mock_api_key ) mock_api_key.is_standalone = True with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(False, Decimal("-1"), "钱包欠费，请先充值"), ): standalone_result = UsageService.check_request_balance_details( mock_db, mock_user, api_key=mock_api_key ) assert normal_result.message == "账户欠费，请先充值" assert standalone_result.message == "Key欠费，请先调账或充值" def test_check_request_balance_exceeded(self) -> None: """测试余额耗尽时拦截新请求""" mock_user = MagicMock() mock_user.role = MagicMock() mock_user.role.value = "user" mock_api_key = MagicMock() mock_api_key.is_standalone = False mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(False, Decimal("0"), "钱包余额不足"), ): is_ok, message = UsageService.check_request_balance( mock_db, mock_user, estimated_cost=5.0, api_key=mock_api_key ) assert is_ok is False assert "余额" in message def test_check_request_balance_no_limit(self) -> None: """测试无配额限制（None）""" mock_user = MagicMock() mock_user.role = MagicMock() mock_user.role.value = "user" mock_api_key = MagicMock() mock_api_key.is_standalone = False mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(True, None, "OK"), ): is_ok, message = UsageService.check_request_balance( mock_db, mock_user, api_key=mock_api_key ) assert is_ok is True def test_check_request_balance_admin_bypass(self) -> None: """测试管理员绕过余额检查""" from src.models.database import UserRole mock_user = MagicMock() mock_user.role = UserRole.ADMIN mock_api_key = MagicMock() mock_api_key.is_standalone = False mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(True, None, "OK"), ): is_ok, message = UsageService.check_request_balance( mock_db, mock_user, api_key=mock_api_key ) assert is_ok is True def test_check_standalone_api_key_balance(self) -> None: """测试独立 API Key 余额充足""" mock_user = MagicMock() mock_user.role = MagicMock() mock_user.role.value = "user" mock_api_key = MagicMock() mock_api_key.is_standalone = True mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(True, Decimal("40"), "OK"), ): is_ok, message = UsageService.check_request_balance( mock_db, mock_user, api_key=mock_api_key ) assert is_ok is True def test_check_standalone_api_key_insufficient_balance(self) -> None: """测试独立 API Key 余额耗尽时拦截""" mock_user = MagicMock() mock_user.role = MagicMock() mock_user.role.value = "user" mock_api_key = MagicMock() mock_api_key.is_standalone = True mock_db = MagicMock() with patch( "src.services.wallet.WalletService.check_request_allowed", return_value=WalletAccessResult(False, Decimal("0"), "钱包余额不足"), ): is_ok, message = UsageService.check_request_balance( mock_db, mock_user, estimated_cost=5.0, api_key=mock_api_key ) assert is_ok is False assert "Key余额不足" in message class TestUsageStatistics: """测试用量统计查询 注意：get_usage_summary 方法内部使用了数据库方言特定的日期函数， 需要真实数据库或更复杂的 mock。这里只测试方法存在性。 """ def test_get_usage_summary_exists(self) -> None: """测试 get_usage_summary 方法存在""" assert hasattr(UsageService, "get_usage_summary") assert callable(getattr(UsageService, "get_usage_summary")) class TestHelperMethods: """测试辅助方法""" @pytest.mark.asyncio async def test_get_rate_multiplier_and_free_tier_default(self) -> None: """测试默认费率倍数""" mock_db = MagicMock() # 模拟未找到 provider_api_key mock_db.query.return_value.filter.return_value.first.return_value = None rate_multiplier, is_free_tier = await UsageService._get_rate_multiplier_and_free_tier( mock_db, provider_api_key_id=None, provider_id=None ) assert rate_multiplier == 1.0 assert is_free_tier is False @pytest.mark.asyncio async def test_get_rate_multiplier_from_provider_api_key(self) -> None: """测试从 ProviderAPIKey 获取费率倍数""" mock_provider_api_key = MagicMock() mock_provider_api_key.rate_multipliers = {"claude:chat": 0.8} mock_endpoint = MagicMock() mock_endpoint.provider_id = "provider-123" mock_provider = MagicMock() mock_provider.billing_type = "standard" mock_db = MagicMock() # 第一次查询返回 provider_api_key mock_db.query.return_value.filter.return_value.first.side_effect = [ mock_provider_api_key, mock_endpoint, mock_provider, ] rate_multiplier, is_free_tier = await UsageService._get_rate_multiplier_and_free_tier( mock_db, provider_api_key_id="pak-123", provider_id=None, api_format="claude:chat" ) assert rate_multiplier == 0.8 assert is_free_tier is False class TestUsageStatusUpdate: """测试进行中状态更新对请求头/体的补写能力""" def test_update_usage_status_can_persist_request_and_provider_payloads(self) -> None: usage = MagicMock() usage.status = "pending" usage.provider_name = "pending" usage.billing_status = "pending" usage.finalized_at = None usage.request_headers = None usage.request_body = None usage.provider_request_headers = None usage.provider_request_body = None mock_db = MagicMock() mock_db.query.return_value.filter.return_value.first.return_value = usage client_headers = {"authorization": "Bearer abc", "x-trace-id": "trace-1"} provider_headers = {"authorization": "Bearer upstream", "x-provider": "demo"} client_body = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}]} provider_body = { "model": "upstream-model", "messages": [{"role": "user", "content": "hello"}], } with ( patch( "src.services.system.config.SystemConfigService.should_log_headers", return_value=True, ), patch( "src.services.system.config.SystemConfigService.should_log_body", return_value=True, ), patch( "src.services.system.config.SystemConfigService.mask_sensitive_headers", side_effect=lambda _db, h: {"masked": h}, ), patch( "src.services.system.config.SystemConfigService.truncate_body", side_effect=lambda _db, b, is_request=True: { "truncated": b, "is_request": is_request, }, ), ): updated = UsageService.update_usage_status( db=mock_db, request_id="req-streaming-1", status="streaming", provider="demo-provider", request_headers=client_headers, request_body=client_body, provider_request_headers=provider_headers, provider_request_body=provider_body, ) assert updated is usage assert usage.request_headers == {"masked": client_headers} assert usage.provider_request_headers == {"masked": provider_headers} assert usage.request_body == {"truncated": client_body, "is_request": True} assert usage.provider_request_body == {"truncated": provider_body, "is_request": True} mock_db.commit.assert_called_once()