opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-24 23:22:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(api): use billing quota for credit pool (#38028)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
zyssyz123
2026-07-08 07:10:18 +00:00
committed by GitHub
co-authored by autofix-ci[bot]
parent 5a210cf03e
commit 5d61318860
4 changed files with 512 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/unit_tests/services/test_billing_service.py
+120 -6
View File
@@ -73,6 +73,23 @@ class TestBillingServiceSendRequest:
assert call_args[1]["headers"]["Billing-Api-Secret-Key"] == "test-secret-key"
assert call_args[1]["headers"]["Content-Type"] == "application/json"
def test_send_request_with_base_url_override(self, mock_httpx_request, mock_billing_config):
"""Quota APIs can use the new billing service without changing legacy billing calls."""
# Arrange
expected_response = {"result": "success"}
mock_response = MagicMock()
mock_response.status_code = httpx.codes.OK
mock_response.json.return_value = expected_response
mock_httpx_request.return_value = mock_response
# Act
result = BillingService._send_request("GET", "/quota/balance", base_url="https://quota.example.com")
# Assert
assert result == expected_response
call_args = mock_httpx_request.call_args
assert call_args[0][1] == "https://quota.example.com/quota/balance"
@pytest.mark.parametrize(
"status_code", [httpx.codes.NOT_FOUND, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.BAD_REQUEST]
)
@@ -393,6 +410,20 @@ class TestBillingServiceSubscriptionInfo:
params={"tenant_id": tenant_id},
)
def test_quota_get_balance_uses_quota_request(self):
tenant_id = "tenant-123"
with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request:
mock_send_quota_request.return_value = {"quota": "200", "usage": "6", "available": "194", "reserved": "0"}
result = BillingService.quota_get_balance(tenant_id, "credit_pool", bucket="trial")
assert result == {"quota": 200, "usage": 6, "available": 194, "reserved": 0}
mock_send_quota_request.assert_called_once_with(
"GET",
"/quota/balance",
params={"tenant_id": tenant_id, "feature_key": "credit_pool", "bucket": "trial"},
)
def test_get_knowledge_rate_limit_with_defaults(self, mock_send_request):
"""Test knowledge rate limit retrieval with default values."""
# Arrange
@@ -518,19 +549,20 @@ class TestBillingServiceUsageCalculation:
assert result == expected_response
mock_send_request.assert_called_once_with("GET", "/tenant-feature-usage/info", params={"tenant_id": tenant_id})
def test_get_quota_info(self, mock_send_request):
def test_get_quota_info(self):
"""Test retrieval of quota info from new endpoint."""
# Arrange
tenant_id = "tenant-123"
expected_response = {"trigger_event": {"limit": 100, "usage": 30}, "api_rate_limit": {"limit": -1, "usage": 0}}
mock_send_request.return_value = expected_response
with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request:
mock_send_quota_request.return_value = expected_response
# Act
result = BillingService.get_quota_info(tenant_id)
# Act
result = BillingService.get_quota_info(tenant_id)
# Assert
assert result == expected_response
mock_send_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id})
mock_send_quota_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id})
def test_update_tenant_feature_plan_usage_positive_delta(self, mock_send_request):
"""Test updating tenant feature usage with positive delta (adding credits)."""
@@ -614,7 +646,7 @@ class TestBillingServiceQuotaOperations:
@pytest.fixture
def mock_send_request(self):
with patch.object(BillingService, "_send_request") as mock:
with patch.object(BillingService, "_send_quota_request") as mock:
yield mock
def test_quota_reserve_success(self, mock_send_request):
@@ -652,6 +684,16 @@ class TestBillingServiceQuotaOperations:
call_json = mock_send_request.call_args[1]["json"]
assert call_json["meta"] == {"source": "webhook"}
def test_quota_reserve_with_bucket(self, mock_send_request):
mock_send_request.return_value = {"reservation_id": "rid-2", "available": 98, "reserved": 1}
BillingService.quota_reserve(
tenant_id="t1", feature_key="credit_pool", request_id="req-2", amount=1, bucket="trial"
)
call_json = mock_send_request.call_args[1]["json"]
assert call_json["bucket"] == "trial"
def test_quota_commit_success(self, mock_send_request):
expected = {"available": 98, "reserved": 0, "refunded": 0}
mock_send_request.return_value = expected
@@ -696,6 +738,20 @@ class TestBillingServiceQuotaOperations:
call_json = mock_send_request.call_args[1]["json"]
assert call_json["meta"] == {"reason": "partial"}
def test_quota_commit_with_bucket(self, mock_send_request):
mock_send_request.return_value = {"available": 97, "reserved": 0, "refunded": 0}
BillingService.quota_commit(
tenant_id="t1",
feature_key="credit_pool",
reservation_id="rid-1",
actual_amount=1,
bucket="paid",
)
call_json = mock_send_request.call_args[1]["json"]
assert call_json["bucket"] == "paid"
def test_quota_release_success(self, mock_send_request):
expected = {"available": 100, "reserved": 0, "released": 1}
mock_send_request.return_value = expected
@@ -720,6 +776,64 @@ class TestBillingServiceQuotaOperations:
assert result["released"] == 1
assert isinstance(result["released"], int)
def test_quota_release_with_bucket(self, mock_send_request):
mock_send_request.return_value = {"available": 100, "reserved": 0, "released": 1}
BillingService.quota_release(tenant_id="t1", feature_key="credit_pool", reservation_id="rid-1", bucket="trial")
call_json = mock_send_request.call_args[1]["json"]
assert call_json["bucket"] == "trial"
def test_quota_consume_capped_success(self, mock_send_request):
mock_send_request.return_value = {
"deducted": "2",
"available": "8",
"reserved": "0",
"quota": "10",
"usage": "2",
}
result = BillingService.quota_consume_capped(
tenant_id="t1",
feature_key="credit_pool",
request_id="req-1",
amount=5,
bucket="paid",
meta={"source": "test"},
)
assert result == {"deducted": 2, "available": 8, "reserved": 0, "quota": 10, "usage": 2}
mock_send_request.assert_called_once_with(
"POST",
"/quota/consume-capped",
json={
"tenant_id": "t1",
"feature_key": "credit_pool",
"request_id": "req-1",
"amount": 5,
"bucket": "paid",
"meta": {"source": "test"},
},
)
def test_send_quota_request_uses_quota_base_url(self):
with (
patch.object(BillingService, "quota_base_url", "https://quota.example.com/v1"),
patch.object(BillingService, "_send_request") as mock_send_request,
):
mock_send_request.return_value = {"ok": True}
result = BillingService._send_quota_request("GET", "/quota/info", params={"tenant_id": "t1"})
assert result == {"ok": True}
mock_send_request.assert_called_once_with(
"GET",
"/quota/info",
json=None,
params={"tenant_id": "t1"},
base_url="https://quota.example.com/v1",
)
def test_get_quota_info_coerces_string_to_int(self, mock_send_request):
"""Test that TypeAdapter coerces string values to int for get_quota_info."""
mock_send_request.return_value = {
api/tests/unit_tests/services/test_credit_pool_service.py
+164 -3
View File
@@ -1,5 +1,6 @@
from collections.abc import Generator
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from unittest.mock import ANY, MagicMock, patch
from uuid import uuid4
import pytest
@@ -13,6 +14,8 @@ from models.enums import ProviderQuotaType
from services.credit_pool_service import (
CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS,
CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS,
FEATURE_KEY_CREDIT_POOL,
CreditPoolBalance,
CreditPoolService,
)
@@ -51,6 +54,12 @@ def _make_redis_lock() -> MagicMock:
return lock
@pytest.fixture(autouse=True)
def _disable_billing_quota_by_default() -> Generator[None, None, None]:
with patch("services.credit_pool_service.dify_config.BILLING_ENABLED", False):
yield
def test_get_pool_uses_provided_session() -> None:
engine, tenant_id, _ = _create_engine_with_pool(quota_limit=10, quota_used=2)
@@ -62,6 +71,13 @@ def test_get_pool_uses_provided_session() -> None:
assert pool.quota_used == 2
def test_credit_pool_balance_unlimited_remaining_and_sufficiency() -> None:
pool = CreditPoolBalance(tenant_id="tenant-1", pool_type="paid", quota_limit=-1, quota_used=999)
assert pool.remaining_credits == -1
assert pool.has_sufficient_credits(10_000)
def test_check_and_deduct_credits_deducts_exact_amount_when_sufficient() -> None:
engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2)
@@ -211,7 +227,7 @@ def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() -
)
redis_lock.acquire.assert_called_once_with(blocking=True)
redis_lock.release.assert_called_once_with()
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL)
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="trial")
def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> None:
@@ -240,7 +256,152 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N
)
redis_lock.acquire.assert_called_once_with(blocking=True)
redis_lock.release.assert_called_once_with()
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID)
get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="paid")
def test_get_pool_uses_billing_quota_balance_when_enabled() -> None:
tenant_id = "tenant-1"
with (
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
patch("services.billing_service.BillingService.quota_get_balance") as quota_get_balance,
):
quota_get_balance.return_value = {"quota": 1000, "usage": 250, "available": 750, "reserved": 0}
pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID)
assert isinstance(pool, CreditPoolBalance)
assert pool.quota_limit == 1000
assert pool.quota_used == 250
assert pool.remaining_credits == 750
quota_get_balance.assert_called_once_with(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket="paid",
)
def test_check_and_deduct_credits_uses_billing_reserve_and_commit_when_enabled() -> None:
tenant_id = "tenant-1"
with (
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
patch("services.billing_service.BillingService.quota_commit") as quota_commit,
patch("services.billing_service.BillingService.quota_release") as quota_release,
):
quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3}
result = CreditPoolService.check_and_deduct_credits(
tenant_id=tenant_id,
credits_required=3,
pool_type=ProviderQuotaType.TRIAL,
)
assert result == 3
quota_reserve.assert_called_once_with(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket="trial",
request_id=ANY,
amount=3,
meta={"source": "credit_pool.check_and_deduct"},
)
quota_commit.assert_called_once_with(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket="trial",
reservation_id="reservation-1",
actual_amount=3,
meta={"source": "credit_pool.check_and_deduct"},
)
quota_release.assert_not_called()
def test_check_and_deduct_credits_raises_when_billing_reserve_is_insufficient() -> None:
with (
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
):
quota_reserve.return_value = {"reservation_id": "", "available": 1, "reserved": 0}
with pytest.raises(QuotaExceededError, match="Insufficient credits remaining"):
CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3)
def test_check_and_deduct_credits_releases_billing_reservation_when_commit_fails() -> None:
with (
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")),
patch("services.billing_service.BillingService.quota_release") as quota_release,
):
quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3}
with pytest.raises(RuntimeError, match="commit failed"):
CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3)
quota_release.assert_called_once_with(
tenant_id="tenant-1",
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket="trial",
reservation_id="reservation-1",
)
def test_check_and_deduct_credits_logs_when_billing_release_fails() -> None:
with (
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
patch("services.billing_service.BillingService.quota_reserve") as quota_reserve,
patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")),
patch(
"services.billing_service.BillingService.quota_release", side_effect=RuntimeError("release failed")
) as quota_release,
patch("services.credit_pool_service.logger.warning") as logger_warning,
):
quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3}
with pytest.raises(RuntimeError, match="commit failed"):
CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3)
quota_release.assert_called_once_with(
tenant_id="tenant-1",
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket="trial",
reservation_id="reservation-1",
)
logger_warning.assert_called_once()
assert logger_warning.call_args.args[3] == "reservation-1"
assert logger_warning.call_args.kwargs["exc_info"] is True
def test_deduct_credits_capped_uses_billing_consume_capped_when_enabled() -> None:
tenant_id = "tenant-1"
with (
patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True),
patch("services.billing_service.BillingService.quota_consume_capped") as quota_consume_capped,
):
quota_consume_capped.return_value = {
"deducted": 2,
"available": 0,
"reserved": 0,
"quota": 10,
"usage": 10,
}
result = CreditPoolService.deduct_credits_capped(
tenant_id=tenant_id,
credits_required=5,
pool_type=ProviderQuotaType.PAID,
)
assert result == 2
quota_consume_capped.assert_called_once_with(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket="paid",
request_id=ANY,
amount=5,
meta={"source": "credit_pool.deduct_capped"},
)
@pytest.mark.parametrize(