opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-21 18:58:35 +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 15:10:18 +08:00 committed by GitHub
parent 5a210cf03e
commit 5d61318860
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 512 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

115
api/services/billing_service.py
View File

@ -50,9 +50,26 @@ class QuotaReleaseResult(TypedDict):
released: int
class QuotaBalanceResult(TypedDict):
available: int
reserved: int
quota: int
usage: int
class QuotaConsumeCappedResult(TypedDict):
deducted: int
available: int
reserved: int
quota: int
usage: int
_quota_reserve_adapter = TypeAdapter(QuotaReserveResult)
_quota_commit_adapter = TypeAdapter(QuotaCommitResult)
_quota_release_adapter = TypeAdapter(QuotaReleaseResult)
_quota_balance_adapter = TypeAdapter(QuotaBalanceResult)
_quota_consume_capped_adapter = TypeAdapter(QuotaConsumeCappedResult)
class _TenantFeatureQuota(TypedDict):
@ -176,6 +193,7 @@ class DismissNotificationDict(TypedDict):
class BillingService:
base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
quota_base_url = os.environ.get("BILLING_QUOTA_API_URL") or base_url
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 4, 60)
@ -215,12 +233,18 @@ class BillingService:
def get_quota_info(cls, tenant_id: str) -> TenantFeatureQuotaInfo:
params = {"tenant_id": tenant_id}
return _tenant_feature_quota_info_adapter.validate_python(
cls._send_request("GET", "/quota/info", params=params)
cls._send_quota_request("GET", "/quota/info", params=params)
)
@classmethod
def quota_reserve(
cls, tenant_id: str, feature_key: str, request_id: str, amount: int = 1, meta: dict | None = None
cls,
tenant_id: str,
feature_key: str,
request_id: str,
amount: int = 1,
meta: dict | None = None,
bucket: str = "",
) -> QuotaReserveResult:
"""Reserve quota before task execution."""
payload: dict = {
@ -229,13 +253,21 @@ class BillingService:
"request_id": request_id,
"amount": amount,
}
if bucket:
payload["bucket"] = bucket
if meta:
payload["meta"] = meta
return _quota_reserve_adapter.validate_python(cls._send_request("POST", "/quota/reserve", json=payload))
return _quota_reserve_adapter.validate_python(cls._send_quota_request("POST", "/quota/reserve", json=payload))
@classmethod
def quota_commit(
cls, tenant_id: str, feature_key: str, reservation_id: str, actual_amount: int, meta: dict | None = None
cls,
tenant_id: str,
feature_key: str,
reservation_id: str,
actual_amount: int,
meta: dict | None = None,
bucket: str = "",
) -> QuotaCommitResult:
"""Commit a reservation with actual consumption."""
payload: dict = {
@ -244,23 +276,57 @@ class BillingService:
"reservation_id": reservation_id,
"actual_amount": actual_amount,
}
if bucket:
payload["bucket"] = bucket
if meta:
payload["meta"] = meta
return _quota_commit_adapter.validate_python(cls._send_request("POST", "/quota/commit", json=payload))
return _quota_commit_adapter.validate_python(cls._send_quota_request("POST", "/quota/commit", json=payload))
@classmethod
def quota_release(cls, tenant_id: str, feature_key: str, reservation_id: str) -> QuotaReleaseResult:
def quota_release(
cls, tenant_id: str, feature_key: str, reservation_id: str, bucket: str = ""
) -> QuotaReleaseResult:
"""Release a reservation (cancel, return frozen quota)."""
return _quota_release_adapter.validate_python(
cls._send_request(
"POST",
"/quota/release",
json={
"tenant_id": tenant_id,
"feature_key": feature_key,
"reservation_id": reservation_id,
},
)
payload = {
"tenant_id": tenant_id,
"feature_key": feature_key,
"reservation_id": reservation_id,
}
if bucket:
payload["bucket"] = bucket
return _quota_release_adapter.validate_python(cls._send_quota_request("POST", "/quota/release", json=payload))
@classmethod
def quota_get_balance(cls, tenant_id: str, feature_key: str, bucket: str = "") -> QuotaBalanceResult:
"""Get quota balance for a feature bucket."""
params = {"tenant_id": tenant_id, "feature_key": feature_key}
if bucket:
params["bucket"] = bucket
return _quota_balance_adapter.validate_python(cls._send_quota_request("GET", "/quota/balance", params=params))
@classmethod
def quota_consume_capped(
cls,
tenant_id: str,
feature_key: str,
request_id: str,
amount: int,
meta: dict | None = None,
bucket: str = "",
) -> QuotaConsumeCappedResult:
"""Consume up to the available quota and return the actual deducted amount."""
payload: dict = {
"tenant_id": tenant_id,
"feature_key": feature_key,
"request_id": request_id,
"amount": amount,
}
if bucket:
payload["bucket"] = bucket
if meta:
payload["meta"] = meta
return _quota_consume_capped_adapter.validate_python(
cls._send_quota_request("POST", "/quota/consume-capped", json=payload)
)
@classmethod
@ -334,6 +400,12 @@ class BillingService:
params = {"tenant_id": tenant_id, "feature_key": feature_key}
return cls._send_request("GET", "/billing/tenant_feature_plan/usage", params=params)
@classmethod
def _send_quota_request(
cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None
):
return cls._send_request(method, endpoint, json=json, params=params, base_url=cls.quota_base_url)
@classmethod
@retry(
wait=wait_fixed(2),
@ -341,10 +413,17 @@ class BillingService:
retry=retry_if_exception_type(httpx.RequestError),
reraise=True,
)
def _send_request(cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None):
def _send_request(
cls,
method: Literal["GET", "POST", "DELETE", "PUT"],
endpoint: str,
json=None,
params=None,
base_url: str | None = None,
):
headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key}
url = f"{cls.base_url}{endpoint}"
url = f"{base_url or cls.base_url}{endpoint}"
response = _http_client.request(method, url, json=json, params=params, headers=headers, follow_redirects=True)
if method == "GET" and response.status_code != httpx.codes.OK:
raise ValueError("Unable to retrieve billing information. Please try again later or contact support.")

143
api/services/credit_pool_service.py
View File

@ -7,6 +7,8 @@ from piling up database transactions while preserving cross-tenant concurrency.
import logging
from collections.abc import Callable
from dataclasses import dataclass
from uuid import uuid4
from sqlalchemy import select
from sqlalchemy.orm import Session
@ -19,11 +21,43 @@ from models.enums import ProviderQuotaType
logger = logging.getLogger(__name__)
FEATURE_KEY_CREDIT_POOL = "credit_pool"
CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS = 10
CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS = 5
@dataclass(frozen=True)
class CreditPoolBalance:
tenant_id: str
pool_type: str
quota_limit: int
quota_used: int
@property
def remaining_credits(self) -> int:
if self.quota_limit == -1:
return -1
return max(0, self.quota_limit - self.quota_used)
def has_sufficient_credits(self, required_credits: int) -> bool:
return self.quota_limit == -1 or self.remaining_credits >= required_credits
class CreditPoolService:
@staticmethod
def _normalize_pool_type(pool_type: str | ProviderQuotaType) -> str:
return pool_type.value if isinstance(pool_type, ProviderQuotaType) else str(pool_type)
@staticmethod
def _use_billing_quota() -> bool:
return bool(dify_config.BILLING_ENABLED)
@staticmethod
def _require_session(session: Session | None) -> Session:
if session is None:
raise ValueError("session is required when billing quota is disabled")
return session
@staticmethod
def _get_tenant_lock_key(tenant_id: str) -> str:
return f"credit_pool:tenant:{tenant_id}:deduct_lock"
@ -77,13 +111,36 @@ class CreditPoolService:
return credit_pool
@classmethod
def get_pool(cls, tenant_id: str, pool_type: str = "trial", *, session: Session) -> TenantCreditPool | None:
def get_pool(
cls,
tenant_id: str,
pool_type: str | ProviderQuotaType = "trial",
*,
session: Session | None = None,
) -> TenantCreditPool | CreditPoolBalance | None:
"""get tenant credit pool"""
normalized_pool_type = cls._normalize_pool_type(pool_type)
if cls._use_billing_quota():
from services.billing_service import BillingService
balance = BillingService.quota_get_balance(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket=normalized_pool_type,
)
return CreditPoolBalance(
tenant_id=tenant_id,
pool_type=normalized_pool_type,
quota_limit=balance["quota"],
quota_used=balance["usage"],
)
session = cls._require_session(session)
return session.scalar(
select(TenantCreditPool)
.where(
TenantCreditPool.tenant_id == tenant_id,
TenantCreditPool.pool_type == pool_type,
TenantCreditPool.pool_type == normalized_pool_type,
)
.limit(1)
)
@ -93,31 +150,77 @@ class CreditPoolService:
cls,
tenant_id: str,
credits_required: int,
pool_type: str = "trial",
pool_type: str | ProviderQuotaType = "trial",
*,
session: Session,
session: Session | None = None,
) -> bool:
"""check if credits are available without deducting"""
pool = cls.get_pool(tenant_id, pool_type, session=session)
if not pool:
return False
return pool.remaining_credits >= credits_required
return pool.has_sufficient_credits(credits_required)
@classmethod
def check_and_deduct_credits(
cls,
tenant_id: str,
credits_required: int,
pool_type: str = "trial",
pool_type: str | ProviderQuotaType = "trial",
*,
session: Session,
session: Session | None = None,
) -> int:
"""Deduct exactly the requested credits or raise without mutating the pool."""
if credits_required <= 0:
return 0
normalized_pool_type = cls._normalize_pool_type(pool_type)
if cls._use_billing_quota():
from services.billing_service import BillingService
request_id = str(uuid4())
result = BillingService.quota_reserve(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket=normalized_pool_type,
request_id=request_id,
amount=credits_required,
meta={"source": "credit_pool.check_and_deduct"},
)
reservation_id = result.get("reservation_id", "")
if not reservation_id:
raise QuotaExceededError("Insufficient credits remaining")
try:
BillingService.quota_commit(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket=normalized_pool_type,
reservation_id=reservation_id,
actual_amount=credits_required,
meta={"source": "credit_pool.check_and_deduct"},
)
except Exception:
try:
BillingService.quota_release(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket=normalized_pool_type,
reservation_id=reservation_id,
)
except Exception:
logger.warning(
"Failed to release reserved credit pool quota, tenant_id=%s, pool_type=%s, reservation_id=%s",
tenant_id,
normalized_pool_type,
reservation_id,
exc_info=True,
)
raise
return credits_required
session = cls._require_session(session)
def deduct() -> int:
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type)
if not pool:
raise QuotaExceededError("Credit pool not found")
@ -144,18 +247,34 @@ class CreditPoolService:
cls,
tenant_id: str,
credits_required: int,
pool_type: str = "trial",
pool_type: str | ProviderQuotaType = "trial",
*,
session: Session,
session: Session | None = None,
) -> int:
"""Deduct up to the available balance and return the actual deducted credits."""
if credits_required <= 0:
return 0
normalized_pool_type = cls._normalize_pool_type(pool_type)
if cls._use_billing_quota():
from services.billing_service import BillingService
result = BillingService.quota_consume_capped(
tenant_id=tenant_id,
feature_key=FEATURE_KEY_CREDIT_POOL,
bucket=normalized_pool_type,
request_id=str(uuid4()),
amount=credits_required,
meta={"source": "credit_pool.deduct_capped"},
)
return result["deducted"]
session = cls._require_session(session)
def deduct() -> int:
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type)
pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type)
if not pool:
logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, pool_type)
logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, normalized_pool_type)
return 0
deducted_credits = min(credits_required, pool.remaining_credits)

126
api/tests/unit_tests/services/test_billing_service.py
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 = {

167
api/tests/unit_tests/services/test_credit_pool_service.py
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(