diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 9720a8e7c3..75dd3519ad 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -55,6 +55,27 @@ _quota_commit_adapter = TypeAdapter(QuotaCommitResult) _quota_release_adapter = TypeAdapter(QuotaReleaseResult) +class _TenantFeatureQuota(TypedDict): + usage: int + limit: int + reset_date: NotRequired[int] + + +class TenantFeatureQuotaInfo(TypedDict): + """Response of /quota/info. + + NOTE (hj24): + - Same convention as BillingInfo: billing may return int fields as str, + always keep non-strict mode to auto-coerce. + """ + + trigger_event: _TenantFeatureQuota + api_rate_limit: _TenantFeatureQuota + + +_tenant_feature_quota_info_adapter = TypeAdapter(TenantFeatureQuotaInfo) + + class _BillingQuota(TypedDict): size: int limit: int @@ -178,9 +199,11 @@ class BillingService: return usage_info @classmethod - def get_quota_info(cls, tenant_id: str): + def get_quota_info(cls, tenant_id: str) -> TenantFeatureQuotaInfo: params = {"tenant_id": tenant_id} - return cls._send_request("GET", "/quota/info", params=params) + return _tenant_feature_quota_info_adapter.validate_python( + cls._send_request("GET", "/quota/info", params=params) + ) @classmethod def quota_reserve( diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index 09806f608e..34f718ba02 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -640,6 +640,38 @@ class TestBillingServiceQuotaOperations: assert result["released"] == 1 assert isinstance(result["released"], int) + def test_get_quota_info_coerces_string_to_int(self, mock_send_request): + """Test that TypeAdapter coerces string values to int for get_quota_info.""" + mock_send_request.return_value = { + "trigger_event": {"usage": "42", "limit": "3000", "reset_date": "1700000000"}, + "api_rate_limit": {"usage": "10", "limit": "-1", "reset_date": "-1"}, + } + + result = BillingService.get_quota_info("t1") + + assert result["trigger_event"]["usage"] == 42 + assert isinstance(result["trigger_event"]["usage"], int) + assert result["trigger_event"]["limit"] == 3000 + assert isinstance(result["trigger_event"]["limit"], int) + assert result["trigger_event"]["reset_date"] == 1700000000 + assert isinstance(result["trigger_event"]["reset_date"], int) + assert result["api_rate_limit"]["limit"] == -1 + assert isinstance(result["api_rate_limit"]["limit"], int) + + def test_get_quota_info_accepts_int_values(self, mock_send_request): + """Test that get_quota_info works with native int values.""" + expected = { + "trigger_event": {"usage": 42, "limit": 3000, "reset_date": 1700000000}, + "api_rate_limit": {"usage": 0, "limit": -1}, + } + mock_send_request.return_value = expected + + result = BillingService.get_quota_info("t1") + + assert result["trigger_event"]["usage"] == 42 + assert result["trigger_event"]["limit"] == 3000 + assert result["api_rate_limit"]["limit"] == -1 + class TestBillingServiceRateLimitEnforcement: """Unit tests for rate limit enforcement mechanisms.