jprdonnelly/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-03-31 18:00:55 -04:00
Code Issues Packages Projects Releases Wiki Activity

refactor: enhance billing info response handling

Browse Source
This commit is contained in:
hj24
2026-03-31 18:23:32 +08:00
parent 6b0c6d0cde
commit 4653ed7ead
3 changed files with 225 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
api/services/billing_service.py
View File

@@ -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):

5
api/services/feature_service.py
View File

@@ -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:

183
api/tests/unit_tests/services/test_billing_service.py
View File

@@ -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")