From 4653ed7ead91c7cf82dfaf398f3cba8db62ab3e6 Mon Sep 17 00:00:00 2001 From: hj24 Date: Tue, 31 Mar 2026 18:23:32 +0800 Subject: [PATCH] refactor: enhance billing info response handling --- api/services/billing_service.py | 56 +++++- api/services/feature_service.py | 5 +- .../services/test_billing_service.py | 183 ++++++++++++++++-- 3 files changed, 225 insertions(+), 19 deletions(-) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 70d4ce1ee6..13f9d1fcf4 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -2,7 +2,7 @@ import json import logging import os from collections.abc import Sequence -from typing import Literal +from typing import Literal, NotRequired import httpx from pydantic import TypeAdapter @@ -26,6 +26,56 @@ class SubscriptionPlan(TypedDict): expiration_date: int +class _BillingQuota(TypedDict): + size: int + limit: int + + +class _VectorSpaceQuota(TypedDict): + size: float + limit: int + + +class _KnowledgeRateLimit(TypedDict): + size: NotRequired[int] + limit: int + + +class _BillingSubscription(TypedDict): + plan: str + interval: str + education: bool + + +class BillingInfo(TypedDict): + """Response of /subscription/info. + + NOTE (hj24): + - Fields not listed here (e.g. trigger_event, api_rate_limit) are stripped by TypeAdapter.validate_python() + - To ensure the precision, billing may convert fields like int as str, be careful when use TypeAdapter: + 1. validate_python in non-strict mode will coerce it to the expected type + 2. In strict mode, it will raise ValidationError + 3. To preserve compatibility, always keep non-strict mode here and avoid strict mode + """ + + enabled: bool + subscription: _BillingSubscription + members: _BillingQuota + apps: _BillingQuota + vector_space: _VectorSpaceQuota + knowledge_rate_limit: _KnowledgeRateLimit + documents_upload_quota: _BillingQuota + annotation_quota_limit: _BillingQuota + docs_processing: str + can_replace_logo: bool + model_load_balancing_enabled: bool + knowledge_pipeline_publish_enabled: bool + next_credit_reset_date: NotRequired[int] + + +_billing_info_adapter = TypeAdapter(BillingInfo) + + class BillingService: base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY") @@ -38,11 +88,11 @@ class BillingService: _PLAN_CACHE_TTL = 600 @classmethod - def get_info(cls, tenant_id: str): + def get_info(cls, tenant_id: str) -> BillingInfo: params = {"tenant_id": tenant_id} billing_info = cls._send_request("GET", "/subscription/info", params=params) - return billing_info + return _billing_info_adapter.validate_python(billing_info) @classmethod def get_tenant_feature_plan_usage_info(cls, tenant_id: str): diff --git a/api/services/feature_service.py b/api/services/feature_service.py index f38e1762d1..56d896629c 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -312,7 +312,10 @@ class FeatureService: features.apps.limit = billing_info["apps"]["limit"] if "vector_space" in billing_info: - features.vector_space.size = billing_info["vector_space"]["size"] + # NOTE (hj24): billing API returns vector_space.size as float (e.g. 0.0) + # but LimitationModel.size is int; truncate here for compatibility + features.vector_space.size = int(billing_info["vector_space"]["size"]) + # NOTE END features.vector_space.limit = billing_info["vector_space"]["limit"] if "documents_upload_quota" in billing_info: diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index 316381f0ca..4eb35e50b7 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -290,9 +290,19 @@ class TestBillingServiceSubscriptionInfo: # Arrange tenant_id = "tenant-123" expected_response = { - "subscription_plan": "professional", - "billing_cycle": "monthly", - "status": "active", + "enabled": True, + "subscription": {"plan": "professional", "interval": "month", "education": False}, + "members": {"size": 1, "limit": 50}, + "apps": {"size": 1, "limit": 200}, + "vector_space": {"size": 0.0, "limit": 20480}, + "knowledge_rate_limit": {"limit": 1000}, + "documents_upload_quota": {"size": 0, "limit": 1000}, + "annotation_quota_limit": {"size": 0, "limit": 5000}, + "docs_processing": "top-priority", + "can_replace_logo": True, + "model_load_balancing_enabled": True, + "knowledge_pipeline_publish_enabled": True, + "next_credit_reset_date": 1775952000, } mock_send_request.return_value = expected_response @@ -1009,17 +1019,14 @@ class TestBillingServiceEdgeCases: yield mock def test_get_info_empty_response(self, mock_send_request): - """Test handling of empty billing info response.""" - # Arrange + """Empty response from billing API should raise ValidationError due to missing required fields.""" + from pydantic import ValidationError + tenant_id = "tenant-empty" mock_send_request.return_value = {} - # Act - result = BillingService.get_info(tenant_id) - - # Assert - assert result == {} - mock_send_request.assert_called_once() + with pytest.raises(ValidationError): + BillingService.get_info(tenant_id) def test_update_tenant_feature_plan_usage_zero_delta(self, mock_send_request): """Test updating tenant feature usage with zero delta (no change).""" @@ -1434,12 +1441,21 @@ class TestBillingServiceIntegrationScenarios: # Step 1: Get current billing info mock_send_request.return_value = { - "subscription_plan": "sandbox", - "billing_cycle": "monthly", - "status": "active", + "enabled": True, + "subscription": {"plan": "sandbox", "interval": "", "education": False}, + "members": {"size": 0, "limit": 1}, + "apps": {"size": 0, "limit": 5}, + "vector_space": {"size": 0.0, "limit": 50}, + "knowledge_rate_limit": {"limit": 10}, + "documents_upload_quota": {"size": 0, "limit": 50}, + "annotation_quota_limit": {"size": 0, "limit": 10}, + "docs_processing": "standard", + "can_replace_logo": False, + "model_load_balancing_enabled": False, + "knowledge_pipeline_publish_enabled": False, } current_info = BillingService.get_info(tenant_id) - assert current_info["subscription_plan"] == "sandbox" + assert current_info["subscription"]["plan"] == "sandbox" # Step 2: Get payment link for upgrade mock_send_request.return_value = {"payment_link": "https://payment.example.com/upgrade"} @@ -1553,3 +1569,140 @@ class TestBillingServiceIntegrationScenarios: mock_send_request.return_value = {"result": "success", "activated": True} activate_result = BillingService.EducationIdentity.activate(account, "token-123", "MIT", "student") assert activate_result["activated"] is True + + +class TestBillingServiceSubscriptionInfoDataType: + """Unit tests for data type coercion in BillingService.get_info + + 1. Verifies the get_info returns correct Python types for numeric fields + 2. Ensure the compatibility regardless of what results the upstream billing API returns + """ + + @pytest.fixture + def mock_send_request(self): + with patch.object(BillingService, "_send_request") as mock: + yield mock + + @pytest.fixture + def normal_billing_response(self) -> dict: + return { + "enabled": True, + "subscription": { + "plan": "team", + "interval": "year", + "education": False, + }, + "members": {"size": 10, "limit": 50}, + "apps": {"size": 80, "limit": 200}, + "vector_space": {"size": 5120.75, "limit": 20480}, + "knowledge_rate_limit": {"limit": 1000}, + "documents_upload_quota": {"size": 450, "limit": 1000}, + "annotation_quota_limit": {"size": 1200, "limit": 5000}, + "docs_processing": "top-priority", + "can_replace_logo": True, + "model_load_balancing_enabled": True, + "knowledge_pipeline_publish_enabled": True, + "next_credit_reset_date": 1745971200, + } + + @pytest.fixture + def string_billing_response(self) -> dict: + return { + "enabled": True, + "subscription": { + "plan": "team", + "interval": "year", + "education": False, + }, + "members": {"size": "10", "limit": "50"}, + "apps": {"size": "80", "limit": "200"}, + "vector_space": {"size": "5120.75", "limit": "20480"}, + "knowledge_rate_limit": {"limit": "1000"}, + "documents_upload_quota": {"size": "450", "limit": "1000"}, + "annotation_quota_limit": {"size": "1200", "limit": "5000"}, + "docs_processing": "top-priority", + "can_replace_logo": True, + "model_load_balancing_enabled": True, + "knowledge_pipeline_publish_enabled": True, + "next_credit_reset_date": "1745971200", + } + + @staticmethod + def _assert_billing_info_types(result: dict): + assert isinstance(result["enabled"], bool) + assert isinstance(result["subscription"]["plan"], str) + assert isinstance(result["subscription"]["interval"], str) + assert isinstance(result["subscription"]["education"], bool) + + assert isinstance(result["members"]["size"], int) + assert isinstance(result["members"]["limit"], int) + + assert isinstance(result["apps"]["size"], int) + assert isinstance(result["apps"]["limit"], int) + + assert isinstance(result["vector_space"]["size"], float) + assert isinstance(result["vector_space"]["limit"], int) + + assert isinstance(result["knowledge_rate_limit"]["limit"], int) + + assert isinstance(result["documents_upload_quota"]["size"], int) + assert isinstance(result["documents_upload_quota"]["limit"], int) + + assert isinstance(result["annotation_quota_limit"]["size"], int) + assert isinstance(result["annotation_quota_limit"]["limit"], int) + + assert isinstance(result["docs_processing"], str) + assert isinstance(result["can_replace_logo"], bool) + assert isinstance(result["model_load_balancing_enabled"], bool) + assert isinstance(result["knowledge_pipeline_publish_enabled"], bool) + if "next_credit_reset_date" in result: + assert isinstance(result["next_credit_reset_date"], int) + + def test_get_info_with_normal_types(self, mock_send_request, normal_billing_response): + """When the billing API returns native numeric types, get_info should preserve them.""" + mock_send_request.return_value = normal_billing_response + + result = BillingService.get_info("tenant-type-test") + + self._assert_billing_info_types(result) + mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": "tenant-type-test"}) + + def test_get_info_with_string_types(self, mock_send_request, string_billing_response): + """When the billing API returns numeric values as strings, get_info should coerce them.""" + mock_send_request.return_value = string_billing_response + + result = BillingService.get_info("tenant-type-test") + + self._assert_billing_info_types(result) + mock_send_request.assert_called_once_with("GET", "/subscription/info", params={"tenant_id": "tenant-type-test"}) + + def test_get_info_without_optional_fields(self, mock_send_request, string_billing_response): + """NotRequired fields can be absent without raising.""" + del string_billing_response["next_credit_reset_date"] + mock_send_request.return_value = string_billing_response + + result = BillingService.get_info("tenant-type-test") + + assert "next_credit_reset_date" not in result + self._assert_billing_info_types(result) + + def test_get_info_with_extra_fields(self, mock_send_request, string_billing_response): + """Undefined fields are silently stripped by validate_python.""" + string_billing_response["new_feature"] = "something" + mock_send_request.return_value = string_billing_response + + result = BillingService.get_info("tenant-type-test") + + # extra fields are dropped by TypeAdapter on TypedDict + assert "new_feature" not in result + self._assert_billing_info_types(result) + + def test_get_info_missing_required_field_raises(self, mock_send_request, string_billing_response): + """Missing a required field should raise ValidationError.""" + from pydantic import ValidationError + + del string_billing_response["members"] + mock_send_request.return_value = string_billing_response + + with pytest.raises(ValidationError): + BillingService.get_info("tenant-type-test")