From bffd67f3a45e613f59fd5e4b4abe0069dbf47317 Mon Sep 17 00:00:00 2001 From: hj24 Date: Thu, 18 Dec 2025 10:33:25 +0800 Subject: [PATCH] feat: add billing subscription plan api --- api/services/billing_service.py | 27 ++++--- .../services/test_billing_service.py | 71 ++++++++++++++++--- 2 files changed, 80 insertions(+), 18 deletions(-) diff --git a/api/services/billing_service.py b/api/services/billing_service.py index a0d45d1487..8bd4e7814e 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -291,15 +291,19 @@ class BillingService: 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 + try: + 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 + except Exception: + logger.exception("Failed to fetch billing info batch for tenants: %s", chunk) + continue return results @@ -307,4 +311,7 @@ class BillingService: def get_expired_subscription_cleanup_whitelist(cls) -> Sequence[str]: resp = cls._send_request("GET", "/subscription/cleanup/whitelist") data = resp.get("data", []) - return data + tenant_whitelist = [] + for item in data: + tenant_whitelist.append(item["tenant_id"]) + return tenant_whitelist diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index fe2b40aaa8..f50f744a75 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -1220,6 +1220,53 @@ class TestBillingServiceSubscriptionOperations: 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 @@ -1250,7 +1297,7 @@ class TestBillingServiceSubscriptionOperations: def test_get_expired_subscription_cleanup_whitelist_success(self, mock_send_request): """Test successful retrieval of expired subscription cleanup whitelist.""" # Arrange - expected_whitelist = [ + api_response = [ { "created_at": "2025-10-16T01:56:17", "tenant_id": "36bd55ec-2ea9-4d75-a9ea-1f26aeb4ffe6", @@ -1267,18 +1314,26 @@ class TestBillingServiceSubscriptionOperations: "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": expected_whitelist} + mock_send_request.return_value = {"data": api_response} # 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" + # 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):