feat: add billing subscription plan api (#29829)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
hj24
2025-12-18 13:11:47 +08:00
committed by GitHub
parent c086aa107c
commit 9a51d2da57
2 changed files with 242 additions and 0 deletions

View File

@@ -1,8 +1,12 @@
import logging
import os
from collections.abc import Sequence
from typing import Literal
import httpx
from pydantic import TypeAdapter
from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
from typing_extensions import TypedDict
from werkzeug.exceptions import InternalServerError
from enums.cloud_plan import CloudPlan
@@ -11,6 +15,15 @@ from extensions.ext_redis import redis_client
from libs.helper import RateLimiter
from models import Account, TenantAccountJoin, TenantAccountRole
logger = logging.getLogger(__name__)
class SubscriptionPlan(TypedDict):
"""Tenant subscriptionplan information."""
plan: str
expiration_date: int
class BillingService:
base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
@@ -239,3 +252,39 @@ class BillingService:
def sync_partner_tenants_bindings(cls, account_id: str, partner_key: str, click_id: str):
payload = {"account_id": account_id, "click_id": click_id}
return cls._send_request("PUT", f"/partners/{partner_key}/tenants", json=payload)
@classmethod
def get_plan_bulk(cls, tenant_ids: Sequence[str]) -> dict[str, SubscriptionPlan]:
"""
Bulk fetch billing subscription plan via billing API.
Payload: {"tenant_ids": ["t1", "t2", ...]} (max 200 per request)
Returns:
Mapping of tenant_id -> {plan: str, expiration_date: int}
"""
results: dict[str, SubscriptionPlan] = {}
subscription_adapter = TypeAdapter(SubscriptionPlan)
chunk_size = 200
for i in range(0, len(tenant_ids), chunk_size):
chunk = tenant_ids[i : i + chunk_size]
try:
resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk})
data = resp.get("data", {})
for tenant_id, plan in data.items():
subscription_plan = subscription_adapter.validate_python(plan)
results[tenant_id] = subscription_plan
except Exception:
logger.exception("Failed to fetch billing info batch for tenants: %s", chunk)
continue
return results
@classmethod
def get_expired_subscription_cleanup_whitelist(cls) -> Sequence[str]:
resp = cls._send_request("GET", "/subscription/cleanup/whitelist")
data = resp.get("data", [])
tenant_whitelist = []
for item in data:
tenant_whitelist.append(item["tenant_id"])
return tenant_whitelist

View File

@@ -1156,6 +1156,199 @@ class TestBillingServiceEdgeCases:
assert "Only team owner or team admin can perform this action" in str(exc_info.value)
class TestBillingServiceSubscriptionOperations:
"""Unit tests for subscription operations in BillingService.
Tests cover:
- Bulk plan retrieval with chunking
- Expired subscription cleanup whitelist retrieval
"""
@pytest.fixture
def mock_send_request(self):
"""Mock _send_request method."""
with patch.object(BillingService, "_send_request") as mock:
yield mock
def test_get_plan_bulk_with_empty_list(self, mock_send_request):
"""Test bulk plan retrieval with empty tenant list."""
# Arrange
tenant_ids = []
# Act
result = BillingService.get_plan_bulk(tenant_ids)
# Assert
assert result == {}
mock_send_request.assert_not_called()
def test_get_plan_bulk_with_chunking(self, mock_send_request):
"""Test bulk plan retrieval with more than 200 tenants (chunking logic)."""
# Arrange - 250 tenants to test chunking (chunk_size = 200)
tenant_ids = [f"tenant-{i}" for i in range(250)]
# First chunk: tenants 0-199
first_chunk_response = {
"data": {f"tenant-{i}": {"plan": "sandbox", "expiration_date": 1735689600} for i in range(200)}
}
# Second chunk: tenants 200-249
second_chunk_response = {
"data": {f"tenant-{i}": {"plan": "professional", "expiration_date": 1767225600} for i in range(200, 250)}
}
mock_send_request.side_effect = [first_chunk_response, second_chunk_response]
# Act
result = BillingService.get_plan_bulk(tenant_ids)
# Assert
assert len(result) == 250
assert result["tenant-0"]["plan"] == "sandbox"
assert result["tenant-199"]["plan"] == "sandbox"
assert result["tenant-200"]["plan"] == "professional"
assert result["tenant-249"]["plan"] == "professional"
assert mock_send_request.call_count == 2
# Verify first chunk call
first_call = mock_send_request.call_args_list[0]
assert first_call[0][0] == "POST"
assert first_call[0][1] == "/subscription/plan/batch"
assert len(first_call[1]["json"]["tenant_ids"]) == 200
# Verify second chunk call
second_call = mock_send_request.call_args_list[1]
assert len(second_call[1]["json"]["tenant_ids"]) == 50
def test_get_plan_bulk_with_partial_batch_failure(self, mock_send_request):
"""Test bulk plan retrieval when one batch fails but others succeed."""
# Arrange - 250 tenants, second batch will fail
tenant_ids = [f"tenant-{i}" for i in range(250)]
# First chunk succeeds
first_chunk_response = {
"data": {f"tenant-{i}": {"plan": "sandbox", "expiration_date": 1735689600} for i in range(200)}
}
# Second chunk fails - need to create a mock that raises when called
def side_effect_func(*args, **kwargs):
if mock_send_request.call_count == 1:
return first_chunk_response
else:
raise ValueError("API error")
mock_send_request.side_effect = side_effect_func
# Act
result = BillingService.get_plan_bulk(tenant_ids)
# Assert - should only have data from first batch
assert len(result) == 200
assert result["tenant-0"]["plan"] == "sandbox"
assert result["tenant-199"]["plan"] == "sandbox"
assert "tenant-200" not in result
assert mock_send_request.call_count == 2
def test_get_plan_bulk_with_all_batches_failing(self, mock_send_request):
"""Test bulk plan retrieval when all batches fail."""
# Arrange
tenant_ids = [f"tenant-{i}" for i in range(250)]
# All chunks fail
def side_effect_func(*args, **kwargs):
raise ValueError("API error")
mock_send_request.side_effect = side_effect_func
# Act
result = BillingService.get_plan_bulk(tenant_ids)
# Assert - should return empty dict
assert result == {}
assert mock_send_request.call_count == 2
def test_get_plan_bulk_with_exactly_200_tenants(self, mock_send_request):
"""Test bulk plan retrieval with exactly 200 tenants (boundary condition)."""
# Arrange
tenant_ids = [f"tenant-{i}" for i in range(200)]
mock_send_request.return_value = {
"data": {f"tenant-{i}": {"plan": "sandbox", "expiration_date": 1735689600} for i in range(200)}
}
# Act
result = BillingService.get_plan_bulk(tenant_ids)
# Assert
assert len(result) == 200
assert mock_send_request.call_count == 1
def test_get_plan_bulk_with_empty_data_response(self, mock_send_request):
"""Test bulk plan retrieval with empty data in response."""
# Arrange
tenant_ids = ["tenant-1", "tenant-2"]
mock_send_request.return_value = {"data": {}}
# Act
result = BillingService.get_plan_bulk(tenant_ids)
# Assert
assert result == {}
def test_get_expired_subscription_cleanup_whitelist_success(self, mock_send_request):
"""Test successful retrieval of expired subscription cleanup whitelist."""
# Arrange
api_response = [
{
"created_at": "2025-10-16T01:56:17",
"tenant_id": "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6",
"contact": "example@dify.ai",
"id": "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe5",
"expired_at": "2026-01-01T01:56:17",
"updated_at": "2025-10-16T01:56:17",
},
{
"created_at": "2025-10-16T02:00:00",
"tenant_id": "tenant-2",
"contact": "test@example.com",
"id": "whitelist-id-2",
"expired_at": "2026-02-01T00:00:00",
"updated_at": "2025-10-16T02:00:00",
},
{
"created_at": "2025-10-16T03:00:00",
"tenant_id": "tenant-3",
"contact": "another@example.com",
"id": "whitelist-id-3",
"expired_at": "2026-03-01T00:00:00",
"updated_at": "2025-10-16T03:00:00",
},
]
mock_send_request.return_value = {"data": api_response}
# Act
result = BillingService.get_expired_subscription_cleanup_whitelist()
# Assert - should return only tenant_ids
assert result == ["36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6", "tenant-2", "tenant-3"]
assert len(result) == 3
assert result[0] == "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6"
assert result[1] == "tenant-2"
assert result[2] == "tenant-3"
mock_send_request.assert_called_once_with("GET", "/subscription/cleanup/whitelist")
def test_get_expired_subscription_cleanup_whitelist_empty_list(self, mock_send_request):
"""Test retrieval of empty cleanup whitelist."""
# Arrange
mock_send_request.return_value = {"data": []}
# Act
result = BillingService.get_expired_subscription_cleanup_whitelist()
# Assert
assert result == []
assert len(result) == 0
class TestBillingServiceIntegrationScenarios:
"""Integration-style tests simulating real-world usage scenarios.