mirror of
https://github.com/langgenius/dify.git
synced 2026-03-26 22:52:38 +08:00
refactor(workspace): optimize /workspaces plan resolution for SaaS and enterprise with resilient fallback (#33788)
This commit is contained in:
L1nSn0w
GitHub
parent
3ce43724df
commit
a1af085736
@ -7,6 +7,7 @@ from sqlalchemy import select
|
||||
from werkzeug.exceptions import Unauthorized
|
||||
|
||||
import services
|
||||
from configs import dify_config
|
||||
from controllers.common.errors import (
|
||||
FilenameNotExistsError,
|
||||
FileTooLargeError,
|
||||
@ -29,6 +30,7 @@ from libs.helper import TimestampField
|
||||
from libs.login import current_account_with_tenant, login_required
|
||||
from models.account import Tenant, TenantStatus
|
||||
from services.account_service import TenantService
|
||||
from services.billing_service import BillingService, SubscriptionPlan
|
||||
from services.enterprise.enterprise_service import EnterpriseService
|
||||
from services.feature_service import FeatureService
|
||||
from services.file_service import FileService
|
||||
@ -108,9 +110,29 @@ class TenantListApi(Resource):
|
||||
current_user, current_tenant_id = current_account_with_tenant()
|
||||
tenants = TenantService.get_join_tenants(current_user)
|
||||
tenant_dicts = []
|
||||
is_enterprise_only = dify_config.ENTERPRISE_ENABLED and not dify_config.BILLING_ENABLED
|
||||
is_saas = dify_config.EDITION == "CLOUD" and dify_config.BILLING_ENABLED
|
||||
tenant_plans: dict[str, SubscriptionPlan] = {}
|
||||
|
||||
if is_saas:
|
||||
tenant_ids = [tenant.id for tenant in tenants]
|
||||
if tenant_ids:
|
||||
tenant_plans = BillingService.get_plan_bulk(tenant_ids)
|
||||
if not tenant_plans:
|
||||
logger.warning("get_plan_bulk returned empty result, falling back to legacy feature path")
|
||||
|
||||
for tenant in tenants:
|
||||
features = FeatureService.get_features(tenant.id)
|
||||
plan: str = CloudPlan.SANDBOX
|
||||
if is_saas:
|
||||
tenant_plan = tenant_plans.get(tenant.id)
|
||||
if tenant_plan:
|
||||
plan = tenant_plan["plan"] or CloudPlan.SANDBOX
|
||||
else:
|
||||
features = FeatureService.get_features(tenant.id)
|
||||
plan = features.billing.subscription.plan or CloudPlan.SANDBOX
|
||||
elif not is_enterprise_only:
|
||||
features = FeatureService.get_features(tenant.id)
|
||||
plan = features.billing.subscription.plan or CloudPlan.SANDBOX
|
||||
|
||||
# Create a dictionary with tenant attributes
|
||||
tenant_dict = {
|
||||
@ -118,7 +140,7 @@ class TenantListApi(Resource):
|
||||
"name": tenant.name,
|
||||
"status": tenant.status,
|
||||
"created_at": tenant.created_at,
|
||||
"plan": features.billing.subscription.plan if features.billing.enabled else CloudPlan.SANDBOX,
|
||||
"plan": plan,
|
||||
"current": tenant.id == current_tenant_id if current_tenant_id else False,
|
||||
}
|
||||
|
||||
|
||||
@ -36,7 +36,115 @@ def unwrap(func):
|
||||
|
||||
|
||||
class TestTenantListApi:
|
||||
def test_get_success(self, app):
|
||||
def test_get_success_saas_path(self, app):
|
||||
api = TenantListApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
tenant1 = MagicMock(
|
||||
id="t1",
|
||||
name="Tenant 1",
|
||||
status="active",
|
||||
created_at=datetime.utcnow(),
|
||||
)
|
||||
tenant2 = MagicMock(
|
||||
id="t2",
|
||||
name="Tenant 2",
|
||||
status="active",
|
||||
created_at=datetime.utcnow(),
|
||||
)
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.current_account_with_tenant", return_value=(MagicMock(), "t1")
|
||||
),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.TenantService.get_join_tenants",
|
||||
return_value=[tenant1, tenant2],
|
||||
),
|
||||
patch("controllers.console.workspace.workspace.dify_config.ENTERPRISE_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.BILLING_ENABLED", True),
|
||||
patch("controllers.console.workspace.workspace.dify_config.EDITION", "CLOUD"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.BillingService.get_plan_bulk",
|
||||
return_value={
|
||||
"t1": {"plan": CloudPlan.TEAM, "expiration_date": 0},
|
||||
"t2": {"plan": CloudPlan.PROFESSIONAL, "expiration_date": 0},
|
||||
},
|
||||
) as get_plan_bulk_mock,
|
||||
patch("controllers.console.workspace.workspace.FeatureService.get_features") as get_features_mock,
|
||||
):
|
||||
result, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert len(result["workspaces"]) == 2
|
||||
assert result["workspaces"][0]["current"] is True
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.TEAM
|
||||
assert result["workspaces"][1]["plan"] == CloudPlan.PROFESSIONAL
|
||||
get_plan_bulk_mock.assert_called_once_with(["t1", "t2"])
|
||||
get_features_mock.assert_not_called()
|
||||
|
||||
def test_get_saas_path_partial_fallback_does_not_gate_plan_on_billing_enabled(self, app):
|
||||
"""Bulk omits a tenant: resolve plan via subscription.plan only; billing.enabled is not used.
|
||||
|
||||
billing.enabled is mocked False to prove the endpoint does not gate on it for this path
|
||||
(SaaS contract treats enabled as on; display follows subscription.plan).
|
||||
"""
|
||||
api = TenantListApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
tenant1 = MagicMock(
|
||||
id="t1",
|
||||
name="Tenant 1",
|
||||
status="active",
|
||||
created_at=datetime.utcnow(),
|
||||
)
|
||||
tenant2 = MagicMock(
|
||||
id="t2",
|
||||
name="Tenant 2",
|
||||
status="active",
|
||||
created_at=datetime.utcnow(),
|
||||
)
|
||||
|
||||
features_t2 = MagicMock()
|
||||
features_t2.billing.enabled = False
|
||||
features_t2.billing.subscription.plan = CloudPlan.PROFESSIONAL
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.current_account_with_tenant", return_value=(MagicMock(), "t1")
|
||||
),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.TenantService.get_join_tenants",
|
||||
return_value=[tenant1, tenant2],
|
||||
),
|
||||
patch("controllers.console.workspace.workspace.dify_config.ENTERPRISE_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.BILLING_ENABLED", True),
|
||||
patch("controllers.console.workspace.workspace.dify_config.EDITION", "CLOUD"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.BillingService.get_plan_bulk",
|
||||
return_value={"t1": {"plan": CloudPlan.TEAM, "expiration_date": 0}},
|
||||
) as get_plan_bulk_mock,
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.FeatureService.get_features",
|
||||
return_value=features_t2,
|
||||
) as get_features_mock,
|
||||
):
|
||||
result, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.TEAM
|
||||
assert result["workspaces"][1]["plan"] == CloudPlan.PROFESSIONAL
|
||||
get_plan_bulk_mock.assert_called_once_with(["t1", "t2"])
|
||||
get_features_mock.assert_called_once_with("t2")
|
||||
|
||||
def test_get_saas_path_falls_back_to_legacy_feature_path_on_bulk_error(self, app):
|
||||
"""Test fallback to FeatureService when bulk billing returns empty result.
|
||||
|
||||
BillingService.get_plan_bulk catches exceptions internally and returns empty dict,
|
||||
so we simulate the real failure mode by returning empty dict for non-empty input.
|
||||
"""
|
||||
api = TenantListApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
@ -54,27 +162,41 @@ class TestTenantListApi:
|
||||
)
|
||||
|
||||
features = MagicMock()
|
||||
features.billing.enabled = True
|
||||
features.billing.subscription.plan = CloudPlan.SANDBOX
|
||||
features.billing.enabled = False
|
||||
features.billing.subscription.plan = CloudPlan.TEAM
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.current_account_with_tenant", return_value=(MagicMock(), "t1")
|
||||
"controllers.console.workspace.workspace.current_account_with_tenant", return_value=(MagicMock(), "t2")
|
||||
),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.TenantService.get_join_tenants",
|
||||
return_value=[tenant1, tenant2],
|
||||
),
|
||||
patch("controllers.console.workspace.workspace.FeatureService.get_features", return_value=features),
|
||||
patch("controllers.console.workspace.workspace.dify_config.ENTERPRISE_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.BILLING_ENABLED", True),
|
||||
patch("controllers.console.workspace.workspace.dify_config.EDITION", "CLOUD"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.BillingService.get_plan_bulk",
|
||||
return_value={}, # Simulates real failure: empty result for non-empty input
|
||||
) as get_plan_bulk_mock,
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.FeatureService.get_features",
|
||||
return_value=features,
|
||||
) as get_features_mock,
|
||||
patch("controllers.console.workspace.workspace.logger.warning") as logger_warning_mock,
|
||||
):
|
||||
result, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert len(result["workspaces"]) == 2
|
||||
assert result["workspaces"][0]["current"] is True
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.TEAM
|
||||
assert result["workspaces"][1]["plan"] == CloudPlan.TEAM
|
||||
get_plan_bulk_mock.assert_called_once_with(["t1", "t2"])
|
||||
assert get_features_mock.call_count == 2
|
||||
logger_warning_mock.assert_called_once()
|
||||
|
||||
def test_get_billing_disabled(self, app):
|
||||
def test_get_billing_disabled_community_path(self, app):
|
||||
api = TenantListApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
@ -87,6 +209,7 @@ class TestTenantListApi:
|
||||
|
||||
features = MagicMock()
|
||||
features.billing.enabled = False
|
||||
features.billing.subscription.plan = CloudPlan.SANDBOX
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces"),
|
||||
@ -98,15 +221,83 @@ class TestTenantListApi:
|
||||
"controllers.console.workspace.workspace.TenantService.get_join_tenants",
|
||||
return_value=[tenant],
|
||||
),
|
||||
patch("controllers.console.workspace.workspace.dify_config.ENTERPRISE_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.BILLING_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.EDITION", "SELF_HOSTED"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.FeatureService.get_features",
|
||||
return_value=features,
|
||||
),
|
||||
) as get_features_mock,
|
||||
):
|
||||
result, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.SANDBOX
|
||||
get_features_mock.assert_called_once_with("t1")
|
||||
|
||||
def test_get_enterprise_only_skips_feature_service(self, app):
|
||||
api = TenantListApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
tenant1 = MagicMock(
|
||||
id="t1",
|
||||
name="Tenant 1",
|
||||
status="active",
|
||||
created_at=datetime.utcnow(),
|
||||
)
|
||||
tenant2 = MagicMock(
|
||||
id="t2",
|
||||
name="Tenant 2",
|
||||
status="active",
|
||||
created_at=datetime.utcnow(),
|
||||
)
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.current_account_with_tenant", return_value=(MagicMock(), "t2")
|
||||
),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.TenantService.get_join_tenants",
|
||||
return_value=[tenant1, tenant2],
|
||||
),
|
||||
patch("controllers.console.workspace.workspace.dify_config.ENTERPRISE_ENABLED", True),
|
||||
patch("controllers.console.workspace.workspace.dify_config.BILLING_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.EDITION", "SELF_HOSTED"),
|
||||
patch("controllers.console.workspace.workspace.FeatureService.get_features") as get_features_mock,
|
||||
):
|
||||
result, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.SANDBOX
|
||||
assert result["workspaces"][1]["plan"] == CloudPlan.SANDBOX
|
||||
assert result["workspaces"][0]["current"] is False
|
||||
assert result["workspaces"][1]["current"] is True
|
||||
get_features_mock.assert_not_called()
|
||||
|
||||
def test_get_enterprise_only_with_empty_tenants(self, app):
|
||||
api = TenantListApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
with (
|
||||
app.test_request_context("/workspaces"),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.current_account_with_tenant", return_value=(MagicMock(), None)
|
||||
),
|
||||
patch(
|
||||
"controllers.console.workspace.workspace.TenantService.get_join_tenants",
|
||||
return_value=[],
|
||||
),
|
||||
patch("controllers.console.workspace.workspace.dify_config.ENTERPRISE_ENABLED", True),
|
||||
patch("controllers.console.workspace.workspace.dify_config.BILLING_ENABLED", False),
|
||||
patch("controllers.console.workspace.workspace.dify_config.EDITION", "SELF_HOSTED"),
|
||||
patch("controllers.console.workspace.workspace.FeatureService.get_features") as get_features_mock,
|
||||
):
|
||||
result, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert result["workspaces"] == []
|
||||
get_features_mock.assert_not_called()
|
||||
|
||||
|
||||
class TestWorkspaceListApi:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user