opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-08-29 03:45:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): standardize console billing portal errors (#41055)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
非法操作
2026-08-24 01:55:48 +00:00
committed by GitHub
parent 9d6c846a49
commit ed45016475
22 changed files with 902 additions and 354 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/.importlinter
+17
View File
@@ -285,3 +285,20 @@ source_modules =
services.recommended_app_catalog_gateway
forbidden_modules =
flask
[importlinter:contract:billing-application-boundary]
name = Billing application services and ports are framework and infrastructure neutral
type = forbidden
source_modules =
services.billing_portal_service
services.partner_tenant_binding_service
forbidden_modules =
configs
controllers
extensions
flask
models
repositories
services.billing_service
sqlalchemy
werkzeug
api/controllers/console/billing/billing.py
+83 -39
View File
@@ -7,20 +7,23 @@ from werkzeug.exceptions import BadRequest
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.console import console_ns
from controllers.console.wraps import (
account_initialization_required,
model_validate,
only_edition_cloud,
setup_required,
with_current_tenant_id,
with_current_user,
from controllers.console.billing.error import (
BillingOperationFailedErrorResponse,
BillingUnavailableErrorResponse,
BillingUnprocessableEntityErrorResponse,
to_billing_request_error,
)
from enums import CloudPlan
from extensions.ext_database import db
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import model_validate
from enums import CloudPlan, DeploymentEdition
from extensions.ext_application_services import application_services
from fields.base import ResponseModel
from libs.login import login_required
from models import Account
from services.billing_service import BillingService
from libs.helper import dump_response
from machinery.context import RequestContext
from models.account import TenantAccountRole
from services.errors.billing import BillingError
_BILLING_PORTAL_ALLOWED_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN})
class SubscriptionQuery(BaseModel):
@@ -45,37 +48,78 @@ class BillingSubscriptionResponse(ResponseModel):
register_schema_models(console_ns, SubscriptionQuery, PartnerTenantsPayload)
register_response_schema_models(console_ns, BillingResponse, BillingInvoiceResponse, BillingSubscriptionResponse)
register_response_schema_models(
console_ns,
BillingOperationFailedErrorResponse,
BillingUnprocessableEntityErrorResponse,
BillingUnavailableErrorResponse,
BillingResponse,
BillingInvoiceResponse,
BillingSubscriptionResponse,
)
@console_ns.route("/billing/subscription")
class Subscription(Resource):
@console_ns.doc(params=query_params_from_model(SubscriptionQuery))
@console_ns.response(200, "Success", console_ns.models[BillingSubscriptionResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@with_current_user
@with_current_tenant_id
@console_ns.response(403, "Forbidden")
@console_ns.response(
422,
"Invalid subscription query",
console_ns.models[BillingUnprocessableEntityErrorResponse.__name__],
)
@console_ns.response(
502,
"Billing operation failed",
console_ns.models[BillingOperationFailedErrorResponse.__name__],
)
@console_ns.response(
503,
"Billing unavailable",
console_ns.models[BillingUnavailableErrorResponse.__name__],
)
@console_account_admission(
editions=frozenset({DeploymentEdition.CLOUD}),
allowed_roles=_BILLING_PORTAL_ALLOWED_ROLES,
)
@model_validate(SubscriptionQuery)
def get(self, req_data: SubscriptionQuery, current_tenant_id: str, current_user: Account):
BillingService.is_tenant_owner_or_admin(current_user, session=db.session())
return BillingService.get_subscription(req_data.plan, req_data.interval, current_user.email, current_tenant_id)
def get(self, req_data: SubscriptionQuery, request_context: RequestContext):
try:
data = application_services().billing_portal.get_subscription(
request_context,
plan=req_data.plan,
interval=req_data.interval,
)
except BillingError as error:
raise to_billing_request_error(error) from error
return dump_response(BillingSubscriptionResponse, data)
@console_ns.route("/billing/invoices")
class Invoices(Resource):
@console_ns.response(200, "Success", console_ns.models[BillingInvoiceResponse.__name__])
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account):
BillingService.is_tenant_owner_or_admin(current_user, session=db.session())
return BillingService.get_invoices(current_user.email, current_tenant_id)
@console_ns.response(403, "Forbidden")
@console_ns.response(
502,
"Billing operation failed",
console_ns.models[BillingOperationFailedErrorResponse.__name__],
)
@console_ns.response(
503,
"Billing unavailable",
console_ns.models[BillingUnavailableErrorResponse.__name__],
)
@console_account_admission(
editions=frozenset({DeploymentEdition.CLOUD}),
allowed_roles=_BILLING_PORTAL_ALLOWED_ROLES,
)
def get(self, request_context: RequestContext):
try:
data = application_services().billing_portal.get_invoices(request_context)
except BillingError as error:
raise to_billing_request_error(error) from error
return dump_response(BillingInvoiceResponse, data)
@console_ns.route("/billing/partners/<string:partner_key>/tenants")
@@ -86,20 +130,20 @@ class PartnerTenants(Resource):
@console_ns.expect(console_ns.models[PartnerTenantsPayload.__name__])
@console_ns.response(200, "Tenants synced to partner successfully", console_ns.models[BillingResponse.__name__])
@console_ns.response(400, "Invalid partner information")
@setup_required
@login_required
@account_initialization_required
@only_edition_cloud
@with_current_user
@console_account_admission(editions=frozenset({DeploymentEdition.CLOUD}))
@model_validate(PartnerTenantsPayload)
def put(self, req_data: PartnerTenantsPayload, current_user: Account, partner_key: str):
def put(self, req_data: PartnerTenantsPayload, request_context: RequestContext, partner_key: str):
try:
click_id = req_data.click_id
decoded_partner_key = base64.b64decode(partner_key).decode("utf-8")
except Exception as e:
raise BadRequest("Invalid partner_key") from e
if not click_id or not decoded_partner_key or not current_user.id:
if not click_id or not decoded_partner_key:
raise BadRequest("Invalid partner information")
return BillingService.sync_partner_tenants_bindings(current_user.id, decoded_partner_key, click_id)
return application_services().partner_tenant_bindings.sync(
account_id=request_context.account_id,
partner_key=decoded_partner_key,
click_id=click_id,
)
api/controllers/console/billing/error.py
+47
View File
@@ -0,0 +1,47 @@
from typing import Literal
from fields.base import ResponseModel
from libs.exception import BaseHTTPException
from services.errors.billing import (
BillingError,
BillingUpstreamInvalidResponseError,
BillingUpstreamUnavailableError,
)
class BillingUnprocessableEntityErrorResponse(ResponseModel):
code: Literal["unprocessable_entity"]
message: str
status: Literal[422]
class BillingOperationFailedErrorResponse(ResponseModel):
code: Literal["billing_operation_failed"]
message: str
status: Literal[502]
class BillingUnavailableErrorResponse(ResponseModel):
code: Literal["billing_unavailable"]
message: str
status: Literal[503]
class BillingOperationFailedError(BaseHTTPException):
error_code = "billing_operation_failed"
description = "We couldn't complete this request. Please try again. If the problem persists, contact support."
code = 502
class BillingUnavailableError(BaseHTTPException):
error_code = "billing_unavailable"
description = "This operation is temporarily unavailable. Please try again later."
code = 503
def to_billing_request_error(error: BillingError) -> BaseHTTPException:
if isinstance(error, BillingUpstreamInvalidResponseError):
return BillingOperationFailedError()
if isinstance(error, BillingUpstreamUnavailableError):
return BillingUnavailableError()
raise TypeError(f"Unsupported billing error: {type(error).__name__}")
api/controllers/console/workspace/model_providers.py
+2 -2
View File
@@ -20,7 +20,6 @@ from controllers.console.wraps import (
with_current_tenant_id,
with_current_user,
)
from extensions.ext_database import db
from fields.base import ResponseModel
from graphon.model_runtime.entities.model_entities import ModelType
from graphon.model_runtime.errors.validate import CredentialsValidateFailedError
@@ -403,6 +402,7 @@ class PreferredProviderTypeUpdateApi(Resource):
@console_ns.route("/workspaces/current/model-providers/<path:provider>/checkout-url")
class ModelProviderPaymentCheckoutUrlApi(Resource):
@console_ns.doc(deprecated=True)
@console_ns.response(
200,
"Model provider checkout URL retrieved successfully",
@@ -410,13 +410,13 @@ class ModelProviderPaymentCheckoutUrlApi(Resource):
)
@setup_required
@login_required
@is_admin_or_owner_required
@account_initialization_required
@with_current_user
@with_current_tenant_id
def get(self, current_tenant_id: str, current_user: Account, provider: str):
if provider != "anthropic":
raise ValueError(f"provider name {provider} is invalid")
BillingService.is_tenant_owner_or_admin(current_user, session=db.session())
data = BillingService.get_model_provider_payment_link(
provider_name=provider,
tenant_id=current_tenant_id,
api/extensions/ext_application_services.py
+13
View File
@@ -73,6 +73,8 @@ from services.auth.data_source_api_key_auth_gateways import (
TenantApiKeyAuthCredentialEncryptor,
)
from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService
from services.billing_portal_service import BillingPortalService
from services.billing_service import BillingService
from services.enterprise.enterprise_service import EnterpriseService
from services.errors.enterprise import EnterpriseServiceError
from services.explore_banner_query_service import ExploreBannerQueryService
@@ -81,6 +83,7 @@ from services.feature_service import FeatureService
from services.feature_service_gateway import FeatureServiceGateway
from services.file_service import FileService
from services.init_validation_service import InitValidationService
from services.partner_tenant_binding_service import PartnerTenantBindingService
from services.recommended_app_catalog_gateway import (
BuiltinRecommendedAppCatalogGateway,
RecommendedAppCatalogRouter,
@@ -141,6 +144,7 @@ class ApplicationServices:
accounts: AccountServices
account_activation: AccountActivationService
app_definitions: AppDefinitionQueryService
billing_portal: BillingPortalService
data_source_api_key_auth: DataSourceApiKeyAuthService
webapp_access: WebAppAccessQueryService
web_app_runtime: WebAppRuntimeQueryService
@@ -149,6 +153,7 @@ class ApplicationServices:
setup: SetupService
feature_queries: FeatureQueryService
init_validation: InitValidationService
partner_tenant_bindings: PartnerTenantBindingService
recommended_app_queries: RecommendedAppQueryService
trial_app_usage: TrialAppUsageRecorder
workspace_queries: WorkspaceQueryService
@@ -258,6 +263,11 @@ def build_application_services(
dify_config.CONSOLE_API_URL + "/console/api/workspaces/current/tool-provider/builtin/"
),
),
billing_portal=BillingPortalService(
accounts=accounts,
get_subscription=BillingService.get_subscription,
get_invoices=BillingService.get_invoices,
),
data_source_api_key_auth=DataSourceApiKeyAuthService(
bindings=data_source_api_key_auth_bindings,
validator=ProviderApiKeyAuthCredentialValidator(),
@@ -296,6 +306,9 @@ def build_application_services(
validation_required=(deployment_edition != DeploymentEdition.CLOUD and bool(initialization_password)),
expected_password=initialization_password,
),
partner_tenant_bindings=PartnerTenantBindingService(
sync_bindings=BillingService.sync_partner_tenants_bindings,
),
recommended_app_queries=RecommendedAppQueryService(
catalog=recommended_app_catalog,
trial_apps=TrialAppQueryRepository(session_factory=database_client),
api/openapi/markdown/console-openapi.md
+35 -1
View File
@@ -4998,6 +4998,9 @@ Restore a published workflow version into the draft workflow
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Success | **application/json**: [BillingInvoiceResponse](#billinginvoiceresponse)<br> |
| 403 | Forbidden | |
| 502 | Billing operation failed | **application/json**: [BillingOperationFailedErrorResponse](#billingoperationfailederrorresponse)<br> |
| 503 | Billing unavailable | **application/json**: [BillingUnavailableErrorResponse](#billingunavailableerrorresponse)<br> |
### [PUT] /billing/partners/{partner_key}/tenants
Sync partner tenants bindings
@@ -5034,6 +5037,10 @@ Sync partner tenants bindings
| Code | Description | Schema |
| ---- | ----------- | ------ |
| 200 | Success | **application/json**: [BillingSubscriptionResponse](#billingsubscriptionresponse)<br> |
| 403 | Forbidden | |
| 422 | Invalid subscription query | **application/json**: [BillingUnprocessableEntityErrorResponse](#billingunprocessableentityerrorresponse)<br> |
| 502 | Billing operation failed | **application/json**: [BillingOperationFailedErrorResponse](#billingoperationfailederrorresponse)<br> |
| 503 | Billing unavailable | **application/json**: [BillingUnavailableErrorResponse](#billingunavailableerrorresponse)<br> |
### [GET] /code-based-extension
Get code-based extension data by module name
@@ -10290,7 +10297,10 @@ Update a plugin endpoint
| ---- | ----------- | ------ |
| 200 | Model provider summaries retrieved successfully | **application/json**: [ModelProviderSummaryListResponse](#modelprovidersummarylistresponse)<br> |
### [GET] /workspaces/current/model-providers/{provider}/checkout-url
### ~~[GET] /workspaces/current/model-providers/{provider}/checkout-url~~
***DEPRECATED***
#### Parameters
| Name | Located in | Description | Required | Schema |
@@ -15452,6 +15462,14 @@ ExporleBanner status
| enabled | boolean | Deprecated. Use system features deployment_edition to determine the product edition. | Yes |
| subscription | [SubscriptionModel](#subscriptionmodel) | | Yes |
#### BillingOperationFailedErrorResponse
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| code | string | | Yes |
| message | string | | Yes |
| status | integer | | Yes |
#### BillingResponse
| Name | Type | Description | Required |
@@ -15464,6 +15482,22 @@ ExporleBanner status
| ---- | ---- | ----------- | -------- |
| url | string | | Yes |
#### BillingUnavailableErrorResponse
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| code | string | | Yes |
| message | string | | Yes |
| status | integer | | Yes |
#### BillingUnprocessableEntityErrorResponse
| Name | Type | Description | Required |
| ---- | ---- | ----------- | -------- |
| code | string | | Yes |
| message | string | | Yes |
| status | integer | | Yes |
#### BinaryFileResponse
| Name | Type | Description | Required |
api/services/billing_portal_service.py
+49
View File
@@ -0,0 +1,49 @@
"""Application service for Console Billing portal links."""
from collections.abc import Callable
from typing import TypedDict
from machinery.context import RequestContext
from services.account_errors import AccountNotFoundError
from services.account_ports import AccountRepository
class BillingPortalLink(TypedDict):
url: str
class BillingPortalService:
def __init__(
self,
*,
accounts: AccountRepository,
get_subscription: Callable[[str, str, str, str], BillingPortalLink],
get_invoices: Callable[[str, str], BillingPortalLink],
) -> None:
self._accounts = accounts
self._get_subscription = get_subscription
self._get_invoices = get_invoices
def get_subscription(
self,
context: RequestContext,
*,
plan: str,
interval: str,
) -> BillingPortalLink:
email, workspace_id = self._resolve_account_email_and_workspace_id(context)
return self._get_subscription(plan, interval, email, workspace_id)
def get_invoices(self, context: RequestContext) -> BillingPortalLink:
email, workspace_id = self._resolve_account_email_and_workspace_id(context)
return self._get_invoices(email, workspace_id)
def _resolve_account_email_and_workspace_id(self, context: RequestContext) -> tuple[str, str]:
workspace_id = context.active_workspace_id
if workspace_id is None:
raise RuntimeError("Console account admission did not resolve an active workspace")
account = self._accounts.get(context.account_id)
if account is None:
raise AccountNotFoundError
return account.email, workspace_id
api/services/billing_service.py
+49 -76
View File
@@ -2,12 +2,10 @@ import json
import logging
import os
from collections.abc import Sequence
from typing import Any, Literal, NotRequired, TypedDict
from typing import Literal, NotRequired, TypedDict
import httpx
from pydantic import TypeAdapter
from sqlalchemy import select
from sqlalchemy.orm import Session
from pydantic import TypeAdapter, ValidationError
from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
from werkzeug.exceptions import InternalServerError
@@ -15,7 +13,12 @@ from core.helper.http_client_pooling import get_pooled_http_client
from enums import CloudPlan
from extensions.ext_redis import redis_client
from libs.helper import RateLimiter
from models import Account, TenantAccountJoin, TenantAccountRole
from models import Account
from services.billing_portal_service import BillingPortalLink
from services.errors.billing import (
BillingUpstreamInvalidResponseError,
BillingUpstreamUnavailableError,
)
logger = logging.getLogger(__name__)
@@ -31,6 +34,12 @@ _http_client: httpx.Client = get_pooled_http_client(
EmailFreezeType = Literal["freeze", "email_domain_suspended"]
class _BillingHTTPStatusError(ValueError):
def __init__(self, message: str, status_code: int):
super().__init__(message)
self.status_code = status_code
class SubscriptionPlan(TypedDict):
"""Tenant subscriptionplan information."""
@@ -38,6 +47,9 @@ class SubscriptionPlan(TypedDict):
expiration_date: int
_billing_portal_link_adapter = TypeAdapter(BillingPortalLink)
class QuotaReserveResult(TypedDict):
reservation_id: str
available: int
@@ -187,14 +199,6 @@ class AccountNotificationDict(TypedDict, total=False):
notifications: list[dict]
class UpsertNotificationDict(TypedDict):
notification_id: str
class BatchAddNotificationAccountsDict(TypedDict):
count: int
class DismissNotificationDict(TypedDict):
success: bool
@@ -359,9 +363,11 @@ class BillingService:
}
@classmethod
def get_subscription(cls, plan: str, interval: str, prefilled_email: str = "", tenant_id: str = ""):
def get_subscription(
cls, plan: str, interval: str, prefilled_email: str = "", tenant_id: str = ""
) -> BillingPortalLink:
params = {"plan": plan, "interval": interval, "prefilled_email": prefilled_email, "tenant_id": tenant_id}
return cls._send_request("GET", "/subscription/payment-link", params=params)
return cls._send_billing_portal_request("/subscription/payment-link", params=params)
@classmethod
def get_model_provider_payment_link(cls, provider_name: str, tenant_id: str, account_id: str, prefilled_email: str):
@@ -374,9 +380,9 @@ class BillingService:
return cls._send_request("GET", "/model-provider/payment-link", params=params)
@classmethod
def get_invoices(cls, prefilled_email: str = "", tenant_id: str = ""):
def get_invoices(cls, prefilled_email: str = "", tenant_id: str = "") -> BillingPortalLink:
params = {"prefilled_email": prefilled_email, "tenant_id": tenant_id}
return cls._send_request("GET", "/invoices", params=params)
return cls._send_billing_portal_request("/invoices", params=params)
@classmethod
def update_tenant_feature_plan_usage(
@@ -444,7 +450,10 @@ class BillingService:
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.")
raise _BillingHTTPStatusError(
"Unable to retrieve billing information. Please try again later or contact support.",
response.status_code,
)
if method == "PUT":
if response.status_code == httpx.codes.INTERNAL_SERVER_ERROR:
raise InternalServerError(
@@ -459,21 +468,28 @@ class BillingService:
raise ValueError(f"Unable to process delete request {url}. Please try again later or contact support.")
return response.json()
@staticmethod
def is_tenant_owner_or_admin(current_user: Account, *, session: Session):
tenant_id = current_user.current_tenant_id
join: TenantAccountJoin | None = session.scalar(
select(TenantAccountJoin)
.where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.account_id == current_user.id)
.limit(1)
)
if not join:
raise ValueError("Tenant account join not found")
if not TenantAccountRole.is_privileged_role(TenantAccountRole(join.role)):
raise ValueError("Only team owner or team admin can perform this action")
@classmethod
def _send_billing_portal_request(
cls,
endpoint: str,
*,
params: dict[str, str],
) -> BillingPortalLink:
try:
response = cls._send_request("GET", endpoint, params=params)
return _billing_portal_link_adapter.validate_python(response)
except _BillingHTTPStatusError as error:
if error.status_code in {httpx.codes.REQUEST_TIMEOUT, httpx.codes.TOO_MANY_REQUESTS} or (
error.status_code >= 500
):
raise BillingUpstreamUnavailableError from error
raise BillingUpstreamInvalidResponseError from error
except httpx.RequestError as error:
raise BillingUpstreamUnavailableError from error
except (json.JSONDecodeError, UnicodeDecodeError, ValidationError) as error:
raise BillingUpstreamInvalidResponseError from error
except ValueError as error:
raise RuntimeError("Unexpected billing service value error") from error
@classmethod
def delete_account(cls, account_id: str):
@@ -744,49 +760,6 @@ class BillingService:
"""
return cls._send_request("GET", "/notifications/active", params={"account_id": account_id})
@classmethod
def upsert_notification(
cls,
contents: list[LangContentDict],
frequency: str = "once",
status: str = "active",
notification_id: str | None = None,
start_time: str | None = None,
end_time: str | None = None,
) -> UpsertNotificationDict:
"""Create or update a notification.
contents: list of {"lang": str, "title": str, "subtitle": str, "body": str, "title_pic_url": str}
start_time / end_time: RFC3339 strings (e.g. "2026-03-01T00:00:00Z"), optional.
Returns {"notification_id": str}.
"""
payload: dict[str, Any] = {
"contents": contents,
"frequency": frequency,
"status": status,
}
if notification_id:
payload["notification_id"] = notification_id
if start_time:
payload["start_time"] = start_time
if end_time:
payload["end_time"] = end_time
return cls._send_request("POST", "/notifications", json=payload)
@classmethod
def batch_add_notification_accounts(
cls, notification_id: str, account_ids: list[str]
) -> BatchAddNotificationAccountsDict:
"""Register target account IDs for a notification (max 1000 per call).
Returns {"count": int}.
"""
return cls._send_request(
"POST",
f"/notifications/{notification_id}/accounts",
json={"account_ids": account_ids},
)
@classmethod
def dismiss_notification(cls, notification_id: str, account_id: str) -> DismissNotificationDict:
"""Mark a notification as dismissed for an account.
api/services/errors/billing.py
+10
View File
@@ -0,0 +1,10 @@
class BillingError(Exception):
pass
class BillingUpstreamInvalidResponseError(BillingError):
pass
class BillingUpstreamUnavailableError(BillingError):
pass
api/services/partner_tenant_binding_service.py
+12
View File
@@ -0,0 +1,12 @@
"""Application service for partner tenant bindings."""
from collections.abc import Callable
from typing import Any
class PartnerTenantBindingService:
def __init__(self, *, sync_bindings: Callable[[str, str, str], dict[str, Any]]) -> None:
self._sync_bindings = sync_bindings
def sync(self, *, account_id: str, partner_key: str, click_id: str) -> dict[str, Any]:
return self._sync_bindings(account_id, partner_key, click_id)
api/tests/test_containers_integration_tests/services/test_billing_service.py
-63
View File
@@ -1,14 +1,10 @@
import json
from collections.abc import Generator
from unittest.mock import patch
from uuid import uuid4
import pytest
from flask import Flask
from sqlalchemy.orm import Session
from extensions.ext_redis import redis_client
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from services.billing_service import BillingService
@@ -368,62 +364,3 @@ class TestBillingServiceGetPlanBulkWithCache:
assert ttl_1_new <= 600
assert ttl_2 > 0
assert ttl_2 <= 600
class TestBillingServiceIsTenantOwnerOrAdmin:
"""
Integration tests for BillingService.is_tenant_owner_or_admin.
Verifies that non-privileged roles (EDITOR, DATASET_OPERATOR) raise ValueError
when checked against real TenantAccountJoin rows in PostgreSQL.
"""
@pytest.fixture(autouse=True)
def _auto_rollback(self, db_session_with_containers: Session) -> Generator[None, None, None]:
yield
db_session_with_containers.rollback()
def _create_account_with_tenant_role(self, db_session: Session, role: TenantAccountRole) -> tuple[Account, Tenant]:
tenant = Tenant(name=f"Tenant {uuid4()}")
db_session.add(tenant)
db_session.flush()
account = Account(
name=f"Account {uuid4()}",
email=f"billing_{uuid4()}@example.com",
password="hashed-password",
password_salt="salt",
interface_language="en-US",
timezone="UTC",
)
db_session.add(account)
db_session.flush()
join = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=role,
current=True,
)
db_session.add(join)
db_session.flush()
# Wire up in-memory reference so current_tenant_id resolves
account._current_tenant = tenant
return account, tenant
def test_is_tenant_owner_or_admin_editor_role_raises_error(self, db_session_with_containers: Session) -> None:
"""is_tenant_owner_or_admin raises ValueError for EDITOR role."""
account, _ = self._create_account_with_tenant_role(db_session_with_containers, TenantAccountRole.EDITOR)
with pytest.raises(ValueError, match="Only team owner or team admin can perform this action"):
BillingService.is_tenant_owner_or_admin(account, session=db_session_with_containers)
def test_is_tenant_owner_or_admin_dataset_operator_raises_error(self, db_session_with_containers: Session) -> None:
"""is_tenant_owner_or_admin raises ValueError for DATASET_OPERATOR role."""
account, _ = self._create_account_with_tenant_role(
db_session_with_containers, TenantAccountRole.DATASET_OPERATOR
)
with pytest.raises(ValueError, match="Only team owner or team admin can perform this action"):
BillingService.is_tenant_owner_or_admin(account, session=db_session_with_containers)
api/tests/unit_tests/controllers/console/billing/test_billing.py
+143 -43
View File
@@ -1,6 +1,9 @@
import base64
import json
from unittest.mock import patch
from collections.abc import Iterator
from inspect import unwrap
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
@@ -8,10 +11,126 @@ from sqlalchemy.orm import Session
from werkzeug.exceptions import BadRequest, UnprocessableEntity
from controllers.console import wraps as console_wraps
from controllers.console.billing.billing import PartnerTenants
from enums import DeploymentEdition
from controllers.console.billing.billing import Invoices, PartnerTenants, Subscription, SubscriptionQuery
from controllers.console.billing.error import (
BillingOperationFailedError,
BillingUnavailableError,
)
from enums import CloudPlan, DeploymentEdition
from machinery.context import RequestContext
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.model import DifySetup
from services.errors.billing import (
BillingUpstreamInvalidResponseError,
BillingUpstreamUnavailableError,
)
class TestBillingPortal:
@pytest.fixture
def app(self) -> Flask:
app = Flask(__name__)
app.config["TESTING"] = True
return app
@pytest.fixture
def request_context(self) -> RequestContext:
return RequestContext(
request_id="request-1",
trace_id="trace-1",
account_id="account-1",
active_workspace_id="tenant-1",
)
@pytest.fixture
def billing_portal(self) -> MagicMock:
return MagicMock()
@pytest.fixture(autouse=True)
def mock_application_services(self, billing_portal: MagicMock) -> Iterator[None]:
with patch(
"controllers.console.billing.billing.application_services",
return_value=SimpleNamespace(billing_portal=billing_portal),
):
yield
def test_get_subscription_uses_admission_context_and_response_contract(
self,
app: Flask,
request_context: RequestContext,
billing_portal: MagicMock,
) -> None:
resource = Subscription()
method = unwrap(resource.get)
query = SubscriptionQuery(plan=CloudPlan.PROFESSIONAL, interval="month")
billing_portal.get_subscription.return_value = {"url": "https://billing.example.com/checkout"}
with app.test_request_context("/billing/subscription"):
result = method(resource, query, request_context)
billing_portal.get_subscription.assert_called_once_with(
request_context,
plan=CloudPlan.PROFESSIONAL,
interval="month",
)
assert result == {"url": "https://billing.example.com/checkout"}
def test_get_invoices_uses_admission_context_and_response_contract(
self,
app: Flask,
request_context: RequestContext,
billing_portal: MagicMock,
) -> None:
resource = Invoices()
method = unwrap(resource.get)
billing_portal.get_invoices.return_value = {"url": "https://billing.example.com/portal"}
with app.test_request_context("/billing/invoices"):
result = method(resource, request_context)
billing_portal.get_invoices.assert_called_once_with(request_context)
assert result == {"url": "https://billing.example.com/portal"}
def test_get_invoices_translates_unavailable_operation(
self,
app: Flask,
request_context: RequestContext,
billing_portal: MagicMock,
) -> None:
resource = Invoices()
method = unwrap(resource.get)
billing_portal.get_invoices.side_effect = BillingUpstreamUnavailableError
with app.test_request_context("/billing/invoices"):
with pytest.raises(BillingUnavailableError) as exc_info:
method(resource, request_context)
assert exc_info.value.data == {
"code": "billing_unavailable",
"message": "This operation is temporarily unavailable. Please try again later.",
"status": 503,
}
def test_get_subscription_translates_invalid_upstream_response(
self,
app: Flask,
request_context: RequestContext,
billing_portal: MagicMock,
) -> None:
resource = Subscription()
method = unwrap(resource.get)
query = SubscriptionQuery(plan=CloudPlan.PROFESSIONAL, interval="month")
billing_portal.get_subscription.side_effect = BillingUpstreamInvalidResponseError
with app.test_request_context("/billing/subscription"):
with pytest.raises(BillingOperationFailedError) as exc_info:
method(resource, query, request_context)
assert exc_info.value.data == {
"code": "billing_operation_failed",
"message": "We couldn't complete this request. Please try again. If the problem persists, contact support.",
"status": 502,
}
@pytest.mark.parametrize(
@@ -55,10 +174,13 @@ class TestPartnerTenants:
return account
@pytest.fixture
def mock_billing_service(self):
"""Mock BillingService."""
with patch("controllers.console.billing.billing.BillingService") as mock_service:
yield mock_service
def partner_tenant_bindings(self):
service = MagicMock()
with patch(
"controllers.console.billing.billing.application_services",
return_value=SimpleNamespace(partner_tenant_bindings=service),
):
yield service
@pytest.fixture
def mock_decorators(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session):
@@ -74,14 +196,14 @@ class TestPartnerTenants:
yield mock_csrf
console_wraps._is_setup_completed.reset_success()
def test_put_success(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
def test_put_success(self, app: Flask, mock_account, partner_tenant_bindings, mock_decorators):
"""Test successful partner tenants bindings sync."""
# Arrange
partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
click_id = "click-id-789"
expected_response = {"result": "success", "data": {"synced": True}}
mock_billing_service.sync_partner_tenants_bindings.return_value = expected_response
partner_tenant_bindings.sync.return_value = expected_response
with app.test_request_context(
method="PUT",
@@ -100,11 +222,13 @@ class TestPartnerTenants:
# Assert
assert result == expected_response
mock_billing_service.sync_partner_tenants_bindings.assert_called_once_with(
mock_account.id, "partner-key-123", click_id
partner_tenant_bindings.sync.assert_called_once_with(
account_id=mock_account.id,
partner_key="partner-key-123",
click_id=click_id,
)
def test_put_invalid_partner_key_base64(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
def test_put_invalid_partner_key_base64(self, app: Flask, mock_account, partner_tenant_bindings, mock_decorators):
"""Test that invalid base64 partner_key raises BadRequest."""
# Arrange
invalid_partner_key = "invalid-base64-!@#$"
@@ -129,7 +253,7 @@ class TestPartnerTenants:
resource.put(invalid_partner_key)
assert "Invalid partner_key" in str(exc_info.value)
def test_put_missing_click_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
def test_put_missing_click_id(self, app: Flask, mock_account, partner_tenant_bindings, mock_decorators):
"""Test that missing click_id raises UnprocessableEntity (422)."""
# Arrange
partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
@@ -154,7 +278,7 @@ class TestPartnerTenants:
resource.put(partner_key_encoded)
def test_put_billing_service_json_decode_error(
self, app: Flask, mock_account, mock_billing_service, mock_decorators
self, app: Flask, mock_account, partner_tenant_bindings, mock_decorators
):
"""Test handling of billing service JSON decode error.
@@ -174,7 +298,7 @@ class TestPartnerTenants:
# Simulate JSON decode error when billing service returns invalid JSON
# This happens when billing service returns non-200 with empty/invalid response body
json_decode_error = json.JSONDecodeError("Expecting value", "", 0)
mock_billing_service.sync_partner_tenants_bindings.side_effect = json_decode_error
partner_tenant_bindings.sync.side_effect = json_decode_error
with app.test_request_context(
method="PUT",
@@ -201,7 +325,7 @@ class TestPartnerTenants:
assert isinstance(exc_info.value, json.JSONDecodeError)
assert "Expecting value" in str(exc_info.value)
def test_put_empty_click_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
def test_put_empty_click_id(self, app: Flask, mock_account, partner_tenant_bindings, mock_decorators):
"""Test that empty click_id raises BadRequest."""
# Arrange
partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
@@ -226,7 +350,9 @@ class TestPartnerTenants:
resource.put(partner_key_encoded)
assert "Invalid partner information" in str(exc_info.value)
def test_put_empty_partner_key_after_decode(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
def test_put_empty_partner_key_after_decode(
self, app: Flask, mock_account, partner_tenant_bindings, mock_decorators
):
"""Test that empty partner_key after decode raises BadRequest."""
# Arrange
# Base64 encode an empty string
@@ -251,29 +377,3 @@ class TestPartnerTenants:
with pytest.raises(BadRequest) as exc_info:
resource.put(empty_partner_key_encoded)
assert "Invalid partner information" in str(exc_info.value)
def test_put_empty_user_id(self, app: Flask, mock_account, mock_billing_service, mock_decorators):
"""Test that empty user id raises BadRequest."""
# Arrange
partner_key_encoded = base64.b64encode(b"partner-key-123").decode("utf-8")
click_id = "click-id-789"
mock_account.id = None # Empty user id
with app.test_request_context(
method="PUT",
json={"click_id": click_id},
path=f"/billing/partners/{partner_key_encoded}/tenants",
):
with (
patch(
"controllers.console.wraps.current_account_with_tenant",
return_value=(mock_account, mock_account.current_tenant_id),
),
patch("libs.login._get_user", return_value=mock_account),
):
resource = PartnerTenants()
# Act & Assert
with pytest.raises(BadRequest) as exc_info:
resource.put(partner_key_encoded)
assert "Invalid partner information" in str(exc_info.value)
api/tests/unit_tests/controllers/console/billing/test_error.py
+9
View File
@@ -0,0 +1,9 @@
import pytest
from controllers.console.billing.error import to_billing_request_error
from services.errors.billing import BillingError
def test_to_billing_request_error_rejects_unknown_error() -> None:
with pytest.raises(TypeError, match="Unsupported billing error"):
to_billing_request_error(BillingError())
api/tests/unit_tests/controllers/console/workspace/test_model_providers.py
+14 -16
View File
@@ -1,8 +1,8 @@
from inspect import unwrap
from unittest.mock import ANY, patch
from unittest.mock import patch
import pytest
from flask import Flask
from flask import Flask, g
from pydantic_core import ValidationError
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
@@ -20,6 +20,7 @@ from controllers.console.workspace.model_providers import (
PreferredProviderTypeUpdateApi,
)
from core.entities.provider_entities import CredentialConfiguration
from enums import DeploymentEdition
from graphon.model_runtime.entities.common_entities import I18nObject
from graphon.model_runtime.entities.model_entities import ModelType
from graphon.model_runtime.entities.provider_entities import ConfigurateMethod
@@ -571,13 +572,8 @@ class TestModelProviderPaymentCheckoutUrlApi:
method = unwrap(api.get)
user = make_account()
with (
app.test_request_context("/"),
patch(
"controllers.console.workspace.model_providers.BillingService.is_tenant_owner_or_admin",
return_value=None,
) as is_tenant_owner_or_admin,
patch(
"controllers.console.workspace.model_providers.BillingService.get_model_provider_payment_link",
return_value={"payment_link": "https://payment.example.com/provider"},
@@ -585,7 +581,6 @@ class TestModelProviderPaymentCheckoutUrlApi:
):
result = method(api, "tenant1", user, provider="anthropic")
is_tenant_owner_or_admin.assert_called_once_with(user, session=ANY)
get_model_provider_payment_link.assert_called_once_with(
provider_name="anthropic",
tenant_id="tenant1",
@@ -602,18 +597,21 @@ class TestModelProviderPaymentCheckoutUrlApi:
with pytest.raises(ValueError):
method(api, "tenant1", make_account(), provider="openai")
def test_permission_denied(self, app: Flask):
def test_checkout_rejects_non_privileged_role(self, app: Flask):
api = ModelProviderPaymentCheckoutUrlApi()
method = unwrap(api.get)
user = make_account()
account = make_account()
with (
app.test_request_context("/"),
patch.object(dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.CLOUD),
patch.object(dify_config, "LOGIN_DISABLED", True),
patch.object(dify_config, "RBAC_ENABLED", False),
patch(
"controllers.console.workspace.model_providers.BillingService.is_tenant_owner_or_admin",
side_effect=Forbidden(),
),
"controllers.console.workspace.model_providers.BillingService.get_model_provider_payment_link",
) as get_model_provider_payment_link,
):
g._login_user = account
with pytest.raises(Forbidden):
method(api, "tenant1", user, provider="anthropic")
api.get(provider="anthropic")
get_model_provider_payment_link.assert_not_called()
api/tests/unit_tests/controllers/test_swagger.py
+63
View File
@@ -618,6 +618,69 @@ def test_console_member_invite_documents_bad_request_response():
}
def test_console_billing_routes_document_error_responses(monkeypatch: pytest.MonkeyPatch):
from configs import dify_config
from controllers.console import bp as console_bp
monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True)
app = Flask(__name__)
app.config["TESTING"] = True
app.config["RESTX_INCLUDE_ALL_MODELS"] = True
app.register_blueprint(console_bp)
payload = app.test_client().get("/console/api/openapi.json").get_json()
expected_responses = {
("/billing/subscription", "get"): {
"422": "BillingUnprocessableEntityErrorResponse",
"502": "BillingOperationFailedErrorResponse",
"503": "BillingUnavailableErrorResponse",
},
("/billing/invoices", "get"): {
"502": "BillingOperationFailedErrorResponse",
"503": "BillingUnavailableErrorResponse",
},
}
for (path, method), responses in expected_responses.items():
operation = payload["paths"][path][method]
for status, model_name in responses.items():
schema = operation["responses"][status]["content"]["application/json"]["schema"]
assert schema["$ref"] == f"#/components/schemas/{model_name}"
forbidden_response = operation["responses"]["403"]
assert forbidden_response["description"] == "Forbidden"
assert "content" not in forbidden_response
expected_error_contracts = {
"BillingUnprocessableEntityErrorResponse": ("unprocessable_entity", 422),
"BillingOperationFailedErrorResponse": ("billing_operation_failed", 502),
"BillingUnavailableErrorResponse": ("billing_unavailable", 503),
}
schemas = payload["components"]["schemas"]
for model_name, (error_code, status) in expected_error_contracts.items():
properties = schemas[model_name]["properties"]
assert properties["code"]["const"] == error_code
assert properties["status"]["const"] == status
def test_console_model_provider_checkout_route_is_deprecated(monkeypatch: pytest.MonkeyPatch):
from configs import dify_config
from controllers.console import bp as console_bp
monkeypatch.setattr(dify_config, "SWAGGER_UI_ENABLED", True)
app = Flask(__name__)
app.config["TESTING"] = True
app.config["RESTX_INCLUDE_ALL_MODELS"] = True
app.register_blueprint(console_bp)
payload = app.test_client().get("/console/api/openapi.json").get_json()
operation = payload["paths"]["/workspaces/current/model-providers/{provider}/checkout-url"]["get"]
assert operation["deprecated"] is True
def test_console_plugin_category_list_exported_schema_uses_typed_items(tmp_path: Path):
from dev.generate_swagger_specs import generate_specs
api/tests/unit_tests/extensions/test_ext_application_services.py
+62
View File
@@ -15,6 +15,8 @@ from sqlalchemy.orm import Session, sessionmaker
from enums import DeploymentEdition, WebAppAccessMode
from extensions import ext_application_services
from extensions.ext_redis import RedisClientWrapper
from machinery.context import RequestContext
from models.account import Account
from models.model import AccountTrialAppRecord, DifySetup
from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository
from repositories.account_integration_repository import SQLAlchemyAccountIntegrationRepository
@@ -28,9 +30,12 @@ from services.account_activation_adapters import (
)
from services.account_avatar_file_gateway import SQLAlchemyAccountAvatarFileGateway
from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService
from services.billing_portal_service import BillingPortalService
from services.billing_service import BillingService
from services.enterprise.enterprise_service import WebAppSettings
from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFoundError
from services.init_validation_service import InvalidInitializationPasswordError
from services.partner_tenant_binding_service import PartnerTenantBindingService
from services.tag_application_service import TagApplicationService
from services.webapp_access_query_service import WebAppAccessUnavailableError
@@ -171,6 +176,63 @@ def test_build_application_services_wires_tag_boundary(
assert isinstance(services.tags, TagApplicationService)
def test_build_application_services_wires_billing_service(
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
) -> None:
account = Account(name="Billing Owner", email="owner@example.com")
account.id = "account-1"
sqlite_session.add(account)
sqlite_session.commit()
with (
patch.object(
BillingService,
"get_subscription",
return_value={"url": "https://billing.example.com/checkout"},
) as get_subscription,
patch.object(
BillingService,
"get_invoices",
return_value={"url": "https://billing.example.com/portal"},
) as get_invoices,
patch.object(
BillingService,
"sync_partner_tenants_bindings",
return_value={"result": "success"},
) as sync_partner_tenants_bindings,
):
services = ext_application_services.build_application_services(
database_client=sqlite_session_factory,
deployment_edition=DeploymentEdition.COMMUNITY,
initialization_password="",
redis=MagicMock(spec=RedisClientWrapper),
)
request_context = RequestContext(
request_id="request-1",
trace_id="trace-1",
account_id="account-1",
active_workspace_id="workspace-1",
)
assert isinstance(services.billing_portal, BillingPortalService)
assert services.billing_portal.get_subscription(
request_context,
plan="professional",
interval="month",
) == {"url": "https://billing.example.com/checkout"}
assert services.billing_portal.get_invoices(request_context) == {"url": "https://billing.example.com/portal"}
assert isinstance(services.partner_tenant_bindings, PartnerTenantBindingService)
assert services.partner_tenant_bindings.sync(
account_id="account-1",
partner_key="partner-key",
click_id="click-1",
) == {"result": "success"}
get_subscription.assert_called_once_with("professional", "month", "owner@example.com", "workspace-1")
get_invoices.assert_called_once_with("owner@example.com", "workspace-1")
sync_partner_tenants_bindings.assert_called_once_with("account-1", "partner-key", "click-1")
def test_build_application_services_wires_account_profile_repository(
sqlite_session_factory: sessionmaker[Session],
) -> None:
api/tests/unit_tests/services/test_account_service.py
+1
View File
@@ -447,6 +447,7 @@ class TestAccountService:
assert result is account
assert result.current_tenant_id == tenant.id
assert result.current_role == TenantAccountRole.NORMAL
mock_refresh_last_active.assert_called_once_with(account, sqlite_session)
def test_load_user_not_found(self, sqlite_session: Session) -> None:
api/tests/unit_tests/services/test_billing_portal_service.py
+116
View File
@@ -0,0 +1,116 @@
from datetime import datetime
from unittest.mock import MagicMock, Mock
import pytest
from machinery.context import RequestContext
from services.account_errors import AccountNotFoundError
from services.account_ports import AccountRepository
from services.billing_portal_service import BillingPortalService
from services.entities.account_entities import AccountSnapshot
def _context(*, workspace_id: str | None = "workspace-1") -> RequestContext:
return RequestContext(
request_id="request-1",
trace_id="trace-1",
account_id="account-1",
active_workspace_id=workspace_id,
)
def _account() -> AccountSnapshot:
return AccountSnapshot(
id="account-1",
name="Account",
email="owner@example.com",
avatar=None,
is_password_set=False,
interface_language="en-US",
interface_theme="light",
timezone="UTC",
last_login_at=None,
last_login_ip=None,
status="active",
initialized_at=None,
created_at=datetime(2026, 1, 1),
)
@pytest.fixture
def get_subscription() -> MagicMock:
return MagicMock()
@pytest.fixture
def get_invoices() -> MagicMock:
return MagicMock()
@pytest.fixture
def accounts() -> Mock:
accounts = Mock(spec=AccountRepository)
accounts.get.return_value = _account()
return accounts
@pytest.fixture
def service(accounts: Mock, get_subscription: MagicMock, get_invoices: MagicMock) -> BillingPortalService:
return BillingPortalService(accounts=accounts, get_subscription=get_subscription, get_invoices=get_invoices)
def test_get_subscription_loads_email_and_delegates(
service: BillingPortalService,
accounts: Mock,
get_subscription: MagicMock,
) -> None:
get_subscription.return_value = {"url": "https://billing.example.com/checkout"}
result = service.get_subscription(
_context(),
plan="professional",
interval="month",
)
assert result == {"url": "https://billing.example.com/checkout"}
accounts.get.assert_called_once_with("account-1")
get_subscription.assert_called_once_with("professional", "month", "owner@example.com", "workspace-1")
def test_get_invoices_loads_email_and_delegates(
service: BillingPortalService,
accounts: Mock,
get_invoices: MagicMock,
) -> None:
get_invoices.return_value = {"url": "https://billing.example.com/portal"}
result = service.get_invoices(_context())
assert result == {"url": "https://billing.example.com/portal"}
accounts.get.assert_called_once_with("account-1")
get_invoices.assert_called_once_with("owner@example.com", "workspace-1")
def test_missing_account_does_not_call_billing(
service: BillingPortalService,
accounts: Mock,
get_invoices: MagicMock,
) -> None:
accounts.get.return_value = None
with pytest.raises(AccountNotFoundError):
service.get_invoices(_context())
get_invoices.assert_not_called()
def test_missing_workspace_does_not_query_account_or_billing(
service: BillingPortalService,
accounts: Mock,
get_invoices: MagicMock,
) -> None:
with pytest.raises(RuntimeError, match="did not resolve an active workspace"):
service.get_invoices(_context(workspace_id=None))
accounts.get.assert_not_called()
get_invoices.assert_not_called()
api/tests/unit_tests/services/test_billing_service.py
+110 -114
View File
@@ -9,8 +9,7 @@ This test module covers all aspects of the billing service including:
- Cache management for billing data
- Partner integration features
Network, billing-provider, and cache boundaries are mocked; database authorization
paths use isolated in-memory SQLite sessions with persisted membership rows.
Network, billing-provider, and cache boundaries are mocked.
Tests follow the Arrange-Act-Assert pattern for clarity.
"""
@@ -20,15 +19,17 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
from sqlalchemy.orm import Session
from werkzeug.exceptions import InternalServerError
from enums import CloudPlan
from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
from services.billing_service import BillingService
from models import Account, Tenant
from services.billing_service import BillingService, _BillingHTTPStatusError
from services.errors.billing import (
BillingUpstreamInvalidResponseError,
BillingUpstreamUnavailableError,
)
TENANT_ID = "11111111-1111-1111-1111-111111111111"
OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222"
ACCOUNT_ID = "33333333-3333-3333-3333-333333333333"
@@ -41,43 +42,6 @@ def _account(*, account_id: str = ACCOUNT_ID, email: str = "user@example.com", t
return account
def _persist_membership(
sqlite_session: Session,
*,
role: TenantAccountRole | None,
add_other_tenant_membership: bool = False,
) -> Account:
account = _account()
tenant = account.current_tenant
assert tenant is not None
sqlite_session.add_all([tenant, account])
if role is not None:
sqlite_session.add(
TenantAccountJoin(
tenant_id=TENANT_ID,
account_id=ACCOUNT_ID,
current=True,
role=role,
)
)
if add_other_tenant_membership:
other_tenant = Tenant(name="Other Tenant")
other_tenant.id = OTHER_TENANT_ID
sqlite_session.add_all(
[
other_tenant,
TenantAccountJoin(
tenant_id=OTHER_TENANT_ID,
account_id=ACCOUNT_ID,
current=False,
role=TenantAccountRole.OWNER,
),
]
)
sqlite_session.commit()
return account
class TestBillingServiceSendRequest:
"""Unit tests for BillingService._send_request method.
@@ -146,15 +110,16 @@ class TestBillingServiceSendRequest:
"status_code", [httpx.codes.NOT_FOUND, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.BAD_REQUEST]
)
def test_get_request_non_200_status_code(self, mock_httpx_request, mock_billing_config, status_code):
"""Test GET request with non-200 status code raises ValueError."""
"""Test GET request preserves the upstream status for its public caller."""
# Arrange
mock_response = MagicMock()
mock_response.status_code = status_code
mock_httpx_request.return_value = mock_response
# Act & Assert
with pytest.raises(ValueError) as exc_info:
with pytest.raises(_BillingHTTPStatusError) as exc_info:
BillingService._send_request("GET", "/test")
assert exc_info.value.status_code == status_code
assert "Unable to retrieve billing information" in str(exc_info.value)
def test_put_request_success(self, mock_httpx_request, mock_billing_config):
@@ -363,6 +328,85 @@ class TestBillingServiceSendRequest:
assert mock_httpx_request.call_count > 1
class TestBillingServicePortalRequest:
def test_sends_get_request(self) -> None:
params = {"tenant_id": "tenant-1"}
with patch.object(
BillingService,
"_send_request",
return_value={"url": "https://example.com", "ignored": True},
) as send_request:
result = BillingService._send_billing_portal_request("/test", params=params)
assert result == {"url": "https://example.com"}
send_request.assert_called_once_with("GET", "/test", params=params)
def test_invalid_response_shape_is_invalid_upstream_response(self) -> None:
with patch.object(BillingService, "_send_request", return_value={}):
with pytest.raises(BillingUpstreamInvalidResponseError):
BillingService._send_billing_portal_request("/test", params={})
@pytest.mark.parametrize(
"status_code",
[httpx.codes.BAD_REQUEST, httpx.codes.UNAUTHORIZED, httpx.codes.NOT_FOUND],
)
def test_terminal_http_response_is_invalid_upstream_response(self, status_code: int) -> None:
with patch.object(
BillingService,
"_send_request",
side_effect=_BillingHTTPStatusError("request failed", status_code),
):
with pytest.raises(BillingUpstreamInvalidResponseError):
BillingService._send_billing_portal_request("/test", params={})
@pytest.mark.parametrize(
"status_code",
[httpx.codes.REQUEST_TIMEOUT, httpx.codes.TOO_MANY_REQUESTS, httpx.codes.INTERNAL_SERVER_ERROR],
)
def test_retryable_http_response_is_unavailable(self, status_code: int) -> None:
with patch.object(
BillingService,
"_send_request",
side_effect=_BillingHTTPStatusError("request failed", status_code),
):
with pytest.raises(BillingUpstreamUnavailableError):
BillingService._send_billing_portal_request("/test", params={})
def test_transport_failure_is_unavailable(self) -> None:
with patch.object(
BillingService,
"_send_request",
side_effect=httpx.RequestError("network error"),
):
with pytest.raises(BillingUpstreamUnavailableError):
BillingService._send_billing_portal_request("/test", params={})
@pytest.mark.parametrize(
"decode_error",
[
json.JSONDecodeError("Expecting value", "", 0),
UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte"),
],
)
def test_invalid_payload_is_invalid_upstream_response(self, decode_error: Exception) -> None:
with patch.object(BillingService, "_send_request", side_effect=decode_error):
with pytest.raises(BillingUpstreamInvalidResponseError):
BillingService._send_billing_portal_request("/test", params={})
def test_unknown_error_is_not_reclassified(self) -> None:
with patch.object(BillingService, "_send_request", side_effect=RuntimeError("programming error")):
with pytest.raises(RuntimeError, match="programming error"):
BillingService._send_billing_portal_request("/test", params={})
def test_unknown_value_error_is_not_exposed_as_invalid_request(self) -> None:
original_error = ValueError("programming error")
with patch.object(BillingService, "_send_request", side_effect=original_error):
with pytest.raises(RuntimeError, match="Unexpected billing service value error") as exc_info:
BillingService._send_billing_portal_request("/test", params={})
assert exc_info.value.__cause__ is original_error
class TestBillingServiceSubscriptionInfo:
"""Unit tests for subscription tier and billing info retrieval.
@@ -587,23 +631,22 @@ class TestBillingServiceSubscriptionInfo:
assert result["limit"] == 100
assert result["subscription_plan"] == CloudPlan.PROFESSIONAL
def test_get_subscription_payment_link(self, mock_send_request):
def test_get_subscription_payment_link(self):
"""Test subscription payment link generation."""
# Arrange
plan = "professional"
interval = "monthly"
interval = "month"
email = "user@example.com"
tenant_id = "tenant-123"
expected_response = {"payment_link": "https://payment.example.com/checkout"}
mock_send_request.return_value = expected_response
# Act
result = BillingService.get_subscription(plan, interval, email, tenant_id)
expected_response = {"url": "https://payment.example.com/checkout"}
with patch.object(
BillingService, "_send_billing_portal_request", return_value=expected_response
) as send_request:
result = BillingService.get_subscription(plan, interval, email, tenant_id)
# Assert
assert result == expected_response
mock_send_request.assert_called_once_with(
"GET",
send_request.assert_called_once_with(
"/subscription/payment-link",
params={"plan": plan, "interval": interval, "prefilled_email": email, "tenant_id": tenant_id},
)
@@ -634,22 +677,20 @@ class TestBillingServiceSubscriptionInfo:
},
)
def test_get_invoices(self, mock_send_request):
def test_get_invoices(self):
"""Test invoice retrieval."""
# Arrange
email = "user@example.com"
tenant_id = "tenant-123"
expected_response = {"invoices": [{"id": "inv-1", "amount": 100}]}
mock_send_request.return_value = expected_response
# Act
result = BillingService.get_invoices(email, tenant_id)
expected_response = {"url": "https://payment.example.com/invoices"}
with patch.object(
BillingService, "_send_billing_portal_request", return_value=expected_response
) as send_request:
result = BillingService.get_invoices(email, tenant_id)
# Assert
assert result == expected_response
mock_send_request.assert_called_once_with(
"GET", "/invoices", params={"prefilled_email": email, "tenant_id": tenant_id}
)
send_request.assert_called_once_with("/invoices", params={"prefilled_email": email, "tenant_id": tenant_id})
class TestBillingServiceUsageCalculation:
@@ -1348,50 +1389,6 @@ class TestBillingServiceAccountManagement:
"POST", "/account/delete-feedback", json={"email": email, "feedback": feedback}
)
@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, TenantAccountJoin)], indirect=True)
def test_is_tenant_owner_or_admin_owner(self, sqlite_session: Session):
"""Test tenant owner/admin check for owner role."""
# Arrange
current_user = _persist_membership(sqlite_session, role=TenantAccountRole.OWNER)
# Act - should not raise exception
BillingService.is_tenant_owner_or_admin(current_user, session=sqlite_session)
@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, TenantAccountJoin)], indirect=True)
def test_is_tenant_owner_or_admin_admin(self, sqlite_session: Session):
"""Test tenant owner/admin check for admin role."""
# Arrange
current_user = _persist_membership(sqlite_session, role=TenantAccountRole.ADMIN)
# Act - should not raise exception
BillingService.is_tenant_owner_or_admin(current_user, session=sqlite_session)
@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, TenantAccountJoin)], indirect=True)
def test_is_tenant_owner_or_admin_normal_user_raises_error(self, sqlite_session: Session):
"""Test tenant owner/admin check raises error for normal user."""
# Arrange
current_user = _persist_membership(sqlite_session, role=TenantAccountRole.NORMAL)
# Act & Assert
with pytest.raises(ValueError) as exc_info:
BillingService.is_tenant_owner_or_admin(current_user, session=sqlite_session)
assert "Only team owner or team admin can perform this action" in str(exc_info.value)
@pytest.mark.parametrize("sqlite_session", [(Account, Tenant, TenantAccountJoin)], indirect=True)
def test_is_tenant_owner_or_admin_no_join_raises_error(self, sqlite_session: Session):
"""Test tenant owner/admin check raises error when join not found."""
# Arrange
current_user = _persist_membership(
sqlite_session,
role=None,
add_other_tenant_membership=True,
)
# Act & Assert
with pytest.raises(ValueError) as exc_info:
BillingService.is_tenant_owner_or_admin(current_user, session=sqlite_session)
assert "Tenant account join not found" in str(exc_info.value)
class TestBillingServiceCacheManagement:
"""Unit tests for billing cache management.
@@ -1540,8 +1537,8 @@ class TestBillingServiceEdgeCases:
"""Test subscription payment link with empty optional parameters."""
# Arrange
plan = "professional"
interval = "yearly"
expected_response = {"payment_link": "https://payment.example.com/checkout"}
interval = "year"
expected_response = {"url": "https://payment.example.com/checkout"}
mock_send_request.return_value = expected_response
# Act - empty email and tenant_id
@@ -1558,7 +1555,7 @@ class TestBillingServiceEdgeCases:
def test_get_invoices_with_empty_params(self, mock_send_request):
"""Test invoice retrieval with empty parameters."""
# Arrange
expected_response = {"invoices": []}
expected_response = {"url": "https://payment.example.com/invoices"}
mock_send_request.return_value = expected_response
# Act
@@ -1566,7 +1563,6 @@ class TestBillingServiceEdgeCases:
# Assert
assert result == expected_response
assert result["invoices"] == []
def test_refund_with_invalid_history_id_format(self, mock_send_request):
"""Test refund with various history ID formats."""
@@ -1871,9 +1867,9 @@ class TestBillingServiceIntegrationScenarios:
assert current_info["subscription"]["plan"] == "sandbox"
# Step 2: Get payment link for upgrade
mock_send_request.return_value = {"payment_link": "https://payment.example.com/upgrade"}
payment_link = BillingService.get_subscription("professional", "monthly", "user@example.com", tenant_id)
assert "payment_link" in payment_link
mock_send_request.return_value = {"url": "https://payment.example.com/upgrade"}
payment_link = BillingService.get_subscription("professional", "month", "user@example.com", tenant_id)
assert "url" in payment_link
# Step 3: Verify new rate limits after upgrade
mock_send_request.return_value = {"limit": 100, "subscription_plan": CloudPlan.PROFESSIONAL}
packages/contracts/generated/api/console/billing/types.gen.ts
+36
View File
@@ -8,6 +8,18 @@ export type BillingInvoiceResponse = {
url: string
}
export type BillingOperationFailedErrorResponse = {
code: 'billing_operation_failed'
message: string
status: 502
}
export type BillingUnavailableErrorResponse = {
code: 'billing_unavailable'
message: string
status: 503
}
export type PartnerTenantsPayload = {
click_id: string
}
@@ -20,6 +32,12 @@ export type BillingSubscriptionResponse = {
url: string
}
export type BillingUnprocessableEntityErrorResponse = {
code: 'unprocessable_entity'
message: string
status: 422
}
export type GetBillingInvoicesData = {
body?: never
path?: never
@@ -27,6 +45,14 @@ export type GetBillingInvoicesData = {
url: '/billing/invoices'
}
export type GetBillingInvoicesErrors = {
403: unknown
502: BillingOperationFailedErrorResponse
503: BillingUnavailableErrorResponse
}
export type GetBillingInvoicesError = GetBillingInvoicesErrors[keyof GetBillingInvoicesErrors]
export type GetBillingInvoicesResponses = {
200: BillingInvoiceResponse
}
@@ -64,6 +90,16 @@ export type GetBillingSubscriptionData = {
url: '/billing/subscription'
}
export type GetBillingSubscriptionErrors = {
403: unknown
422: BillingUnprocessableEntityErrorResponse
502: BillingOperationFailedErrorResponse
503: BillingUnavailableErrorResponse
}
export type GetBillingSubscriptionError =
GetBillingSubscriptionErrors[keyof GetBillingSubscriptionErrors]
export type GetBillingSubscriptionResponses = {
200: BillingSubscriptionResponse
}
packages/contracts/generated/api/console/billing/zod.gen.ts
+27
View File
@@ -9,6 +9,24 @@ export const zBillingInvoiceResponse = z.object({
url: z.string(),
})
/**
* BillingOperationFailedErrorResponse
*/
export const zBillingOperationFailedErrorResponse = z.object({
code: z.literal('billing_operation_failed'),
message: z.string(),
status: z.literal(502),
})
/**
* BillingUnavailableErrorResponse
*/
export const zBillingUnavailableErrorResponse = z.object({
code: z.literal('billing_unavailable'),
message: z.string(),
status: z.literal(503),
})
/**
* PartnerTenantsPayload
*/
@@ -28,6 +46,15 @@ export const zBillingSubscriptionResponse = z.object({
url: z.string(),
})
/**
* BillingUnprocessableEntityErrorResponse
*/
export const zBillingUnprocessableEntityErrorResponse = z.object({
code: z.literal('unprocessable_entity'),
message: z.string(),
status: z.literal(422),
})
/**
* Success
*/
packages/contracts/generated/api/console/workspaces/orpc.gen.ts
+4
View File
@@ -1098,8 +1098,12 @@ export const summary = {
get: get13,
}
/**
* @deprecated
*/
export const get14 = oc
.route({
deprecated: true,
inputStructure: 'detailed',
method: 'GET',
operationId: 'getWorkspacesCurrentModelProvidersByProviderCheckoutUrl',