feat: add billing subscription plan api

This commit is contained in:
hj24 2025-12-18 10:33:25 +08:00
parent 03cc3868ef
commit 0bbf8f72b1
2 changed files with 177 additions and 1 deletions

View File

@ -1,7 +1,7 @@
import logging
import os
from collections.abc import Sequence
from typing import Literal
from typing import Literal, TypedDict
import httpx
from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
@ -16,6 +16,13 @@ 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")
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
@ -270,3 +277,34 @@ 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] = {}
chunk_size = 200
for i in range(0, len(tenant_ids), chunk_size):
chunk = tenant_ids[i : i + chunk_size]
resp = cls._send_request("POST", "/subscription/plan/batch", json={"tenant_ids": chunk})
data = resp.get("data", {})
for tenant_id, plan in data.items():
if isinstance(plan, dict) and "plan" in plan and "expiration_date" in plan:
subscription_plan: SubscriptionPlan = {
"plan": str(plan["plan"]),
"expiration_date": int(plan["expiration_date"]),
}
results[tenant_id] = subscription_plan
return results
@classmethod
def get_expired_subscription_cleanup_whitelist(cls) -> Sequence[str]:
resp = cls._send_request("GET", "/subscription/cleanup/whitelist")
data = resp.get("data", [])
return data

View File

@ -1156,6 +1156,144 @@ 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_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
expected_whitelist = [
{
"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",
},
]
mock_send_request.return_value = {"data": expected_whitelist}
# Act
result = BillingService.get_expired_subscription_cleanup_whitelist()
# Assert
assert result == expected_whitelist
assert len(result) == 2
assert result[0]["tenant_id"] == "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6"
assert result[0]["contact"] == "example@dify.ai"
assert result[1]["tenant_id"] == "tenant-2"
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.