From eba589db8975e6a8ae3a9ead9985f708e106d3aa Mon Sep 17 00:00:00 2001 From: yyh <92089059+lyzno1@users.noreply.github.com> Date: Mon, 7 Sep 2026 08:29:57 +0000 Subject: [PATCH] refactor(api): remove legacy billing enabled state (#41917) --- api/controllers/console/wraps.py | 65 ++++++++-------- api/controllers/openapi/workspaces.py | 3 +- api/controllers/service_api/app/completion.py | 2 +- api/controllers/service_api/app/workflow.py | 2 +- api/controllers/service_api/wraps.py | 45 +++++------ api/controllers/web/site.py | 3 +- api/extensions/ext_application_services.py | 1 + api/openapi/markdown/console-openapi.md | 1 - api/services/annotation_service.py | 5 +- api/services/billing_service.py | 1 - api/services/dataset_service.py | 21 +++--- api/services/document_indexing_proxy/base.py | 15 ++-- api/services/entities/feature_entities.py | 7 -- api/services/feature_service.py | 5 +- .../rag_pipeline/rag_pipeline_task_proxy.py | 14 ++-- .../vector_space_admission_service.py | 16 ++-- api/services/web_app_runtime_query_service.py | 5 +- api/services/workspace_service.py | 8 +- api/tasks/document_indexing_task.py | 32 ++++---- api/tasks/duplicate_document_indexing_task.py | 36 ++++----- api/tasks/retry_document_indexing_task.py | 34 +++++---- .../sync_website_document_indexing_task.py | 24 +++--- .../services/test_annotation_service.py | 8 +- .../test_api_based_extension_service.py | 3 - .../test_dataset_service_permissions.py | 16 ++-- .../services/test_feature_service.py | 22 ------ .../services/test_message_service.py | 1 - .../services/test_workspace_service.py | 3 +- .../tasks/test_dataset_indexing_task.py | 18 ++--- .../tasks/test_document_indexing_task.py | 31 ++++---- .../test_duplicate_document_indexing_task.py | 27 ++++--- .../console/app/test_annotation_security.py | 8 +- .../console/test_workflow_run_archive.py | 1 - .../console/test_workspace_members.py | 1 - .../controllers/console/test_wraps.py | 17 ++--- .../openapi/test_workspaces_members.py | 15 ++-- .../service_api/app/test_completion.py | 16 ++-- .../service_api/app/test_workflow.py | 16 ++-- .../dataset/test_dataset_segment.py | 4 - .../service_api/dataset/test_document.py | 6 +- .../controllers/service_api/test_wraps.py | 16 ++-- .../services/dataset_service_test_helpers.py | 3 +- .../test_rag_pipeline_task_proxy.py | 19 ++--- .../services/test_annotation_service.py | 31 ++++---- .../services/test_batch_indexing_base.py | 30 ++++---- .../services/test_billing_service.py | 8 -- .../services/test_dataset_service_document.py | 74 +++++++++++++------ .../test_dataset_service_lock_not_owned.py | 3 +- .../test_document_indexing_task_proxy.py | 42 +++++------ ..._duplicate_document_indexing_task_proxy.py | 49 ++++++------ .../services/test_feature_entities.py | 10 ++- ...ture_service_human_input_email_delivery.py | 16 ---- ...ature_service_knowledge_file_size_limit.py | 18 ++--- .../test_feature_service_vector_space.py | 1 - .../services/test_rag_pipeline_task_proxy.py | 44 ++++++----- .../test_vector_space_admission_service.py | 15 +++- .../test_web_app_runtime_query_service.py | 24 +++++- .../services/test_workspace_credit_pool.py | 2 - .../services/test_workspace_service.py | 2 - .../tasks/test_dataset_indexing_task.py | 20 ++--- .../test_retry_document_indexing_task.py | 4 +- ...est_sync_website_document_indexing_task.py | 4 +- .../api/console/features/types.gen.ts | 1 - .../generated/api/console/features/zod.gen.ts | 6 +- .../billing/pricing-modal-flow.test.tsx | 1 - .../billing/utils/__tests__/index.spec.ts | 2 - 66 files changed, 485 insertions(+), 518 deletions(-) diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 749d3243750..813c8a75a41 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -163,7 +163,7 @@ def cloud_edition_billing_paid_plan_required[**P, R](view: Callable[P, R]) -> Ca def decorated(*args: P.args, **kwargs: P.kwargs): _, current_tenant_id = current_account_with_tenant() billing_info = BillingService.get_info(current_tenant_id, exclude_vector_space=True) - if not billing_info["enabled"] or billing_info["subscription"]["plan"] not in ( + if billing_info["subscription"]["plan"] not in ( CloudPlan.PROFESSIONAL, CloudPlan.TEAM, ): @@ -178,10 +178,10 @@ def cloud_edition_billing_resource_check[**P, R](resource: str) -> Callable[[Cal @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): _, current_tenant_id = current_account_with_tenant() - if resource == "vector_space": - if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: - return view(*args, **kwargs) + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: + return view(*args, **kwargs) + if resource == "vector_space": vector_space = application_services().feature_queries.get_workspace_vector_space(current_tenant_id) if 0 < vector_space.limit <= vector_space.size: abort( @@ -191,30 +191,26 @@ def cloud_edition_billing_resource_check[**P, R](resource: str) -> Callable[[Cal return view(*args, **kwargs) features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) - if features.billing.enabled: - members = features.members - apps = features.apps - documents_upload_quota = features.documents_upload_quota - annotation_quota_limit = features.annotation_quota_limit - if resource == "members" and 0 < members.limit <= members.size: - abort(403, "The number of members has reached the limit of your subscription.") - elif resource == "apps" and 0 < apps.limit <= apps.size: - abort(403, "The number of apps has reached the limit of your subscription.") - elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size: - # The api of file upload is used in the multiple places, - # so we need to check the source of the request from datasets - source = request.args.get("source") or request.form.get("source") - if source == "datasets": - abort(403, "The number of documents has reached the limit of your subscription.") - else: - return view(*args, **kwargs) - elif resource == "workspace_custom" and not features.can_replace_logo: - abort(403, "The workspace custom feature has reached the limit of your subscription.") - elif resource == "annotation" and 0 < annotation_quota_limit.limit < annotation_quota_limit.size: - abort(403, "The annotation quota has reached the limit of your subscription.") + members = features.members + apps = features.apps + documents_upload_quota = features.documents_upload_quota + annotation_quota_limit = features.annotation_quota_limit + if resource == "members" and 0 < members.limit <= members.size: + abort(403, "The number of members has reached the limit of your subscription.") + elif resource == "apps" and 0 < apps.limit <= apps.size: + abort(403, "The number of apps has reached the limit of your subscription.") + elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size: + # The api of file upload is used in the multiple places, + # so we need to check the source of the request from datasets + source = request.args.get("source") or request.form.get("source") + if source == "datasets": + abort(403, "The number of documents has reached the limit of your subscription.") else: return view(*args, **kwargs) - + elif resource == "workspace_custom" and not features.can_replace_logo: + abort(403, "The workspace custom feature has reached the limit of your subscription.") + elif resource == "annotation" and 0 < annotation_quota_limit.limit < annotation_quota_limit.size: + abort(403, "The annotation quota has reached the limit of your subscription.") return view(*args, **kwargs) return decorated @@ -229,16 +225,15 @@ def cloud_edition_billing_knowledge_limit_check[**P, R]( @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): _, current_tenant_id = current_account_with_tenant() + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD or resource != "add_segment": + return view(*args, **kwargs) + features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) - if features.billing.enabled: - if resource == "add_segment": - if features.billing.subscription.plan == CloudPlan.SANDBOX: - abort( - 403, - "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan.", - ) - else: - return view(*args, **kwargs) + if features.billing.subscription.plan == CloudPlan.SANDBOX: + abort( + 403, + "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan.", + ) return view(*args, **kwargs) diff --git a/api/controllers/openapi/workspaces.py b/api/controllers/openapi/workspaces.py index b5fec3bd4ca..6bff0349f78 100644 --- a/api/controllers/openapi/workspaces.py +++ b/api/controllers/openapi/workspaces.py @@ -37,6 +37,7 @@ from controllers.openapi._models import ( ) from controllers.openapi.auth.composition import auth_router from controllers.openapi.auth.data import AuthData +from enums import DeploymentEdition from libs.oauth_bearer import Scope, TokenType from models import Account, Tenant, TenantAccountJoin from models.account import TenantAccountRole, TenantStatus @@ -82,7 +83,7 @@ def _load_account(session: Session, account_id: object) -> Account: def _check_member_invite_quota(tenant_id: str) -> None: features = FeatureService.get_features(tenant_id) - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: members = features.members if 0 < members.limit <= members.size: raise MemberLimitExceeded() diff --git a/api/controllers/service_api/app/completion.py b/api/controllers/service_api/app/completion.py index 50bb033164b..5365a1ab551 100644 --- a/api/controllers/service_api/app/completion.py +++ b/api/controllers/service_api/app/completion.py @@ -399,7 +399,7 @@ class ChatApi(Resource): and dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD ): billing_info = BillingService.get_info(app_model.tenant_id, exclude_vector_space=True) - if billing_info["enabled"] and billing_info["subscription"]["plan"] == CloudPlan.SANDBOX: + if billing_info["subscription"]["plan"] == CloudPlan.SANDBOX: raise WorkflowVersionExecutionNotAllowedError() external_trace_id = get_external_trace_id(request) diff --git a/api/controllers/service_api/app/workflow.py b/api/controllers/service_api/app/workflow.py index d40438e6fa3..939058ba271 100644 --- a/api/controllers/service_api/app/workflow.py +++ b/api/controllers/service_api/app/workflow.py @@ -476,7 +476,7 @@ class WorkflowRunByIdApi(Resource): if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: billing_info = BillingService.get_info(app_model.tenant_id, exclude_vector_space=True) - if billing_info["enabled"] and billing_info["subscription"]["plan"] == CloudPlan.SANDBOX: + if billing_info["subscription"]["plan"] == CloudPlan.SANDBOX: raise WorkflowVersionExecutionNotAllowedError() payload = WorkflowRunPayload.model_validate(omit_trace_session_id_from_payload(service_api_ns.payload) or {}) diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py index 9d29d498a33..77201f96f0c 100644 --- a/api/controllers/service_api/wraps.py +++ b/api/controllers/service_api/wraps.py @@ -194,14 +194,14 @@ def cloud_edition_billing_resource_check[**P, R]( @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): api_token = validate_and_get_api_token(api_token_type) - if resource == "vector_space": - if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: - return view(*args, **kwargs) + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: + return view(*args, **kwargs) + if resource == "vector_space": vector_space = application_services().feature_queries.get_workspace_vector_space(api_token.tenant_id) if vector_space.usage_unknown: features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True) - if features.billing.enabled and features.billing.subscription.plan == CloudPlan.SANDBOX: + if features.billing.subscription.plan == CloudPlan.SANDBOX: raise ServiceUnavailable( "Unable to verify vector space usage right now. Please try again later." ) @@ -211,20 +211,16 @@ def cloud_edition_billing_resource_check[**P, R]( features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True) - if features.billing.enabled: - members = features.members - apps = features.apps - documents_upload_quota = features.documents_upload_quota - - if resource == "members" and 0 < members.limit <= members.size: - raise Forbidden("The number of members has reached the limit of your subscription.") - elif resource == "apps" and 0 < apps.limit <= apps.size: - raise Forbidden("The number of apps has reached the limit of your subscription.") - elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size: - raise Forbidden("The number of documents has reached the limit of your subscription.") - else: - return view(*args, **kwargs) + members = features.members + apps = features.apps + documents_upload_quota = features.documents_upload_quota + if resource == "members" and 0 < members.limit <= members.size: + raise Forbidden("The number of members has reached the limit of your subscription.") + elif resource == "apps" and 0 < apps.limit <= apps.size: + raise Forbidden("The number of apps has reached the limit of your subscription.") + elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size: + raise Forbidden("The number of documents has reached the limit of your subscription.") return view(*args, **kwargs) if resource == "vector_space": @@ -245,15 +241,14 @@ def cloud_edition_billing_knowledge_limit_check[**P, R]( @wraps(view) def decorated(*args: P.args, **kwargs: P.kwargs): api_token = validate_and_get_api_token(api_token_type) + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD or resource != "add_segment": + return view(*args, **kwargs) + features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True) - if features.billing.enabled: - if resource == "add_segment": - if features.billing.subscription.plan == CloudPlan.SANDBOX: - raise Forbidden( - "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan." - ) - else: - return view(*args, **kwargs) + if features.billing.subscription.plan == CloudPlan.SANDBOX: + raise Forbidden( + "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan." + ) return view(*args, **kwargs) diff --git a/api/controllers/web/site.py b/api/controllers/web/site.py index 54914b9bfb4..7b2afe16bcc 100644 --- a/api/controllers/web/site.py +++ b/api/controllers/web/site.py @@ -7,6 +7,7 @@ from configs import dify_config from controllers.common.schema import register_response_schema_models from controllers.web import web_ns from controllers.web.wraps import WebApiResource +from enums import DeploymentEdition from extensions.ext_application_services import application_services from fields.base import ResponseModel from libs.helper import build_icon_url, dump_response @@ -102,7 +103,7 @@ class WebAppSiteResponse(ResponseModel): site_response = WebSiteResponse.model_validate(site, from_attributes=True) site_response.icon_url = icon_url if icon_url is not None else build_icon_url(site.icon_type, site.icon) - if features.billing.enabled and not features.webapp_copyright_enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD and not features.webapp_copyright_enabled: site_response.copyright = None site_response.input_placeholder = None diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index a8050b74ad9..9c57bb557f6 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -635,6 +635,7 @@ def build_application_services( file_service=file_service, workspace_features=feature_gateway.get_workspace_features, files_url=dify_config.FILES_URL, + deployment_edition=deployment_edition, ), explore_banner_queries=ExploreBannerQueryService( banners=ExploreBannerQueryRepository(session_factory=database_client), diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index 74a65d44382..151dcef37f2 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -16118,7 +16118,6 @@ ExporleBanner status | Name | Type | Description | Required | | ---- | ---- | ----------- | -------- | -| enabled | boolean | Deprecated. Use system features deployment_edition to determine the product edition. | Yes | | subscription | [SubscriptionModel](#subscriptionmodel) | | Yes | #### BillingOperationFailedErrorResponse diff --git a/api/services/annotation_service.py b/api/services/annotation_service.py index ac0de983eb9..7efa26803a2 100644 --- a/api/services/annotation_service.py +++ b/api/services/annotation_service.py @@ -9,6 +9,7 @@ from werkzeug.datastructures import FileStorage from werkzeug.exceptions import NotFound from core.helper.csv_sanitizer import CSVSanitizer +from enums import DeploymentEdition from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from libs.login import current_account_with_tenant @@ -532,8 +533,8 @@ class AppAnnotationService: ) # Check annotation quota limit - features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True) annotation_quota_limit = features.annotation_quota_limit if annotation_quota_limit.limit < len(result) + annotation_quota_limit.size: raise ValueError("The number of annotations exceeds the limit of your subscription.") diff --git a/api/services/billing_service.py b/api/services/billing_service.py index f89547f4558..25965bbccfe 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -173,7 +173,6 @@ class BillingInfo(TypedDict): 3. To preserve compatibility, always keep non-strict mode here and avoid strict mode """ - enabled: bool subscription: _BillingSubscription members: _BillingQuota apps: _BillingQuota diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index 37296e4b1ec..fe36c270496 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -23,7 +23,7 @@ from core.model_manager import ModelManager from core.rag.index_processor.constant.built_in_field import BuiltInField from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.retrieval.retrieval_methods import RetrievalMethod -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from events.dataset_event import dataset_was_deleted from events.document_event import document_was_deleted from extensions.ext_redis import redis_client @@ -1455,8 +1455,11 @@ class DatasetService: @staticmethod def get_dataset_auto_disable_logs(dataset_ref: DatasetRef, session: Session) -> AutoDisableLogsDict: + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: + return {"document_ids": [], "count": 0} + features = FeatureService.get_features(dataset_ref.tenant_id, exclude_vector_space=True) - if not features.billing.enabled or features.billing.subscription.plan == CloudPlan.SANDBOX: + if features.billing.subscription.plan == CloudPlan.SANDBOX: return { "document_ids": [], "count": 0, @@ -2208,9 +2211,8 @@ class DocumentService: assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None - features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) - - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) if not knowledge_config.original_document_id: count = 0 if knowledge_config.data_source: @@ -2520,7 +2522,7 @@ class DocumentService: # # check document limit # features = FeatureService.get_features(current_user.current_tenant_id) - # if features.billing.enabled: + # if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: # if not knowledge_config.original_document_id: # count = 0 # if knowledge_config.data_source: @@ -2797,7 +2799,7 @@ class DocumentService: @staticmethod def check_document_creation_limits(count: int, features: FeatureModel): """Validate billing-backed document creation limits before document rows are created.""" - if not features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION != DeploymentEdition.CLOUD: return if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: @@ -3014,9 +3016,8 @@ class DocumentService: assert current_user.current_tenant_id is not None assert knowledge_config.data_source - features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) - - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True) count = 0 if knowledge_config.data_source.info_list.data_source_type == "upload_file": upload_file_list = ( diff --git a/api/services/document_indexing_proxy/base.py b/api/services/document_indexing_proxy/base.py index 91835c8abb2..b38d6e64a6a 100644 --- a/api/services/document_indexing_proxy/base.py +++ b/api/services/document_indexing_proxy/base.py @@ -4,7 +4,8 @@ from collections.abc import Callable from functools import cached_property from typing import Any, ClassVar -from enums import CloudPlan +from configs import dify_config +from enums import CloudPlan, DeploymentEdition from services.feature_service import FeatureService logger = logging.getLogger(__name__) @@ -88,14 +89,10 @@ class DocumentTaskProxyBase(ABC): - Paid plans → priority queue + tenant isolation - Self-hosted → priority queue, no isolation """ - logger.info( - "dispatch args: %s - %s - %s", - self._tenant_id, - self.features.billing.enabled, - self.features.billing.subscription.plan, - ) - # dispatch to different indexing queue with tenant isolation when billing enabled - if self.features.billing.enabled: + logger.info("Dispatching tenant %s in %s", self._tenant_id, dify_config.DEPLOYMENT_EDITION) + + # Cloud queues isolate tenants and prioritize paid plans. + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: if self.features.billing.subscription.plan == CloudPlan.SANDBOX: # dispatch to normal pipeline queue with tenant self sub queue for sandbox plan self._send_to_default_tenant_queue() diff --git a/api/services/entities/feature_entities.py b/api/services/entities/feature_entities.py index f35c558d745..a963240e2ad 100644 --- a/api/services/entities/feature_entities.py +++ b/api/services/entities/feature_entities.py @@ -17,13 +17,6 @@ class SubscriptionModel(FeatureResponseModel): class BillingModel(FeatureResponseModel): - # Deprecated compatibility field. Deployment edition is the only source of truth for product edition. - # TODO: Remove after clients migrate to `SystemFeatureModel.deployment_edition`. - enabled: bool = Field( - default=False, - deprecated=True, - description="Deprecated. Use system features deployment_edition to determine the product edition.", - ) subscription: SubscriptionModel = SubscriptionModel() diff --git a/api/services/feature_service.py b/api/services/feature_service.py index b7627e9df4a..dd9f84f8cd9 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -12,8 +12,6 @@ class FeatureService: return CloudPlan.SANDBOX billing_info = BillingService.get_info(tenant_id, exclude_vector_space=True) - if not billing_info["enabled"]: - return CloudPlan.SANDBOX return CloudPlan(billing_info["subscription"]["plan"]) @classmethod @@ -86,7 +84,7 @@ class FeatureService: return True if not tenant_id: return False - return features.billing.enabled and features.billing.subscription.plan.is_paid + return features.billing.subscription.plan.is_paid @classmethod def _fulfill_trial_models_from_env(cls, quota_types: tuple[str, ...] | None = None) -> list[str]: @@ -140,7 +138,6 @@ class FeatureService: features_usage_info = BillingService.get_quota_info(tenant_id) - features.billing.enabled = billing_info["enabled"] features.billing.subscription.plan = CloudPlan(billing_info["subscription"]["plan"]) features.billing.subscription.interval = billing_info["subscription"]["interval"] features.education.activated = billing_info["subscription"].get("education", False) diff --git a/api/services/rag_pipeline/rag_pipeline_task_proxy.py b/api/services/rag_pipeline/rag_pipeline_task_proxy.py index 4c466e37d0b..ab3b7338c08 100644 --- a/api/services/rag_pipeline/rag_pipeline_task_proxy.py +++ b/api/services/rag_pipeline/rag_pipeline_task_proxy.py @@ -3,9 +3,10 @@ import logging from collections.abc import Callable, Sequence from functools import cached_property +from configs import dify_config from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from extensions.ext_database import db from services.feature_service import FeatureService from services.file_service import FileService @@ -79,15 +80,10 @@ class RagPipelineTaskProxy: if not upload_file_id: raise ValueError("upload_file_id is empty") - logger.info( - "dispatch args: %s - %s - %s", - self._dataset_tenant_id, - self.features.billing.enabled, - self.features.billing.subscription.plan, - ) + logger.info("Dispatching tenant %s in %s", self._dataset_tenant_id, dify_config.DEPLOYMENT_EDITION) - # dispatch to different pipeline queue with tenant isolation when billing enabled - if self.features.billing.enabled: + # Cloud queues isolate tenants and prioritize paid plans. + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: if self.features.billing.subscription.plan == CloudPlan.SANDBOX: # dispatch to normal pipeline queue with tenant isolation for sandbox plan self._send_to_default_tenant_queue(upload_file_id) diff --git a/api/services/vector_space_admission_service.py b/api/services/vector_space_admission_service.py index 8d3da92bb1c..98d3a8aef5f 100644 --- a/api/services/vector_space_admission_service.py +++ b/api/services/vector_space_admission_service.py @@ -395,15 +395,13 @@ class VectorSpaceAdmissionService: ) from error plan = None - if billing_info["enabled"]: - try: - plan = CloudPlan(billing_info["subscription"]["plan"]) - except ValueError: - logger.warning( - "Skipping TiDB vector-space admission for unknown plan tenant_id=%s plan=%s", - tenant_id, - billing_info["subscription"]["plan"], - ) + try: + plan = CloudPlan(billing_info["subscription"]["plan"]) + except ValueError: + logger.warning( + "Skipping TiDB vector-space admission for unknown plan tenant_id=%s", + tenant_id, + ) self._plan_by_tenant[tenant_id] = plan return plan diff --git a/api/services/web_app_runtime_query_service.py b/api/services/web_app_runtime_query_service.py index 559b0685d85..5b64c946be5 100644 --- a/api/services/web_app_runtime_query_service.py +++ b/api/services/web_app_runtime_query_service.py @@ -4,6 +4,7 @@ import json from collections.abc import Callable, Mapping from typing import NamedTuple, Protocol, cast +from enums import DeploymentEdition from services.app_definition_query_service import AppSiteConfiguration from services.entities.feature_entities import FeatureModel from services.file_service import FileService @@ -50,11 +51,13 @@ class WebAppRuntimeQueryService: file_service: FileService, workspace_features: Callable[[str], FeatureModel], files_url: str, + deployment_edition: DeploymentEdition, ) -> None: self._runtime = runtime self._file_service = file_service self._workspace_features = workspace_features self._files_url = files_url + self._deployment_edition = deployment_edition def get_bootstrap(self, app_id: str) -> WebAppBootstrap: record = self._runtime.get_runtime_record(app_id) @@ -70,7 +73,7 @@ class WebAppRuntimeQueryService: site = cast(dict[str, str | bool | None], record.site._asdict()) site["icon_url"] = site_icon_url - if features.billing.enabled and not features.webapp_copyright_enabled: + if self._deployment_edition == DeploymentEdition.CLOUD and not features.webapp_copyright_enabled: site["copyright"] = None site["input_placeholder"] = None diff --git a/api/services/workspace_service.py b/api/services/workspace_service.py index 855dcdd1131..9cfd620867a 100644 --- a/api/services/workspace_service.py +++ b/api/services/workspace_service.py @@ -73,7 +73,7 @@ class WorkspaceService: if effective_pool is None: return EffectiveCreditPool( - plan=subscription_plan if billing_info["enabled"] else None, + plan=subscription_plan, next_credit_reset_date=billing_info.get("next_credit_reset_date"), ) @@ -87,7 +87,7 @@ class WorkspaceService: exhausted_at = None return EffectiveCreditPool( - plan=subscription_plan if billing_info["enabled"] else None, + plan=subscription_plan, pool_type=effective_pool_type, quota_limit=effective_pool.quota_limit, quota_used=effective_pool.quota_used, @@ -137,7 +137,9 @@ class WorkspaceService: tenant_info["role"] = tenant_account_join.role feature = FeatureService.get_features(tenant.id, exclude_vector_space=True) - tenant_info["plan"] = feature.billing.subscription.plan if feature.billing.enabled else None + tenant_info["plan"] = ( + feature.billing.subscription.plan if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD else None + ) can_replace_logo = feature.can_replace_logo if can_replace_logo and TenantService.has_roles( diff --git a/api/tasks/document_indexing_task.py b/api/tasks/document_indexing_task.py index 2ad17d31577..7d41c542bba 100644 --- a/api/tasks/document_indexing_task.py +++ b/api/tasks/document_indexing_task.py @@ -13,7 +13,7 @@ from core.entities.document_task import DocumentTask from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document from models.enums import IndexingStatus @@ -62,9 +62,9 @@ def _document_indexing(dataset_id: str, document_ids: Sequence[str]): logger.info(click.style(f"Dataset is not found: {dataset_id}", fg="yellow")) return # check document limit - features = FeatureService.get_features(dataset.tenant_id) - try: - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(dataset.tenant_id) + try: vector_space = features.vector_space assert vector_space is not None count = len(document_ids) @@ -78,18 +78,18 @@ def _document_indexing(dataset_id: str, document_ids: Sequence[str]): "Your total number of documents plus the number of uploads have over the limit of " "your subscription." ) - except Exception as e: - for document_id in document_ids: - document = session.scalar( - select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1) - ) - if document: - document.indexing_status = IndexingStatus.ERROR - document.error = str(e) - document.stopped_at = naive_utc_now() - session.add(document) - session.commit() - return + except Exception as e: + for document_id in document_ids: + document = session.scalar( + select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1) + ) + if document: + document.indexing_status = IndexingStatus.ERROR + document.error = str(e) + document.stopped_at = naive_utc_now() + session.add(document) + session.commit() + return # Phase 1: Persist parsing status before slow extraction and vector operations. with session_factory.create_session() as session, session.begin(): diff --git a/api/tasks/duplicate_document_indexing_task.py b/api/tasks/duplicate_document_indexing_task.py index d062794df33..c083dbdbbad 100644 --- a/api/tasks/duplicate_document_indexing_task.py +++ b/api/tasks/duplicate_document_indexing_task.py @@ -12,7 +12,7 @@ from core.entities.document_task import DocumentTask from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document, DocumentSegment from models.enums import IndexingStatus @@ -88,9 +88,9 @@ def _duplicate_document_indexing_task(dataset_id: str, document_ids: Sequence[st return # check document limit - features = FeatureService.get_features(dataset.tenant_id) - try: - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(dataset.tenant_id) + try: vector_space = features.vector_space assert vector_space is not None count = len(document_ids) @@ -106,20 +106,20 @@ def _duplicate_document_indexing_task(dataset_id: str, document_ids: Sequence[st "Your total number of documents plus the number of uploads have exceeded the limit of " "your subscription." ) - except Exception as e: - documents = list( - session.scalars( - select(Document).where(Document.id.in_(document_ids), Document.dataset_id == dataset_id) - ).all() - ) - for document in documents: - if document is not None: - document.indexing_status = IndexingStatus.ERROR - document.error = str(e) - document.stopped_at = naive_utc_now() - session.add(document) - session.commit() - return + except Exception as e: + documents = list( + session.scalars( + select(Document).where(Document.id.in_(document_ids), Document.dataset_id == dataset_id) + ).all() + ) + for document in documents: + if document is not None: + document.indexing_status = IndexingStatus.ERROR + document.error = str(e) + document.stopped_at = naive_utc_now() + session.add(document) + session.commit() + return documents = list( session.scalars( diff --git a/api/tasks/retry_document_indexing_task.py b/api/tasks/retry_document_indexing_task.py index dcf74737544..08e0aaea5ab 100644 --- a/api/tasks/retry_document_indexing_task.py +++ b/api/tasks/retry_document_indexing_task.py @@ -5,9 +5,11 @@ import click from celery import shared_task from sqlalchemy import delete, select +from configs import dify_config from core.db.session_factory import session_factory from core.indexing_runner import IndexingRunner from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from enums import DeploymentEdition from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models import Account, Tenant @@ -48,9 +50,9 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str], user_ for document_id in document_ids: retry_indexing_cache_key = f"document_{document_id}_is_retried" # check document limit - features = FeatureService.get_features(tenant.id) - try: - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(tenant.id) + try: vector_space = features.vector_space assert vector_space is not None if 0 < vector_space.limit <= vector_space.size: @@ -58,18 +60,20 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str], user_ "Your total number of documents plus the number of uploads have over the limit of " "your subscription." ) - except Exception as e: - document = session.scalar( - select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1) - ) - if document: - document.indexing_status = IndexingStatus.ERROR - document.error = str(e) - document.stopped_at = naive_utc_now() - session.add(document) - session.commit() - redis_client.delete(retry_indexing_cache_key) - return + except Exception as e: + document = session.scalar( + select(Document) + .where(Document.id == document_id, Document.dataset_id == dataset_id) + .limit(1) + ) + if document: + document.indexing_status = IndexingStatus.ERROR + document.error = str(e) + document.stopped_at = naive_utc_now() + session.add(document) + session.commit() + redis_client.delete(retry_indexing_cache_key) + return logger.info(click.style(f"Start retry document: {document_id}", fg="green")) document = session.scalar( diff --git a/api/tasks/sync_website_document_indexing_task.py b/api/tasks/sync_website_document_indexing_task.py index 6cd6d40459a..cceafa9623c 100644 --- a/api/tasks/sync_website_document_indexing_task.py +++ b/api/tasks/sync_website_document_indexing_task.py @@ -5,9 +5,11 @@ import click from celery import shared_task from sqlalchemy import delete, select +from configs import dify_config from core.db.session_factory import session_factory from core.indexing_runner import IndexingRunner from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from enums import DeploymentEdition from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, DocumentSegment @@ -45,9 +47,9 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str): sync_indexing_cache_key = f"document_{document_id}_is_sync" # check document limit - features = FeatureService.get_features(dataset.tenant_id) - try: - if features.billing.enabled: + if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD: + features = FeatureService.get_features(dataset.tenant_id) + try: vector_space = features.vector_space assert vector_space is not None if 0 < vector_space.limit <= vector_space.size: @@ -55,14 +57,14 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str): "Your total number of documents plus the number of uploads have over the limit of " "your subscription." ) - except Exception as e: - document.indexing_status = IndexingStatus.ERROR - document.error = str(e) - document.stopped_at = naive_utc_now() - session.add(document) - session.commit() - redis_client.delete(sync_indexing_cache_key) - return + except Exception as e: + document.indexing_status = IndexingStatus.ERROR + document.error = str(e) + document.stopped_at = naive_utc_now() + session.add(document) + session.commit() + redis_client.delete(sync_indexing_cache_key) + return logger.info(click.style(f"Start sync website document: {document_id}", fg="green")) try: diff --git a/api/tests/test_containers_integration_tests/services/test_annotation_service.py b/api/tests/test_containers_integration_tests/services/test_annotation_service.py index f9e66d91e78..d5ba5b20d68 100644 --- a/api/tests/test_containers_integration_tests/services/test_annotation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_annotation_service.py @@ -5,6 +5,7 @@ from faker import Faker from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound +from enums import DeploymentEdition from models import Account from models.enums import ConversationFromSource, InvokeFrom from models.model import MessageAnnotation @@ -12,6 +13,7 @@ from services.annotation_service import AppAnnotationService from services.app_ref_service import AnnotationRef, AppRef from services.app_service import AppService, CreateAppParams from tests.test_containers_integration_tests.helpers import generate_valid_password +from tests.unit_tests.config_override import config_overrides_context class TestAnnotationService: @@ -32,7 +34,6 @@ class TestAnnotationService: patch("services.annotation_service.current_account_with_tenant") as mock_current_account_with_tenant, ): # Setup default mock returns - mock_account_feature_service.get_features.return_value.billing.enabled = False mock_add_task.delay.return_value = None mock_update_task.delay.return_value = None mock_delete_task.delay.return_value = None @@ -879,6 +880,7 @@ class TestAnnotationService: assert retrieved_annotation.content == annotation_args["answer"] assert retrieved_annotation.account_id == account.id + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_batch_import_app_annotations_success( self, db_session_with_containers: Session, mock_external_service_dependencies ): @@ -900,8 +902,6 @@ class TestAnnotationService: stream=BytesIO(csv_content.encode("utf-8")), filename="annotations.csv", content_type="text/csv" ) - mock_external_service_dependencies["feature_service"].get_features.return_value.billing.enabled = False - # Mock pandas to return expected DataFrame import pandas as pd @@ -958,6 +958,7 @@ class TestAnnotationService: assert "error_msg" in result assert "empty" in result["error_msg"].lower() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_batch_import_app_annotations_quota_exceeded( self, db_session_with_containers: Session, mock_external_service_dependencies ): @@ -989,7 +990,6 @@ class TestAnnotationService: mock_pd.read_csv.return_value = mock_df # Mock FeatureService to return billing enabled with quota exceeded - mock_external_service_dependencies["feature_service"].get_features.return_value.billing.enabled = True mock_external_service_dependencies[ "feature_service" ].get_features.return_value.annotation_quota_limit.limit = 1 diff --git a/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py b/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py index ba0a3c2707a..1312911df1a 100644 --- a/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py +++ b/api/tests/test_containers_integration_tests/services/test_api_based_extension_service.py @@ -20,9 +20,6 @@ class TestAPIBasedExtensionService: patch("services.account_service.SystemFeatureService") as mock_account_feature_service, patch("services.api_based_extension_service.APIBasedExtensionRequestor") as mock_requestor, ): - # Setup default mock returns - mock_account_feature_service.get_features.return_value.billing.enabled = False - # Mock successful ping response mock_requestor_instance = mock_requestor.return_value mock_requestor_instance.request.return_value = {"result": "pong"} diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py index 58a129a15e4..e0dd574a175 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py @@ -10,6 +10,7 @@ from flask import Flask from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexTechniqueType +from enums import DeploymentEdition from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import ( AppDatasetJoin, @@ -23,6 +24,7 @@ from models.enums import DataSourceType from services.dataset_ref_service import DatasetRef, DatasetRefService from services.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService from services.errors.account import NoPermissionError +from tests.unit_tests.config_override import config_overrides_context class DatasetPermissionIntegrationFactory: @@ -406,13 +408,10 @@ class TestDatasetServicePermissionsAndLifecycle: assert dataset.updated_by == owner.id assert dataset.updated_at == now - def test_get_dataset_auto_disable_logs_returns_empty_when_billing_is_disabled( - self, db_session_with_containers: Session - ): + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_get_dataset_auto_disable_logs_returns_empty_outside_cloud(self, db_session_with_containers: Session): owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers) - features = SimpleNamespace( - billing=SimpleNamespace(enabled=False, subscription=SimpleNamespace(plan="professional")) - ) + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="professional"))) dataset_ref = DatasetRef(tenant_id=tenant.id, dataset_id=str(uuid4())) with patch("services.dataset_service.FeatureService.get_features", return_value=features): @@ -420,6 +419,7 @@ class TestDatasetServicePermissionsAndLifecycle: assert result == {"document_ids": [], "count": 0} + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_get_dataset_auto_disable_logs_returns_recent_document_ids(self, db_session_with_containers: Session): owner, tenant = DatasetPermissionIntegrationFactory.create_account_with_tenant(db_session_with_containers) dataset = DatasetPermissionIntegrationFactory.create_dataset( @@ -439,9 +439,7 @@ class TestDatasetServicePermissionsAndLifecycle: dataset_id=dataset.id, document_id=str(uuid4()), ) - features = SimpleNamespace( - billing=SimpleNamespace(enabled=True, subscription=SimpleNamespace(plan="professional")) - ) + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="professional"))) dataset_ref = DatasetRefService.create_dataset_ref(dataset) with patch("services.dataset_service.FeatureService.get_features", return_value=features): diff --git a/api/tests/test_containers_integration_tests/services/test_feature_service.py b/api/tests/test_containers_integration_tests/services/test_feature_service.py index 99c8dd9b322..edbca68e194 100644 --- a/api/tests/test_containers_integration_tests/services/test_feature_service.py +++ b/api/tests/test_containers_integration_tests/services/test_feature_service.py @@ -30,7 +30,6 @@ class TestFeatureService: ): # Setup default mock returns for BillingService mock_billing_service.get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly", "education": True}, "members": {"size": 5, "limit": 10}, "apps": {"size": 3, "limit": 20}, @@ -118,7 +117,6 @@ class TestFeatureService: assert isinstance(result, FeatureModel) # Verify billing features - assert result.billing.enabled is True assert result.billing.subscription.plan == CloudPlan.PROFESSIONAL assert result.billing.subscription.interval == "monthly" assert result.education.activated is True @@ -184,7 +182,6 @@ class TestFeatureService: # Set mock return value inside the patch context mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.SANDBOX, "interval": "monthly", "education": False}, "members": {"size": 1, "limit": 3}, "apps": {"size": 1, "limit": 5}, @@ -510,9 +507,6 @@ class TestFeatureService: assert result is not None assert isinstance(result, FeatureModel) - # Verify billing is disabled - assert result.billing.enabled is False - # Verify environment-based features assert result.can_replace_logo is True assert result.model_load_balancing_enabled is True @@ -598,9 +592,6 @@ class TestFeatureService: assert result is not None assert isinstance(result, FeatureModel) - # Cloud billing is not loaded in the Enterprise edition. - assert result.billing.enabled is False - # Verify enterprise features assert result.webapp_copyright_enabled is True @@ -710,9 +701,6 @@ class TestFeatureService: assert result is not None assert isinstance(result, FeatureModel) - # Billing data is not loaded without a tenant ID. - assert result.billing.enabled is False - # Verify environment-based features assert result.can_replace_logo is True assert result.model_load_balancing_enabled is False @@ -753,7 +741,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "yearly"}, # Missing members, apps, vector_space, etc. } @@ -766,7 +753,6 @@ class TestFeatureService: assert isinstance(result, FeatureModel) # Verify billing features - assert result.billing.enabled is True assert result.billing.subscription.plan == CloudPlan.PROFESSIONAL assert result.billing.subscription.interval == "yearly" @@ -814,7 +800,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly"}, "vector_space": {"size": 0, "limit": 0}, "apps": {"size": 5, "limit": 10}, @@ -931,7 +916,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "yearly"}, "members": {"size": 10, "limit": 10}, "vector_space": {"size": 3, "limit": 5}, @@ -1251,7 +1235,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.TEAM, "interval": "yearly"}, "members": {"size": 0, "limit": 0}, "apps": {"size": 0, "limit": -1}, @@ -1355,7 +1338,6 @@ class TestFeatureService: # Arrange: Setup edge case education mock tenant_id = self._create_test_tenant_id() mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "semester", "education": True}, "members": {"size": 100, "limit": 200}, "apps": {"size": 50, "limit": 100}, @@ -1509,7 +1491,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.TEAM, "interval": "monthly"}, "docs_processing": "advanced", "can_replace_logo": True, @@ -1627,7 +1608,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.TEAM, "interval": "yearly"}, "annotation_quota_limit": {"size": 999, "limit": 1000}, "knowledge_rate_limit": {"limit": 500}, @@ -1688,7 +1668,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly"}, "documents_upload_quota": { "size": 0, # Edge case: zero current size @@ -1803,7 +1782,6 @@ class TestFeatureService: mock_config.EDUCATION_ENABLED = False mock_external_service_dependencies["billing_service"].get_info.return_value = { - "enabled": True, "subscription": { "plan": CloudPlan.PROFESSIONAL, "interval": "monthly", diff --git a/api/tests/test_containers_integration_tests/services/test_message_service.py b/api/tests/test_containers_integration_tests/services/test_message_service.py index 62d03690ee9..8c935876fe2 100644 --- a/api/tests/test_containers_integration_tests/services/test_message_service.py +++ b/api/tests/test_containers_integration_tests/services/test_message_service.py @@ -33,7 +33,6 @@ class TestMessageService: patch("services.message_service.TokenBufferMemory") as mock_token_buffer_memory, ): # Setup default mock returns - mock_account_feature_service.get_features.return_value.billing.enabled = False # Mock ModelManager mock_model_instance = mock_model_manager.return_value.get_default_model_instance.return_value diff --git a/api/tests/test_containers_integration_tests/services/test_workspace_service.py b/api/tests/test_containers_integration_tests/services/test_workspace_service.py index d6edcd25f85..e8a14c68acf 100644 --- a/api/tests/test_containers_integration_tests/services/test_workspace_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workspace_service.py @@ -27,9 +27,9 @@ class TestWorkspaceService: # Setup default mock returns feature = mock_feature_service.get_features.return_value feature.can_replace_logo = True - feature.billing.enabled = True feature.billing.subscription.plan = "professional" mock_tenant_service.has_roles.return_value = True + mock_dify_config.DEPLOYMENT_EDITION = DeploymentEdition.CLOUD mock_dify_config.FILES_URL = "https://example.com/files" yield { @@ -611,7 +611,6 @@ class TestWorkspaceService: mock_external_service_dependencies["dify_config"].DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY feature = mock_external_service_dependencies["feature_service"].get_features.return_value feature.can_replace_logo = False - feature.billing.enabled = False mock_external_service_dependencies["tenant_service"].has_roles.return_value = False with patch("services.workspace_service.current_user", account): diff --git a/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py index 0c36794cb8c..b0700ea9709 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_dataset_indexing_task.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import Session from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexTechniqueType -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus @@ -22,6 +22,7 @@ from tasks.document_indexing_task import ( normal_document_indexing_task, priority_document_indexing_task, ) +from tests.unit_tests.config_override import config_overrides_context class _TrackedSessionContext: @@ -89,7 +90,6 @@ def patched_external_dependencies(): ): mock_runner_instance = mock_indexing_runner.return_value mock_features = MagicMock() - mock_features.billing.enabled = False mock_features.billing.subscription.plan = CloudPlan.PROFESSIONAL mock_features.vector_space.limit = 100 mock_features.vector_space.size = 0 @@ -249,6 +249,7 @@ class TestDatasetIndexingTaskIntegration: assert len(run_args) == len(document_ids) self._assert_documents_parsing(db_session_with_containers, document_ids) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_batch_processing_with_limit_check( self, db_session_with_containers: Session, patched_external_dependencies ): @@ -260,7 +261,6 @@ class TestDatasetIndexingTaskIntegration: dataset, documents = self._create_test_dataset_and_documents(db_session_with_containers, document_count=3) document_ids = [doc.id for doc in documents] features = patched_external_dependencies["features"] - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.PROFESSIONAL features.vector_space.limit = 100 features.vector_space.size = 50 @@ -273,6 +273,7 @@ class TestDatasetIndexingTaskIntegration: patched_external_dependencies["indexing_runner_instance"].run.assert_not_called() self._assert_documents_error_contains(db_session_with_containers, document_ids, "batch upload limit") + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_batch_processing_sandbox_plan_single_document_only( self, db_session_with_containers: Session, patched_external_dependencies ): @@ -281,7 +282,6 @@ class TestDatasetIndexingTaskIntegration: dataset, documents = self._create_test_dataset_and_documents(db_session_with_containers, document_count=2) document_ids = [doc.id for doc in documents] features = patched_external_dependencies["features"] - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.SANDBOX # Act @@ -375,6 +375,7 @@ class TestDatasetIndexingTaskIntegration: task_dispatch_spy.apply_async.assert_not_called() delete_key_spy.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_validation_failure_sets_error_status_when_vector_space_at_limit( self, db_session_with_containers: Session, patched_external_dependencies ): @@ -383,7 +384,6 @@ class TestDatasetIndexingTaskIntegration: dataset, documents = self._create_test_dataset_and_documents(db_session_with_containers, document_count=3) document_ids = [doc.id for doc in documents] features = patched_external_dependencies["features"] - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.PROFESSIONAL features.vector_space.limit = 100 features.vector_space.size = 100 @@ -590,6 +590,7 @@ class TestDatasetIndexingTaskIntegration: call_kwargs = task_dispatch_spy.apply_async.call_args_list[index].kwargs.get("kwargs", {}) assert call_kwargs.get("document_ids") == expected_task["document_ids"] + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_billing_disabled_skips_limit_checks( self, db_session_with_containers: Session, patched_external_dependencies ): @@ -601,7 +602,6 @@ class TestDatasetIndexingTaskIntegration: document_ids=large_document_ids, ) features = patched_external_dependencies["features"] - features.billing.enabled = False # Act _document_indexing(dataset.id, large_document_ids) @@ -688,6 +688,7 @@ class TestDatasetIndexingTaskIntegration: # Assert self._assert_documents_parsing(db_session_with_containers, [special_document_id]) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_zero_vector_space_limit_allows_unlimited( self, db_session_with_containers: Session, patched_external_dependencies ): @@ -696,7 +697,6 @@ class TestDatasetIndexingTaskIntegration: dataset, documents = self._create_test_dataset_and_documents(db_session_with_containers, document_count=3) document_ids = [doc.id for doc in documents] features = patched_external_dependencies["features"] - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.PROFESSIONAL features.vector_space.limit = 0 features.vector_space.size = 1000 @@ -708,6 +708,7 @@ class TestDatasetIndexingTaskIntegration: patched_external_dependencies["indexing_runner_instance"].run.assert_called_once() self._assert_documents_parsing(db_session_with_containers, document_ids) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_negative_vector_space_values_handled_gracefully( self, db_session_with_containers: Session, patched_external_dependencies ): @@ -716,7 +717,6 @@ class TestDatasetIndexingTaskIntegration: dataset, documents = self._create_test_dataset_and_documents(db_session_with_containers, document_count=3) document_ids = [doc.id for doc in documents] features = patched_external_dependencies["features"] - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.PROFESSIONAL features.vector_space.limit = -1 features.vector_space.size = 100 @@ -728,6 +728,7 @@ class TestDatasetIndexingTaskIntegration: patched_external_dependencies["indexing_runner_instance"].run.assert_called_once() self._assert_documents_parsing(db_session_with_containers, document_ids) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_large_document_batch_processing(self, db_session_with_containers: Session, patched_external_dependencies): """Process a batch exactly at configured upload limit. @@ -741,7 +742,6 @@ class TestDatasetIndexingTaskIntegration: document_ids=document_ids, ) features = patched_external_dependencies["features"] - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.PROFESSIONAL features.vector_space.limit = 10000 features.vector_space.size = 0 diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py index bb45e2be34a..85b3993b1fd 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py @@ -7,7 +7,7 @@ from sqlalchemy.orm import Session from core.entities.document_task import DocumentTask from core.rag.index_processor.constant.index_type import IndexTechniqueType -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus @@ -18,6 +18,7 @@ from tasks.document_indexing_task import ( normal_document_indexing_task, # New normal task priority_document_indexing_task, # New priority task ) +from tests.unit_tests.config_override import config_overrides_context class TestDocumentIndexingTasks: @@ -41,7 +42,6 @@ class TestDocumentIndexingTasks: # Setup mock indexing runner mock_runner_instance = mock_indexing_runner.return_value # Setup mock feature service mock_features = MagicMock() - mock_features.billing.enabled = False mock_feature_service.get_features.return_value = mock_features yield { @@ -147,7 +147,9 @@ class TestDocumentIndexingTasks: return dataset, documents def _create_test_dataset_with_billing_features( - self, db_session_with_containers: Session, mock_external_service_dependencies, billing_enabled=True + self, + db_session_with_containers: Session, + mock_external_service_dependencies, ): """ Helper method to create a test dataset with billing features configured. @@ -155,7 +157,6 @@ class TestDocumentIndexingTasks: Args: db_session_with_containers: Database session from testcontainers infrastructure mock_external_service_dependencies: Mock dependencies - billing_enabled: Whether billing is enabled Returns: tuple: (dataset, documents) - Created dataset and document instances @@ -224,11 +225,9 @@ class TestDocumentIndexingTasks: db_session_with_containers.commit() # Configure billing features - mock_external_service_dependencies["features"].billing.enabled = billing_enabled - if billing_enabled: - mock_external_service_dependencies["features"].billing.subscription.plan = CloudPlan.SANDBOX - mock_external_service_dependencies["features"].vector_space.limit = 100 - mock_external_service_dependencies["features"].vector_space.size = 50 + mock_external_service_dependencies["features"].billing.subscription.plan = CloudPlan.SANDBOX + mock_external_service_dependencies["features"].vector_space.limit = 100 + mock_external_service_dependencies["features"].vector_space.size = 50 # Refresh dataset to ensure it's properly loaded db_session_with_containers.refresh(dataset) @@ -467,6 +466,7 @@ class TestDocumentIndexingTasks: processed_documents = self._runner_documents_arg(mock_external_service_dependencies) assert len(processed_documents) == 4 + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_document_indexing_task_billing_sandbox_plan_batch_limit( self, db_session_with_containers: Session, mock_external_service_dependencies ): @@ -481,7 +481,8 @@ class TestDocumentIndexingTasks: """ # Arrange: Create test data with billing enabled dataset, documents = self._create_test_dataset_with_billing_features( - db_session_with_containers, mock_external_service_dependencies, billing_enabled=True + db_session_with_containers, + mock_external_service_dependencies, ) # Configure sandbox plan with batch limit @@ -529,21 +530,23 @@ class TestDocumentIndexingTasks: # Verify no indexing runner was called mock_external_service_dependencies["indexing_runner"].assert_not_called() - def test_document_indexing_task_billing_disabled_success( + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_document_indexing_task_community_success( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ - Test successful processing when billing is disabled. + Test successful processing outside Cloud. This test verifies: - - Processing continues normally when billing is disabled + - Processing continues normally outside Cloud - No billing validation occurs - Documents are processed successfully - IndexingRunner is called correctly """ # Arrange: Create test data with billing disabled dataset, documents = self._create_test_dataset_with_billing_features( - db_session_with_containers, mock_external_service_dependencies, billing_enabled=False + db_session_with_containers, + mock_external_service_dependencies, ) document_ids = [doc.id for doc in documents] diff --git a/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py index d5393435bfa..1dacc2bba10 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_duplicate_document_indexing_task.py @@ -7,7 +7,7 @@ from sqlalchemy.orm import Session from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus @@ -18,6 +18,7 @@ from tasks.duplicate_document_indexing_task import ( normal_duplicate_document_indexing_task, # New normal task priority_duplicate_document_indexing_task, # New priority task ) +from tests.unit_tests.config_override import config_overrides_context class TestDuplicateDocumentIndexingTasks: @@ -45,7 +46,6 @@ class TestDuplicateDocumentIndexingTasks: # Setup mock indexing runner mock_runner_instance = mock_indexing_runner.return_value # Setup mock feature service mock_features = MagicMock() - mock_features.billing.enabled = False mock_feature_service.get_features.return_value = mock_features # Setup mock index processor factory @@ -214,7 +214,9 @@ class TestDuplicateDocumentIndexingTasks: return dataset, documents, segments def _create_test_dataset_with_billing_features( - self, db_session_with_containers: Session, mock_external_service_dependencies, billing_enabled=True + self, + db_session_with_containers: Session, + mock_external_service_dependencies, ): """ Helper method to create a test dataset with billing features configured. @@ -222,7 +224,6 @@ class TestDuplicateDocumentIndexingTasks: Args: db_session_with_containers: Database session from testcontainers infrastructure mock_external_service_dependencies: Mock dependencies - billing_enabled: Whether billing is enabled Returns: tuple: (dataset, documents) - Created dataset and document instances @@ -292,11 +293,9 @@ class TestDuplicateDocumentIndexingTasks: db_session_with_containers.commit() # Configure billing features - mock_external_service_dependencies["features"].billing.enabled = billing_enabled - if billing_enabled: - mock_external_service_dependencies["features"].billing.subscription.plan = CloudPlan.SANDBOX - mock_external_service_dependencies["features"].vector_space.limit = 100 - mock_external_service_dependencies["features"].vector_space.size = 50 + mock_external_service_dependencies["features"].billing.subscription.plan = CloudPlan.SANDBOX + mock_external_service_dependencies["features"].vector_space.limit = 100 + mock_external_service_dependencies["features"].vector_space.size = 50 # Refresh dataset to ensure it's properly loaded db_session_with_containers.refresh(dataset) @@ -517,7 +516,8 @@ class TestDuplicateDocumentIndexingTasks: """ # Arrange: Create test data with billing enabled dataset, documents = self._create_test_dataset_with_billing_features( - db_session_with_containers, mock_external_service_dependencies, billing_enabled=True + db_session_with_containers, + mock_external_service_dependencies, ) # Configure sandbox plan with batch limit @@ -580,7 +580,8 @@ class TestDuplicateDocumentIndexingTasks: """ # Arrange: Create test data with billing enabled dataset, documents = self._create_test_dataset_with_billing_features( - db_session_with_containers, mock_external_service_dependencies, billing_enabled=True + db_session_with_containers, + mock_external_service_dependencies, ) # Configure TEAM plan with vector space limit exceeded @@ -818,7 +819,8 @@ class TestDuplicateDocumentIndexingTasks: db_session_with_containers, mock_external_service_dependencies ) - def test_duplicate_document_indexing_with_billing_enabled_sandbox_plan( + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_duplicate_document_indexing_with_cloud_sandbox_plan( self, db_session_with_containers: Session, mock_external_service_dependencies ): """Test duplicate document indexing with billing enabled and sandbox plan.""" @@ -826,6 +828,7 @@ class TestDuplicateDocumentIndexingTasks: db_session_with_containers, mock_external_service_dependencies ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_duplicate_document_indexing_with_billing_limit_exceeded( self, db_session_with_containers: Session, mock_external_service_dependencies ): diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation_security.py b/api/tests/unit_tests/controllers/console/app/test_annotation_security.py index 273cfbb8417..0aa896337e3 100644 --- a/api/tests/unit_tests/controllers/console/app/test_annotation_security.py +++ b/api/tests/unit_tests/controllers/console/app/test_annotation_security.py @@ -15,10 +15,12 @@ from werkzeug.datastructures import FileStorage from configs import dify_config from controllers.console.wraps import annotation_import_concurrency_limit, annotation_import_rate_limit +from enums import DeploymentEdition from models.account import Account from models.model import App, AppMode, IconType from services.annotation_service import AppAnnotationService from tasks.annotation.batch_import_annotations_task import batch_import_annotations_task +from tests.unit_tests.config_override import config_overrides_context def _account() -> Account: @@ -212,6 +214,7 @@ class TestAnnotationImportFileValidation: class TestAnnotationImportServiceValidation: """Test service layer validation for annotation import.""" + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_max_records_limit_enforced(self, sqlite_session: Session): """Test that files with too many records are rejected.""" @@ -228,8 +231,6 @@ class TestAnnotationImportServiceValidation: mock_auth.return_value = (_account(), "tenant_id") with patch("services.annotation_service.FeatureService") as mock_features: - mock_features.get_features.return_value.billing.enabled = False - result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) # Should return error about too many records @@ -273,6 +274,7 @@ class TestAnnotationImportServiceValidation: assert "error_msg" in result assert "malformed" in result["error_msg"].lower() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_valid_import_succeeds(self, sqlite_session: Session): """Test that valid import request succeeds.""" @@ -286,8 +288,6 @@ class TestAnnotationImportServiceValidation: mock_auth.return_value = (_account(), "tenant_id") with patch("services.annotation_service.FeatureService") as mock_features: - mock_features.get_features.return_value.billing.enabled = False - with patch("services.annotation_service.batch_import_annotations_task") as mock_task: with patch("services.annotation_service.redis_client"): result = AppAnnotationService.batch_import_app_annotations("app_id", file, sqlite_session) diff --git a/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py b/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py index 0a33e67163b..6348a2f8ce9 100644 --- a/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py +++ b/api/tests/unit_tests/controllers/console/test_workflow_run_archive.py @@ -128,7 +128,6 @@ def test_workflow_run_archive_endpoint_allows_admitted_role_when_rbac_is_enabled assert tenant_id == "tenant-1" assert exclude_vector_space return { - "enabled": True, "subscription": {"plan": CloudPlan.TEAM}, } diff --git a/api/tests/unit_tests/controllers/console/test_workspace_members.py b/api/tests/unit_tests/controllers/console/test_workspace_members.py index 6557fba1107..2528b387623 100644 --- a/api/tests/unit_tests/controllers/console/test_workspace_members.py +++ b/api/tests/unit_tests/controllers/console/test_workspace_members.py @@ -29,7 +29,6 @@ def _build_feature_flags(): placeholder_quota = SimpleNamespace(limit=0, size=0) workspace_members = SimpleNamespace(enabled=False, is_available=lambda count: True) return SimpleNamespace( - billing=SimpleNamespace(enabled=False), workspace_members=workspace_members, members=placeholder_quota, apps=placeholder_quota, diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index ee8926bd4d3..54002f37246 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/test_wraps.py @@ -784,7 +784,7 @@ class TestBillingPaidPlanRequired: def paid_view(): return "paid_success" - billing_info = {"enabled": True, "subscription": {"plan": plan}} + billing_info = {"subscription": {"plan": plan}} with ( patch( "controllers.console.wraps.current_account_with_tenant", @@ -797,18 +797,15 @@ class TestBillingPaidPlanRequired: assert result == "paid_success" get_info.assert_called_once_with("tenant123", exclude_vector_space=True) - @pytest.mark.parametrize( - ("enabled", "plan"), - [(False, "professional"), (True, "sandbox"), (True, "unknown")], - ) - def test_should_reject_non_paid_plan(self, enabled: bool, plan: str): + @pytest.mark.parametrize("plan", ["sandbox", "unknown"]) + def test_should_reject_non_paid_plan(self, plan: str): app = create_app_with_login() @cloud_edition_billing_paid_plan_required def paid_view(): return "paid_success" - billing_info = {"enabled": enabled, "subscription": {"plan": plan}} + billing_info = {"subscription": {"plan": plan}} with app.test_request_context(): with ( patch( @@ -827,11 +824,11 @@ class TestBillingPaidPlanRequired: class TestBillingResourceLimits: """Test billing resource limit decorators""" + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_should_allow_when_under_resource_limit(self): """Test that requests are allowed when under resource limits""" # Arrange mock_features = MagicMock() - mock_features.billing.enabled = True mock_features.members.limit = 10 mock_features.members.size = 5 @@ -881,12 +878,12 @@ class TestBillingResourceLimits: get_vector_space.assert_called_once_with("tenant123") get_features.assert_not_called() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_should_reject_when_over_resource_limit(self): """Test that requests are rejected when over resource limits""" # Arrange app = create_app_with_login() mock_features = MagicMock() - mock_features.billing.enabled = True mock_features.members.limit = 10 mock_features.members.size = 10 @@ -906,12 +903,12 @@ class TestBillingResourceLimits: assert exc_info.value.code == 403 assert "members has reached the limit" in str(exc_info.value.description) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_should_check_source_for_documents_limit(self): """Test document limit checks request source""" # Arrange app = create_app_with_login() mock_features = MagicMock() - mock_features.billing.enabled = True mock_features.documents_upload_quota.limit = 100 mock_features.documents_upload_quota.size = 100 diff --git a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py b/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py index a5cd7ca5b9c..33a726e9c41 100644 --- a/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py +++ b/api/tests/unit_tests/controllers/openapi/test_workspaces_members.py @@ -42,6 +42,7 @@ from controllers.openapi.workspaces import ( WorkspaceMembersApi, WorkspaceSwitchApi, ) +from enums import DeploymentEdition from libs.oauth_bearer import AuthContext, Scope, SubjectType, TokenType, reset_auth_ctx, set_auth_ctx from models import Account, Tenant, TenantAccountJoin from models.account import AccountStatus, TenantAccountRole, TenantStatus @@ -55,6 +56,7 @@ from services.errors.account import ( NoPermissionError, RoleAlreadyAssignedError, ) +from tests.unit_tests.config_override import config_overrides_context if not hasattr(builtins, "MethodView"): builtins.MethodView = MethodView # type: ignore[attr-defined] @@ -444,7 +446,6 @@ def test_invite_happy_path_returns_invite_url_and_member_id( def _features( *, - billing_enabled: bool = False, members_size: int = 0, members_limit: int = 0, workspace_members_enabled: bool = False, @@ -452,17 +453,16 @@ def _features( workspace_members_limit: int = 0, ) -> SimpleNamespace: """Build a feature object matching the surface `_check_member_invite_quota` - reads: `.billing.enabled`, `.members.{size,limit}`, + reads: `.members.{size,limit}`, `.workspace_members.{enabled, is_available(N)}`. - Defaults model CE (both flags off, both caps inert). + Defaults leave both quotas unrestricted. """ def _is_available(n: int) -> bool: return workspace_members_size + n <= workspace_members_limit return SimpleNamespace( - billing=SimpleNamespace(enabled=billing_enabled), members=SimpleNamespace(size=members_size, limit=members_limit), workspace_members=SimpleNamespace( enabled=workspace_members_enabled, @@ -482,6 +482,7 @@ def _invite_request(app, ws_id: str, acct_id: uuid.UUID): ) +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_invite_blocked_by_saas_members_cap( app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch, database_session: Session ): @@ -507,7 +508,7 @@ def test_invite_blocked_by_saas_members_cap( "FeatureService", SimpleNamespace( get_features=Mock( - return_value=_features(billing_enabled=True, members_size=10, members_limit=10), + return_value=_features(members_size=10, members_limit=10), ), ), ) @@ -520,13 +521,13 @@ def test_invite_blocked_by_saas_members_cap( invite_mock.assert_not_called() +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.ENTERPRISE) def test_invite_blocked_by_ee_workspace_members_license( app: Flask, bypass_pipeline, monkeypatch: pytest.MonkeyPatch, database_session: Session ): """EE License workspace_members cap → MemberLicenseExceeded (403). - Note: billing.enabled is False (EE without SaaS billing); only the - license cap fires. + Enterprise member limits come from the license. """ ws_id = str(uuid.uuid4()) acct_id = uuid.uuid4() diff --git a/api/tests/unit_tests/controllers/service_api/app/test_completion.py b/api/tests/unit_tests/controllers/service_api/app/test_completion.py index a83e8b9f734..3a479a1d351 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_completion.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_completion.py @@ -559,7 +559,7 @@ class TestChatApiController: completion_module = sys.modules["controllers.service_api.app.completion"] apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) - billing_get_info = Mock(return_value={"enabled": True, "subscription": {"plan": CloudPlan.SANDBOX}}) + billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) generate = Mock() monkeypatch.setattr(BillingService, "get_info", billing_get_info) monkeypatch.setattr(AppGenerateService, "generate", generate) @@ -583,13 +583,12 @@ class TestChatApiController: assert exc_info.value.error_code == "workflow_version_execution_not_allowed" @pytest.mark.parametrize( - ("deployment_edition", "billing_enabled", "plan", "workflow_id"), + ("deployment_edition", "plan", "workflow_id"), [ - (DeploymentEdition.COMMUNITY, True, CloudPlan.SANDBOX, str(uuid.uuid4())), - (DeploymentEdition.ENTERPRISE, True, CloudPlan.SANDBOX, str(uuid.uuid4())), - (DeploymentEdition.CLOUD, False, CloudPlan.SANDBOX, str(uuid.uuid4())), - (DeploymentEdition.CLOUD, True, CloudPlan.PROFESSIONAL, str(uuid.uuid4())), - (DeploymentEdition.CLOUD, True, CloudPlan.SANDBOX, None), + (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX, str(uuid.uuid4())), + (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX, str(uuid.uuid4())), + (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL, str(uuid.uuid4())), + (DeploymentEdition.CLOUD, CloudPlan.SANDBOX, None), ], ) def test_allows_default_or_entitled_workflow_version_execution( @@ -598,14 +597,13 @@ class TestChatApiController: monkeypatch: pytest.MonkeyPatch, orm_session: Session, deployment_edition: DeploymentEdition, - billing_enabled: bool, plan: CloudPlan, workflow_id: str | None, ) -> None: completion_module = sys.modules["controllers.service_api.app.completion"] apply_config_overrides(monkeypatch, DEPLOYMENT_EDITION=deployment_edition) - billing_get_info = Mock(return_value={"enabled": billing_enabled, "subscription": {"plan": plan}}) + billing_get_info = Mock(return_value={"subscription": {"plan": plan}}) generate = Mock(return_value={"result": "ok"}) monkeypatch.setattr(BillingService, "get_info", billing_get_info) monkeypatch.setattr(AppGenerateService, "generate", generate) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py index 1c44c2397db..f631735796b 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py @@ -608,7 +608,7 @@ class TestWorkflowRunApi: config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) workflow_module = sys.modules["controllers.service_api.app.workflow"] - billing_get_info = Mock(return_value={"enabled": True, "subscription": {"plan": CloudPlan.SANDBOX}}) + billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) generate = Mock(return_value={"result": "ok"}) monkeypatch.setattr(BillingService, "get_info", billing_get_info) monkeypatch.setattr(AppGenerateService, "generate", generate) @@ -667,7 +667,7 @@ class TestWorkflowRunByIdApi: config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) workflow_module = sys.modules["controllers.service_api.app.workflow"] - billing_get_info = Mock(return_value={"enabled": True, "subscription": {"plan": CloudPlan.SANDBOX}}) + billing_get_info = Mock(return_value={"subscription": {"plan": CloudPlan.SANDBOX}}) generate = Mock() monkeypatch.setattr(BillingService, "get_info", billing_get_info) monkeypatch.setattr(AppGenerateService, "generate", generate) @@ -700,12 +700,11 @@ class TestWorkflowRunByIdApi: } @pytest.mark.parametrize( - ("deployment_edition", "billing_enabled", "plan"), + ("deployment_edition", "plan"), [ - (DeploymentEdition.COMMUNITY, True, CloudPlan.SANDBOX), - (DeploymentEdition.ENTERPRISE, True, CloudPlan.SANDBOX), - (DeploymentEdition.CLOUD, False, CloudPlan.SANDBOX), - (DeploymentEdition.CLOUD, True, CloudPlan.PROFESSIONAL), + (DeploymentEdition.COMMUNITY, CloudPlan.SANDBOX), + (DeploymentEdition.ENTERPRISE, CloudPlan.SANDBOX), + (DeploymentEdition.CLOUD, CloudPlan.PROFESSIONAL), ], ) def test_allows_execution_outside_enabled_sandbox_plan( @@ -713,14 +712,13 @@ class TestWorkflowRunByIdApi: app: Flask, monkeypatch: pytest.MonkeyPatch, deployment_edition: DeploymentEdition, - billing_enabled: bool, plan: CloudPlan, sqlite_session: Session, config_overrides: Callable[..., None], ) -> None: config_overrides(DEPLOYMENT_EDITION=deployment_edition) - billing_get_info = Mock(return_value={"enabled": billing_enabled, "subscription": {"plan": plan}}) + billing_get_info = Mock(return_value={"subscription": {"plan": plan}}) generate = Mock(return_value={"result": "ok"}) monkeypatch.setattr(BillingService, "get_info", billing_get_info) monkeypatch.setattr(AppGenerateService, "generate", generate) diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py index ca2fe02c02d..ace817a4bf3 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py @@ -1212,7 +1212,6 @@ class TestSegmentApiPost(SQLiteEndpointTest): mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() - mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features mock_vector_space = Mock() @@ -1555,7 +1554,6 @@ class TestDatasetSegmentApiUpdate(SQLiteEndpointTest): """Configure mocks to neutralise billing/auth decorators.""" mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() - mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features mock_vector_space = Mock() mock_vector_space.limit = 10 @@ -2054,7 +2052,6 @@ class TestChildChunkApiPost(SQLiteEndpointTest): def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() - mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features mock_vector_space = Mock() mock_vector_space.limit = 10 @@ -2368,7 +2365,6 @@ class TestModelValidateDecorator(SQLiteEndpointTest): mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() - mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features mock_vector_space = Mock() diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py index c297b95987e..5fc0837cc9f 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_document.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_document.py @@ -48,6 +48,7 @@ from controllers.service_api.dataset.document import ( ) from controllers.service_api.dataset.error import ArchivedDocumentImmutableError from core.rag.index_processor.constant.index_type import IndexStructureType +from enums import DeploymentEdition from extensions.storage.storage_type import StorageType from models.account import Account from models.dataset import Dataset, Document, DocumentSegment @@ -65,6 +66,7 @@ from services.dataset_ref_service import DatasetRef from services.dataset_service import DocumentService from services.entities.knowledge_entities.knowledge_entities import ProcessRule, RetrievalModel from services.errors.file import FileTooLargeError as FileTooLargeServiceError +from tests.unit_tests.config_override import config_overrides_context def _document_data_source_info() -> dict[str, str]: @@ -658,6 +660,7 @@ class TestDocumentServiceFileOperations: class TestDocumentServiceSaveValidation: """Test validations during document saving.""" + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) @patch("services.dataset_service.DatasetService.check_doc_form") @patch("services.dataset_service.FeatureService.get_features") def test_save_document_validates_doc_form(self, mock_features, mock_check_form, sqlite_session: Session): @@ -665,7 +668,6 @@ class TestDocumentServiceSaveValidation: dataset = make_dataset(tenant_id="tenant_id") config = Mock() features = Mock() - features.billing.enabled = False mock_features.return_value = features class TestStopError(Exception): @@ -1416,7 +1418,6 @@ class TestDocumentAddByTextApi(SQLiteControllerTest): mock_validate_token.return_value = api_token mock_features = Mock() - mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features mock_vector_space = Mock() @@ -1591,7 +1592,6 @@ def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): api_token = ApiToken(tenant_id=tenant_id, type=ApiTokenType.DATASET, token="dataset-token") mock_validate_token.return_value = api_token mock_features = Mock() - mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features mock_vector_space = Mock() mock_vector_space.limit = 10 diff --git a/api/tests/unit_tests/controllers/service_api/test_wraps.py b/api/tests/unit_tests/controllers/service_api/test_wraps.py index fc872c6e44f..997017fddaa 100644 --- a/api/tests/unit_tests/controllers/service_api/test_wraps.py +++ b/api/tests/unit_tests/controllers/service_api/test_wraps.py @@ -314,6 +314,7 @@ class TestCloudEditionBillingResourceCheck: app.config["TESTING"] = True return app + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_allows_when_under_limit(self, mock_get_features, mock_validate_token, app: Flask): @@ -322,7 +323,6 @@ class TestCloudEditionBillingResourceCheck: mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() - mock_features.billing.enabled = True mock_features.members.limit = 10 mock_features.members.size = 5 mock_get_features.return_value = mock_features @@ -381,7 +381,6 @@ class TestCloudEditionBillingResourceCheck: mock_get_vector_space.return_value = Mock(size=0, limit=50, usage_unknown=True) mock_get_features.return_value = SimpleNamespace( billing=SimpleNamespace( - enabled=True, subscription=SimpleNamespace(plan=CloudPlan.SANDBOX), ) ) @@ -411,7 +410,6 @@ class TestCloudEditionBillingResourceCheck: mock_get_vector_space.return_value = Mock(size=0, limit=50, usage_unknown=True) mock_get_features.return_value = SimpleNamespace( billing=SimpleNamespace( - enabled=True, subscription=SimpleNamespace(plan=plan), ) ) @@ -429,6 +427,7 @@ class TestCloudEditionBillingResourceCheck: assert result == "document_uploaded" mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_loads_features_when_checking_non_vector_space_limit( @@ -439,7 +438,6 @@ class TestCloudEditionBillingResourceCheck: mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() - mock_features.billing.enabled = True mock_features.documents_upload_quota.limit = 10 mock_features.documents_upload_quota.size = 5 mock_get_features.return_value = mock_features @@ -456,6 +454,7 @@ class TestCloudEditionBillingResourceCheck: assert result == "document_uploaded" mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_rejects_when_at_limit(self, mock_get_features, mock_validate_token, app: Flask): @@ -464,7 +463,6 @@ class TestCloudEditionBillingResourceCheck: mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() - mock_features.billing.enabled = True mock_features.members.limit = 10 mock_features.members.size = 10 mock_get_features.return_value = mock_features @@ -479,15 +477,15 @@ class TestCloudEditionBillingResourceCheck: add_member() assert "members has reached the limit" in str(exc_info.value) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_allows_when_billing_disabled(self, mock_get_features, mock_validate_token, app: Flask): - """Test that request is allowed when billing is disabled.""" + """Test that request is allowed outside Cloud.""" # Arrange mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() - mock_features.billing.enabled = False mock_get_features.return_value = mock_features @cloud_edition_billing_resource_check("members", "app") @@ -512,6 +510,7 @@ class TestCloudEditionBillingKnowledgeLimitCheck: app.config["TESTING"] = True return app + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_rejects_add_segment_in_sandbox(self, mock_get_features, mock_validate_token, app: Flask): @@ -520,7 +519,6 @@ class TestCloudEditionBillingKnowledgeLimitCheck: mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() - mock_features.billing.enabled = True mock_features.billing.subscription.plan = CloudPlan.SANDBOX mock_get_features.return_value = mock_features @@ -534,6 +532,7 @@ class TestCloudEditionBillingKnowledgeLimitCheck: add_segment() assert "upgrade to a paid plan" in str(exc_info.value) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("controllers.service_api.wraps.validate_and_get_api_token") @patch("controllers.service_api.wraps.FeatureService.get_features") def test_allows_other_operations_in_sandbox(self, mock_get_features, mock_validate_token, app: Flask): @@ -542,7 +541,6 @@ class TestCloudEditionBillingKnowledgeLimitCheck: mock_validate_token.return_value = Mock(tenant_id="tenant123") mock_features = Mock() - mock_features.billing.enabled = True mock_features.billing.subscription.plan = CloudPlan.SANDBOX mock_get_features.return_value = mock_features diff --git a/api/tests/unit_tests/services/dataset_service_test_helpers.py b/api/tests/unit_tests/services/dataset_service_test_helpers.py index 99157c63cab..6153afa3396 100644 --- a/api/tests/unit_tests/services/dataset_service_test_helpers.py +++ b/api/tests/unit_tests/services/dataset_service_test_helpers.py @@ -273,10 +273,9 @@ def _make_lock_context() -> MagicMock: return context_manager -def _make_features(*, enabled: bool, plan: str = CloudPlan.PROFESSIONAL) -> SimpleNamespace: +def _make_features(*, plan: str = CloudPlan.PROFESSIONAL) -> SimpleNamespace: return SimpleNamespace( billing=SimpleNamespace( - enabled=enabled, subscription=SimpleNamespace(plan=plan), ), documents_upload_quota=SimpleNamespace(limit=1000, size=0), diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py index ac6d4d5fd01..b9e5963c3f9 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_task_proxy.py @@ -5,11 +5,12 @@ from unittest.mock import Mock import pytest from pytest_mock import MockerFixture -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from extensions.storage.storage_type import StorageType from models.enums import CreatorUserRole from models.model import UploadFile from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy +from tests.unit_tests.config_override import config_overrides_context @pytest.fixture @@ -53,13 +54,12 @@ def test_delay_with_entities_calls_dispatch(mocker: MockerFixture, proxy) -> Non # --- _dispatch --- +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_dispatch_billing_sandbox_uses_default_tenant_queue(mocker: MockerFixture, proxy) -> None: upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") send_mock = mocker.patch.object(proxy, "_send_to_default_tenant_queue") - features = SimpleNamespace( - billing=SimpleNamespace(enabled=True, subscription=SimpleNamespace(plan=CloudPlan.SANDBOX)) - ) + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan=CloudPlan.SANDBOX))) mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) proxy._dispatch() @@ -68,13 +68,12 @@ def test_dispatch_billing_sandbox_uses_default_tenant_queue(mocker: MockerFixtur send_mock.assert_called_once_with("file-1") +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_dispatch_billing_non_sandbox_uses_priority_tenant_queue(mocker: MockerFixture, proxy) -> None: upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") send_mock = mocker.patch.object(proxy, "_send_to_priority_tenant_queue") - features = SimpleNamespace( - billing=SimpleNamespace(enabled=True, subscription=SimpleNamespace(plan=CloudPlan.PROFESSIONAL)) - ) + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan=CloudPlan.PROFESSIONAL))) mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) proxy._dispatch() @@ -83,11 +82,12 @@ def test_dispatch_billing_non_sandbox_uses_priority_tenant_queue(mocker: MockerF send_mock.assert_called_once_with("file-1") +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_dispatch_no_billing_uses_priority_direct_queue(mocker: MockerFixture, proxy) -> None: upload_mock = mocker.patch.object(proxy, "_upload_invoke_entities", return_value="file-1") send_mock = mocker.patch.object(proxy, "_send_to_priority_direct_queue") - features = SimpleNamespace(billing=SimpleNamespace(enabled=False, subscription=SimpleNamespace(plan="free"))) + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="free"))) mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) proxy._dispatch() @@ -96,10 +96,11 @@ def test_dispatch_no_billing_uses_priority_direct_queue(mocker: MockerFixture, p send_mock.assert_called_once_with("file-1") +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_dispatch_raises_on_empty_upload_file_id(mocker: MockerFixture, proxy) -> None: mocker.patch.object(proxy, "_upload_invoke_entities", return_value="") - features = SimpleNamespace(billing=SimpleNamespace(enabled=False, subscription=SimpleNamespace(plan="free"))) + features = SimpleNamespace(billing=SimpleNamespace(subscription=SimpleNamespace(plan="free"))) mocker.patch.object(type(proxy), "features", new_callable=lambda: property(lambda self: features)) with pytest.raises(ValueError, match="upload_file_id is empty"): diff --git a/api/tests/unit_tests/services/test_annotation_service.py b/api/tests/unit_tests/services/test_annotation_service.py index 097f12e9482..8e5c305fa34 100644 --- a/api/tests/unit_tests/services/test_annotation_service.py +++ b/api/tests/unit_tests/services/test_annotation_service.py @@ -22,6 +22,7 @@ from werkzeug.datastructures import FileStorage from werkzeug.exceptions import NotFound import services.annotation_service as annotation_service_module +from enums import DeploymentEdition from models.account import Account from models.dataset import DatasetCollectionBinding from models.enums import CollectionBindingType @@ -37,6 +38,7 @@ from models.model import ( ) from services.annotation_service import AppAnnotationService from services.app_ref_service import AnnotationRef, AppRef +from tests.unit_tests.config_override import config_overrides_context TENANT_ID = "tenant-1" OTHER_TENANT_ID = "tenant-2" @@ -593,16 +595,13 @@ class TestAppAnnotationServiceBatchImport: features: Any | None = None, ) -> dict[str, Any]: if features is None: - features = SimpleNamespace(billing=SimpleNamespace(enabled=False), annotation_quota_limit=None) + features = SimpleNamespace(annotation_quota_limit=None) with ( patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), - patch( - "configs.dify_config", - new=SimpleNamespace( - ANNOTATION_IMPORT_MAX_RECORDS=maximum, - ANNOTATION_IMPORT_MIN_RECORDS=minimum, - ), + config_overrides_context( + ANNOTATION_IMPORT_MAX_RECORDS=maximum, + ANNOTATION_IMPORT_MIN_RECORDS=minimum, ), ): return AppAnnotationService.batch_import_app_annotations(app.id, _file(content), sqlite_session) @@ -678,10 +677,10 @@ class TestAppAnnotationServiceBatchImport: assert "at least" in cast(str, result["error_msg"]) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_rejects_subscription_quota_overflow(self, sqlite_session: Session, current_user: Account) -> None: app = _persist_app(sqlite_session) features = SimpleNamespace( - billing=SimpleNamespace(enabled=True), annotation_quota_limit=SimpleNamespace(limit=1, size=1), ) @@ -694,10 +693,11 @@ class TestAppAnnotationServiceBatchImport: assert "exceeds the limit" in cast(str, result["error_msg"]) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_valid_import_enqueues_job(self, sqlite_session: Session, current_user: Account) -> None: app = _persist_app(sqlite_session) dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]}) - features = SimpleNamespace(billing=SimpleNamespace(enabled=False), annotation_quota_limit=None) + features = SimpleNamespace(annotation_quota_limit=None) with ( patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), @@ -705,10 +705,7 @@ class TestAppAnnotationServiceBatchImport: patch.object(annotation_service_module, "redis_client") as redis, patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-3"), patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)), - patch( - "configs.dify_config", - new=SimpleNamespace(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), - ), + config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), ): result = AppAnnotationService.batch_import_app_annotations( app.id, _file(b"question,answer\nq,a\n"), sqlite_session @@ -722,22 +719,20 @@ class TestAppAnnotationServiceBatchImport: "uuid-3", [{"question": "q1", "answer": "a1"}], app.id, TENANT_ID, current_user.id ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_unexpected_error_cleans_active_job( self, sqlite_session: Session, current_user: Account, caplog: pytest.LogCaptureFixture ) -> None: app = _persist_app(sqlite_session) dataframe = pd.DataFrame({"q": ["q1"], "a": ["a1"]}) - features = SimpleNamespace(billing=SimpleNamespace(enabled=False), annotation_quota_limit=None) + features = SimpleNamespace(annotation_quota_limit=None) with ( patch.object(annotation_service_module.pd, "read_csv", return_value=dataframe), patch.object(annotation_service_module.FeatureService, "get_features", return_value=features), patch.object(annotation_service_module, "redis_client") as redis, patch.object(annotation_service_module.uuid, "uuid4", return_value="uuid-4"), patch.object(annotation_service_module, "naive_utc_now", return_value=datetime.fromtimestamp(1)), - patch( - "configs.dify_config", - new=SimpleNamespace(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), - ), + config_overrides_context(ANNOTATION_IMPORT_MAX_RECORDS=5, ANNOTATION_IMPORT_MIN_RECORDS=1), ): redis.zadd.side_effect = RuntimeError("boom") redis.zrem.side_effect = RuntimeError("cleanup-failed") diff --git a/api/tests/unit_tests/services/test_batch_indexing_base.py b/api/tests/unit_tests/services/test_batch_indexing_base.py index 17420379427..929cd6251db 100644 --- a/api/tests/unit_tests/services/test_batch_indexing_base.py +++ b/api/tests/unit_tests/services/test_batch_indexing_base.py @@ -5,8 +5,9 @@ from unittest.mock import MagicMock, patch import pytest from core.entities.document_task import DocumentTask -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from services.document_indexing_proxy.batch_indexing_base import BatchDocumentIndexingProxy +from tests.unit_tests.config_override import config_overrides_context # --------------------------------------------------------------------------- # Concrete subclass for testing (the base class is abstract) @@ -275,13 +276,13 @@ class TestSendToTenantQueue: class TestDispatchRouting: """Tests for the _dispatch / delay routing logic inherited from the base class.""" - def _mock_features(self, enabled: bool, plan: CloudPlan) -> MagicMock: + def _mock_features(self, plan: CloudPlan) -> MagicMock: features = MagicMock() - features.billing.enabled = enabled features.billing.subscription.plan = plan return features - def test_should_send_to_normal_tenant_queue_when_billing_enabled_and_sandbox_plan(self) -> None: + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_should_send_to_normal_tenant_queue_in_cloud_with_sandbox_plan(self) -> None: """Sandbox plan routes to normal priority queue with tenant isolation.""" # Arrange proxy = make_proxy() @@ -289,7 +290,7 @@ class TestDispatchRouting: proxy._tenant_isolated_task_queue.get_task_key.return_value = None with patch("services.document_indexing_proxy.base.FeatureService.get_features") as mock_features: - mock_features.return_value = self._mock_features(enabled=True, plan=CloudPlan.SANDBOX) + mock_features.return_value = self._mock_features(plan=CloudPlan.SANDBOX) # Act with patch.object(proxy, "_send_to_default_tenant_queue") as mock_method: @@ -298,13 +299,14 @@ class TestDispatchRouting: # Assert mock_method.assert_called_once() - def test_should_send_to_priority_tenant_queue_when_billing_enabled_and_paid_plan(self) -> None: + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) + def test_should_send_to_priority_tenant_queue_in_cloud_with_paid_plan(self) -> None: """Non-sandbox paid plan routes to priority queue with tenant isolation.""" # Arrange proxy = make_proxy() with patch("services.document_indexing_proxy.base.FeatureService.get_features") as mock_features: - mock_features.return_value = self._mock_features(enabled=True, plan=CloudPlan.PROFESSIONAL) + mock_features.return_value = self._mock_features(plan=CloudPlan.PROFESSIONAL) # Act with patch.object(proxy, "_send_to_priority_tenant_queue") as mock_method: @@ -313,13 +315,14 @@ class TestDispatchRouting: # Assert mock_method.assert_called_once() - def test_should_send_to_priority_direct_queue_when_billing_not_enabled(self) -> None: + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_should_send_to_priority_direct_queue_outside_cloud(self) -> None: """Self-hosted / no billing → priority direct queue (no tenant isolation).""" # Arrange proxy = make_proxy() with patch("services.document_indexing_proxy.base.FeatureService.get_features") as mock_features: - mock_features.return_value = self._mock_features(enabled=False, plan=CloudPlan.SANDBOX) + mock_features.return_value = self._mock_features(plan=CloudPlan.SANDBOX) # Act with patch.object(proxy, "_send_to_priority_direct_queue") as mock_method: @@ -340,19 +343,20 @@ class TestDispatchRouting: # Assert mock_dispatch.assert_called_once() - def test_should_use_feature_service_for_billing_info(self) -> None: - """Verify that FeatureService.get_features is consulted during dispatch.""" + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) + def test_should_skip_feature_service_outside_cloud(self) -> None: + """Self-hosted dispatch does not load Cloud plan data.""" # Arrange proxy = make_proxy() with patch("services.document_indexing_proxy.base.FeatureService.get_features") as mock_features: - mock_features.return_value = self._mock_features(enabled=False, plan=CloudPlan.SANDBOX) + mock_features.return_value = self._mock_features(plan=CloudPlan.SANDBOX) with patch.object(proxy, "_send_to_priority_direct_queue"): # Act proxy._dispatch() # Assert - mock_features.assert_called_once_with(TENANT_ID, exclude_vector_space=True) + mock_features.assert_not_called() class TestBaseRouterHelpers: diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index cf587aed3fc..e8027fcda06 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -430,7 +430,6 @@ class TestBillingServiceSubscriptionInfo: # Arrange tenant_id = "tenant-123" expected_response = { - "enabled": True, "subscription": {"plan": "professional", "interval": "month", "education": False}, "members": {"size": 1, "limit": 50}, "apps": {"size": 1, "limit": 200}, @@ -458,7 +457,6 @@ class TestBillingServiceSubscriptionInfo: # Arrange tenant_id = "tenant-123" expected_response = { - "enabled": True, "subscription": {"plan": "professional", "interval": "month", "education": False}, "members": {"size": 1, "limit": 50}, "apps": {"size": 1, "limit": 200}, @@ -488,7 +486,6 @@ class TestBillingServiceSubscriptionInfo: # Arrange tenant_id = "tenant-123" expected_response = { - "enabled": True, "subscription": {"plan": "professional", "interval": "month", "education": False}, "members": {"size": 1, "limit": 50}, "apps": {"size": 1, "limit": 200}, @@ -544,7 +541,6 @@ class TestBillingServiceSubscriptionInfo: def test_get_info_preserves_unknown_vector_space_usage(self, mock_send_request): tenant_id = "tenant-123" expected_response = { - "enabled": True, "subscription": {"plan": "sandbox", "interval": "", "education": False}, "members": {"size": 1, "limit": 1}, "apps": {"size": 1, "limit": 10}, @@ -1750,7 +1746,6 @@ class TestBillingServiceIntegrationScenarios: # Step 1: Get current billing info mock_send_request.return_value = { - "enabled": True, "subscription": {"plan": "sandbox", "interval": "", "education": False}, "members": {"size": 0, "limit": 1}, "apps": {"size": 0, "limit": 5}, @@ -1822,7 +1817,6 @@ class TestBillingServiceSubscriptionInfoDataType: @pytest.fixture def normal_billing_response(self) -> dict: return { - "enabled": True, "subscription": { "plan": "team", "interval": "year", @@ -1844,7 +1838,6 @@ class TestBillingServiceSubscriptionInfoDataType: @pytest.fixture def string_billing_response(self) -> dict: return { - "enabled": True, "subscription": { "plan": "team", "interval": "year", @@ -1865,7 +1858,6 @@ class TestBillingServiceSubscriptionInfoDataType: @staticmethod def _assert_billing_info_types(result: dict): - assert isinstance(result["enabled"], bool) assert isinstance(result["subscription"]["plan"], str) assert isinstance(result["subscription"]["interval"], str) assert isinstance(result["subscription"]["education"], bool) diff --git a/api/tests/unit_tests/services/test_dataset_service_document.py b/api/tests/unit_tests/services/test_dataset_service_document.py index 18229ac3fb1..c70468e5f65 100644 --- a/api/tests/unit_tests/services/test_dataset_service_document.py +++ b/api/tests/unit_tests/services/test_dataset_service_document.py @@ -6,12 +6,14 @@ from datetime import datetime from sqlalchemy import event, select from sqlalchemy.orm import Session +from enums import DeploymentEdition from models.account import Tenant from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import UploadFile from models.source import DataSourceOauthBinding from services.dataset_ref_service import DatasetRefService +from tests.unit_tests.config_override import config_overrides_context from .dataset_service_test_helpers import ( Account, @@ -555,6 +557,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: with patch("services.dataset_service.current_user", account): yield account + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_without_dataset_id_creates_high_quality_dataset_with_default_retrieval_model( self, account_context, sqlite_session: Session ): @@ -581,7 +584,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: first_document = _document_row(name="VeryLongDocumentNameForDataset.txt") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch( "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", return_value=binding, @@ -617,6 +620,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: session=sqlite_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_without_dataset_id_uses_provided_retrieval_model( self, account_context, sqlite_session: Session ): @@ -644,7 +648,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: first_document = _document_row(name="Doc") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object( DocumentService, "save_document_with_dataset_id", @@ -662,6 +666,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: assert dataset.collection_binding_id is None assert sqlite_session.get(Dataset, dataset.id) is dataset + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_save_document_without_dataset_id_rejects_sandbox_batch_upload( self, account_context, unbound_session: Session ): @@ -678,7 +683,7 @@ class TestDocumentServiceSaveDocumentWithoutDatasetId: with ( patch( "services.dataset_service.FeatureService.get_features", - return_value=_make_features(enabled=True, plan=CloudPlan.SANDBOX), + return_value=_make_features(plan=CloudPlan.SANDBOX), ), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): @@ -1080,13 +1085,14 @@ class TestDocumentServiceSaveDocumentWithDatasetId: ): yield account + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_save_document_with_dataset_id_requires_file_info_for_upload_source( self, account_context, unbound_session: Session ): dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(file_ids=None) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=True)): + with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="File source info is required"): DocumentService.save_document_with_dataset_id( dataset, @@ -1095,6 +1101,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: session=unbound_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_save_document_with_dataset_id_blocks_batch_upload_for_sandbox_plan( self, account_context, unbound_session: Session ): @@ -1104,7 +1111,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: with ( patch( "services.dataset_service.FeatureService.get_features", - return_value=_make_features(enabled=True, plan=CloudPlan.SANDBOX), + return_value=_make_features(plan=CloudPlan.SANDBOX), ), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): @@ -1118,6 +1125,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: check_quota.assert_not_called() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_save_document_with_dataset_id_enforces_batch_upload_limit( self, account_context, @@ -1129,7 +1137,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: knowledge_config = _make_upload_knowledge_config(file_ids=["file-1", "file-2"]) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=True)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): with pytest.raises(ValueError, match="batch upload limit of 1"): @@ -1142,6 +1150,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: check_quota.assert_not_called() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_updates_existing_document_and_data_source_type( self, account_context, sqlite_session: Session ): @@ -1151,7 +1160,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: updated_document.batch = "batch-existing" with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object( DocumentService, "update_document_with_dataset_id", return_value=updated_document ) as update_document, @@ -1168,13 +1177,14 @@ class TestDocumentServiceSaveDocumentWithDatasetId: assert batch == "batch-existing" update_document.assert_called_once_with(dataset, knowledge_config, account_context, session=sqlite_session) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_requires_data_source_for_new_documents( self, account_context, unbound_session: Session ): dataset = _dataset_row() knowledge_config = _make_upload_knowledge_config(data_source=None) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): + with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="Data source is required when creating new documents"): DocumentService.save_document_with_dataset_id( dataset, @@ -1183,6 +1193,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: session=unbound_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_requires_existing_process_rule_for_custom_mode( self, account_context, sqlite_session: Session ): @@ -1194,7 +1205,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: process_rule=ProcessRule(mode="custom"), ) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): + with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="No process rule found"): DocumentService.save_document_with_dataset_id( dataset, @@ -1203,6 +1214,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: session=sqlite_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_rejects_invalid_indexing_technique( self, account_context, unbound_session: Session ): @@ -1214,7 +1226,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: indexing_technique="broken-technique", ) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): + with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): with pytest.raises(ValueError, match="Indexing technique is invalid"): DocumentService.save_document_with_dataset_id( dataset, @@ -1223,6 +1235,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: session=unbound_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_returns_empty_for_invalid_process_rule_mode( self, account_context, unbound_session: Session ): @@ -1230,7 +1243,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"]) knowledge_config.process_rule = SimpleNamespace(mode="unsupported-mode", rules=None) - with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)): + with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()): documents, batch = DocumentService.save_document_with_dataset_id( dataset, knowledge_config, @@ -1241,6 +1254,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: assert documents == [] assert batch == "" + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents( self, account_context, sqlite_session: Session ): @@ -1254,7 +1268,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, patch("services.dataset_service.DuplicateDocumentIndexingTaskProxy") as duplicate_proxy_cls, @@ -1285,6 +1299,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: duplicate_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-duplicate"]) duplicate_proxy_cls.return_value.delay.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_notion_import_truncates_names_and_cleans_removed_pages( self, account_context, sqlite_session: Session ): @@ -1330,7 +1345,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.clean_notion_document_task") as clean_task, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, @@ -1353,6 +1368,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: document_proxy_cls.assert_called_once_with(dataset.tenant_id, dataset.id, ["doc-new"]) document_proxy_cls.return_value.delay.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_website_crawl_truncates_long_urls( self, account_context, sqlite_session: Session ): @@ -1379,7 +1395,7 @@ class TestDocumentServiceSaveDocumentWithDatasetId: doc_language="English", ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, ): @@ -1711,6 +1727,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: with patch("services.dataset_service.current_user", account): yield account + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_save_document_without_dataset_id_counts_notion_pages_for_quota( self, account_context, @@ -1741,7 +1758,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: ) ), ) - features = _make_features(enabled=True) + features = _make_features() document = _document_row(name="Doc") with ( @@ -1763,6 +1780,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: check_quota.assert_called_once_with(3, features) assert sqlite_session.get(Dataset, dataset.id) is dataset + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_save_document_without_dataset_id_enforces_batch_limit_for_website_urls( self, account_context, @@ -1786,7 +1804,7 @@ class TestDocumentServiceSaveWithoutDatasetBilling: ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=True)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch.object(DocumentService, "check_documents_upload_quota") as check_quota, ): with pytest.raises(ValueError, match="batch upload limit of 1"): @@ -1936,6 +1954,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: ): yield account + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_initializes_high_quality_dataset_from_default_embedding_model( self, account_context, sqlite_session: Session ): @@ -1955,7 +1974,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: binding.id = "binding-1" with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.ModelManager") as model_manager_cls, patch( "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", @@ -1991,6 +2010,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: } get_binding.assert_called_once_with("default-provider", "default-embedding", sqlite_session) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_uses_explicit_embedding_and_retrieval_model( self, account_context, sqlite_session: Session ): @@ -2020,7 +2040,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: updated_document = _document_row(document_id="doc-1") with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.ModelManager") as model_manager_cls, patch( "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding", @@ -2038,6 +2058,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: assert dataset.embedding_model_provider == "explicit-provider" assert dataset.retrieval_model == knowledge_config.retrieval_model.model_dump() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_creates_custom_process_rule_for_new_upload_document( self, account_context, sqlite_session: Session ): @@ -2057,7 +2078,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.DocumentIndexingTaskProxy") as document_proxy_cls, patch("services.dataset_service.time.strftime", return_value="20260101010101"), @@ -2082,6 +2103,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: document_proxy_cls.assert_called_once_with("tenant-1", "dataset-1", [created_document.id]) document_proxy_cls.return_value.delay.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_creates_automatic_process_rule_for_new_upload_document( self, account_context, sqlite_session: Session ): @@ -2094,7 +2116,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.DocumentIndexingTaskProxy"), patch("services.dataset_service.time.strftime", return_value="20260101010101"), @@ -2114,6 +2136,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: assert created_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) assert sqlite_session.get(Document, documents[0].id) is documents[0] + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_creates_fallback_automatic_process_rule_when_latest_is_missing( self, account_context, sqlite_session: Session ): @@ -2123,7 +2146,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.DocumentIndexingTaskProxy"), patch("services.dataset_service.time.strftime", return_value="20260101010101"), @@ -2142,6 +2165,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: assert created_rule.mode == "automatic" assert created_rule.rules == json.dumps(DatasetProcessRule.AUTOMATIC_RULES) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_raises_when_upload_file_lookup_is_incomplete( self, account_context, sqlite_session: Session ): @@ -2151,7 +2175,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: sqlite_session.commit() with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, patch("services.dataset_service.time.strftime", return_value="20260101010101"), patch("services.dataset_service.secrets.randbelow", return_value=23), @@ -2165,6 +2189,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: session=sqlite_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_requires_notion_info_list_for_notion_import( self, account_context, sqlite_session: Session ): @@ -2180,7 +2205,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, ): mock_redis.lock.return_value = _make_lock_context() @@ -2193,6 +2218,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: session=sqlite_session, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_save_document_with_dataset_id_requires_website_info_list_for_website_crawl( self, account_context, sqlite_session: Session ): @@ -2208,7 +2234,7 @@ class TestDocumentServiceSaveDocumentAdditionalBranches: ) with ( - patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)), + patch("services.dataset_service.FeatureService.get_features", return_value=_make_features()), patch("services.dataset_service.redis_client") as mock_redis, ): mock_redis.lock.return_value = _make_lock_context() diff --git a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py index 76dc9f584d2..6049e76227a 100644 --- a/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py +++ b/api/tests/unit_tests/services/test_dataset_service_lock_not_owned.py @@ -42,9 +42,8 @@ def fake_current_user(monkeypatch: pytest.MonkeyPatch): @pytest.fixture def fake_features(monkeypatch: pytest.MonkeyPatch): - """Features.billing.enabled == False to skip quota logic.""" features = types.SimpleNamespace( - billing=types.SimpleNamespace(enabled=False, subscription=types.SimpleNamespace(plan="ENTERPRISE")), + billing=types.SimpleNamespace(subscription=types.SimpleNamespace(plan="ENTERPRISE")), documents_upload_quota=types.SimpleNamespace(limit=10_000, size=0), ) monkeypatch.setattr( diff --git a/api/tests/unit_tests/services/test_document_indexing_task_proxy.py b/api/tests/unit_tests/services/test_document_indexing_task_proxy.py index c012b132cd9..ace9abcb49e 100644 --- a/api/tests/unit_tests/services/test_document_indexing_task_proxy.py +++ b/api/tests/unit_tests/services/test_document_indexing_task_proxy.py @@ -2,19 +2,19 @@ from unittest.mock import Mock, patch from core.entities.document_task import DocumentTask from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy +from tests.unit_tests.config_override import config_overrides_context class DocumentIndexingTaskProxyTestDataFactory: """Factory class for creating test data and mock objects for DocumentIndexingTaskProxy tests.""" @staticmethod - def create_mock_features(billing_enabled: bool = False, plan: CloudPlan | str | None = CloudPlan.SANDBOX) -> Mock: + def create_mock_features(plan: CloudPlan | str | None = CloudPlan.SANDBOX) -> Mock: """Create mock features with billing configuration.""" features = Mock() features.billing = Mock() - features.billing.enabled = billing_enabled features.billing.subscription = Mock() features.billing.subscription.plan = plan return features @@ -171,13 +171,12 @@ class TestDocumentIndexingTaskProxy: # Assert proxy._send_to_direct_queue.assert_called_once_with(proxy.PRIORITY_TASK_FUNC) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") - def test_dispatch_with_billing_enabled_sandbox_plan(self, mock_feature_service): - """Test _dispatch method when billing is enabled with sandbox plan.""" + def test_dispatch_with_cloud_sandbox_plan(self, mock_feature_service): + """Test _dispatch method in Cloud with Sandbox plan.""" # Arrange - mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.SANDBOX - ) + mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) mock_feature_service.get_features.return_value = mock_features proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy() proxy._send_to_default_tenant_queue = Mock() @@ -188,13 +187,12 @@ class TestDocumentIndexingTaskProxy: # Assert proxy._send_to_default_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") - def test_dispatch_with_billing_enabled_non_sandbox_plan(self, mock_feature_service): - """Test _dispatch method when billing is enabled with non-sandbox plan.""" + def test_dispatch_with_cloud_paid_plan(self, mock_feature_service): + """Test _dispatch method in Cloud with a paid plan.""" # Arrange - mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.TEAM - ) + mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.TEAM) mock_feature_service.get_features.return_value = mock_features proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -205,11 +203,12 @@ class TestDocumentIndexingTaskProxy: # If billing enabled with non sandbox plan, should send to priority tenant queue proxy._send_to_priority_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) @patch("services.document_indexing_proxy.base.FeatureService") - def test_dispatch_with_billing_disabled(self, mock_feature_service): - """Test _dispatch method when billing is disabled.""" + def test_dispatch_outside_cloud(self, mock_feature_service): + """Test _dispatch method outside Cloud.""" # Arrange - mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(billing_enabled=False) + mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features() mock_feature_service.get_features.return_value = mock_features proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy() proxy._send_to_priority_direct_queue = Mock() @@ -220,13 +219,12 @@ class TestDocumentIndexingTaskProxy: # If billing disabled, for example: self-hosted or enterprise, should send to priority direct queue proxy._send_to_priority_direct_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_delay_method(self, mock_feature_service): """Test delay method integration.""" # Arrange - mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.SANDBOX - ) + mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) mock_feature_service.get_features.return_value = mock_features proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy() proxy._send_to_default_tenant_queue = Mock() @@ -253,11 +251,12 @@ class TestDocumentIndexingTaskProxy: assert task.dataset_id == dataset_id assert task.document_ids == document_ids + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_dispatch_edge_case_empty_plan(self, mock_feature_service): """Test _dispatch method with empty plan string.""" # Arrange - mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(billing_enabled=True, plan="") + mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan="") mock_feature_service.get_features.return_value = mock_features proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -268,11 +267,12 @@ class TestDocumentIndexingTaskProxy: # Assert proxy._send_to_priority_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_dispatch_edge_case_none_plan(self, mock_feature_service): """Test _dispatch method with None plan.""" # Arrange - mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(billing_enabled=True, plan=None) + mock_features = DocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=None) mock_feature_service.get_features.return_value = mock_features proxy = DocumentIndexingTaskProxyTestDataFactory.create_document_task_proxy() proxy._send_to_priority_tenant_queue = Mock() diff --git a/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py b/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py index 2af411b0d76..7a5039e94cf 100644 --- a/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py +++ b/api/tests/unit_tests/services/test_duplicate_document_indexing_task_proxy.py @@ -2,21 +2,21 @@ from unittest.mock import Mock, patch from core.entities.document_task import DocumentTask from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import ( DuplicateDocumentIndexingTaskProxy, ) +from tests.unit_tests.config_override import config_overrides_context class DuplicateDocumentIndexingTaskProxyTestDataFactory: """Factory class for creating test data and mock objects for DuplicateDocumentIndexingTaskProxy tests.""" @staticmethod - def create_mock_features(billing_enabled: bool = False, plan: CloudPlan | str | None = CloudPlan.SANDBOX) -> Mock: + def create_mock_features(plan: CloudPlan | str | None = CloudPlan.SANDBOX) -> Mock: """Create mock features with billing configuration.""" features = Mock() features.billing = Mock() - features.billing.enabled = billing_enabled features.billing.subscription = Mock() features.billing.subscription.plan = plan return features @@ -196,13 +196,12 @@ class TestDuplicateDocumentIndexingTaskProxy: # Assert proxy._send_to_direct_queue.assert_called_once_with(proxy.PRIORITY_TASK_FUNC) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") - def test_dispatch_with_billing_enabled_sandbox_plan(self, mock_feature_service): - """Test _dispatch method when billing is enabled with sandbox plan.""" + def test_dispatch_with_cloud_sandbox_plan(self, mock_feature_service): + """Test _dispatch method in Cloud with Sandbox plan.""" # Arrange - mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.SANDBOX - ) + mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() proxy._send_to_default_tenant_queue = Mock() @@ -213,13 +212,12 @@ class TestDuplicateDocumentIndexingTaskProxy: # Assert proxy._send_to_default_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") - def test_dispatch_with_billing_enabled_non_sandbox_plan(self, mock_feature_service): - """Test _dispatch method when billing is enabled with non-sandbox plan.""" + def test_dispatch_with_cloud_paid_plan(self, mock_feature_service): + """Test _dispatch method in Cloud with a paid plan.""" # Arrange - mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.TEAM - ) + mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.TEAM) mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -231,11 +229,12 @@ class TestDuplicateDocumentIndexingTaskProxy: # If billing enabled with non sandbox plan, should send to priority tenant queue proxy._send_to_priority_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) @patch("services.document_indexing_proxy.base.FeatureService") - def test_dispatch_with_billing_disabled(self, mock_feature_service): - """Test _dispatch method when billing is disabled.""" + def test_dispatch_outside_cloud(self, mock_feature_service): + """Test _dispatch method outside Cloud.""" # Arrange - mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features(billing_enabled=False) + mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features() mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() proxy._send_to_priority_direct_queue = Mock() @@ -247,13 +246,12 @@ class TestDuplicateDocumentIndexingTaskProxy: # If billing disabled, for example: self-hosted or enterprise, should send to priority direct queue proxy._send_to_priority_direct_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_delay_method(self, mock_feature_service): """Test delay method integration.""" # Arrange - mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.SANDBOX - ) + mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() proxy._send_to_default_tenant_queue = Mock() @@ -265,13 +263,12 @@ class TestDuplicateDocumentIndexingTaskProxy: # If billing enabled with sandbox plan, should send to default tenant queue proxy._send_to_default_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_dispatch_edge_case_empty_plan(self, mock_feature_service): """Test _dispatch method with empty plan string.""" # Arrange - mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan="" - ) + mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan="") mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -282,13 +279,12 @@ class TestDuplicateDocumentIndexingTaskProxy: # Assert proxy._send_to_priority_tenant_queue.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_dispatch_edge_case_none_plan(self, mock_feature_service): """Test _dispatch method with None plan.""" # Arrange - mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=None - ) + mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features(plan=None) mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -345,12 +341,13 @@ class TestDuplicateDocumentIndexingTaskProxy: assert proxy._document_ids == document_ids assert len(proxy._document_ids) == 100 + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.document_indexing_proxy.base.FeatureService") def test_dispatch_with_professional_plan(self, mock_feature_service): """Test _dispatch method when billing is enabled with professional plan.""" # Arrange mock_features = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.PROFESSIONAL + plan=CloudPlan.PROFESSIONAL ) mock_feature_service.get_features.return_value = mock_features proxy = DuplicateDocumentIndexingTaskProxyTestDataFactory.create_duplicate_document_task_proxy() diff --git a/api/tests/unit_tests/services/test_feature_entities.py b/api/tests/unit_tests/services/test_feature_entities.py index 4b04d837036..f3c833e63cb 100644 --- a/api/tests/unit_tests/services/test_feature_entities.py +++ b/api/tests/unit_tests/services/test_feature_entities.py @@ -2,7 +2,7 @@ import pytest from pydantic import ValidationError from enums import CloudPlan -from services.entities.feature_entities import LicenseLimitationModel, SubscriptionModel +from services.entities.feature_entities import FeatureModel, LicenseLimitationModel, SubscriptionModel def test_subscription_model_uses_the_cloud_plan_value_set() -> None: @@ -40,3 +40,11 @@ def test_license_limitation_availability( limitation = LicenseLimitationModel(enabled=enabled, size=size, limit=limit) assert limitation.is_available(required) is expected + + +def test_feature_response_exposes_subscription_without_legacy_billing_switch() -> None: + features = FeatureModel(billing={"subscription": {"plan": "professional", "interval": "month"}}) + + assert features.model_dump(mode="json")["billing"] == { + "subscription": {"plan": "professional", "interval": "month"}, + } diff --git a/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py b/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py index d790121f35c..cb16e88a8b0 100644 --- a/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py +++ b/api/tests/unit_tests/services/test_feature_service_human_input_email_delivery.py @@ -14,7 +14,6 @@ class HumanInputEmailDeliveryCase: name: str deployment_edition: DeploymentEdition tenant_id: str | None - billing_feature_enabled: bool plan: str expected: bool @@ -24,7 +23,6 @@ CASES = [ name="enterprise_edition", deployment_edition=DeploymentEdition.ENTERPRISE, tenant_id=None, - billing_feature_enabled=False, plan=CloudPlan.SANDBOX, expected=True, ), @@ -32,7 +30,6 @@ CASES = [ name="community_edition", deployment_edition=DeploymentEdition.COMMUNITY, tenant_id=None, - billing_feature_enabled=False, plan=CloudPlan.SANDBOX, expected=True, ), @@ -40,15 +37,6 @@ CASES = [ name="cloud_edition_requires_tenant", deployment_edition=DeploymentEdition.CLOUD, tenant_id=None, - billing_feature_enabled=True, - plan=CloudPlan.PROFESSIONAL, - expected=False, - ), - HumanInputEmailDeliveryCase( - name="billing_feature_off", - deployment_edition=DeploymentEdition.CLOUD, - tenant_id="tenant-1", - billing_feature_enabled=False, plan=CloudPlan.PROFESSIONAL, expected=False, ), @@ -56,7 +44,6 @@ CASES = [ name="professional_plan", deployment_edition=DeploymentEdition.CLOUD, tenant_id="tenant-1", - billing_feature_enabled=True, plan=CloudPlan.PROFESSIONAL, expected=True, ), @@ -64,7 +51,6 @@ CASES = [ name="team_plan", deployment_edition=DeploymentEdition.CLOUD, tenant_id="tenant-1", - billing_feature_enabled=True, plan=CloudPlan.TEAM, expected=True, ), @@ -72,7 +58,6 @@ CASES = [ name="sandbox_plan", deployment_edition=DeploymentEdition.CLOUD, tenant_id="tenant-1", - billing_feature_enabled=True, plan=CloudPlan.SANDBOX, expected=False, ), @@ -86,7 +71,6 @@ def test_resolve_human_input_email_delivery_enabled_matrix( ): config_overrides(DEPLOYMENT_EDITION=case.deployment_edition) features = FeatureModel() - features.billing.enabled = case.billing_feature_enabled features.billing.subscription.plan = case.plan result = FeatureService._resolve_human_input_email_delivery_enabled( diff --git a/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py b/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py index 5a38f1f71fd..cf6105373bb 100644 --- a/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py +++ b/api/tests/unit_tests/services/test_feature_service_knowledge_file_size_limit.py @@ -9,15 +9,14 @@ from services.feature_service import FeatureService @pytest.mark.parametrize( - ("deployment_edition", "tenant_id", "billing_feature_enabled", "plan", "expected"), + ("deployment_edition", "tenant_id", "plan", "expected"), [ - (DeploymentEdition.COMMUNITY, "tenant-1", True, CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.ENTERPRISE, "tenant-1", True, CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.CLOUD, None, True, CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.CLOUD, "tenant-1", False, CloudPlan.PROFESSIONAL, 15), - (DeploymentEdition.CLOUD, "tenant-1", True, CloudPlan.SANDBOX, 15), - (DeploymentEdition.CLOUD, "tenant-1", True, CloudPlan.PROFESSIONAL, 50), - (DeploymentEdition.CLOUD, "tenant-1", True, CloudPlan.TEAM, 50), + (DeploymentEdition.COMMUNITY, "tenant-1", CloudPlan.PROFESSIONAL, 15), + (DeploymentEdition.ENTERPRISE, "tenant-1", CloudPlan.PROFESSIONAL, 15), + (DeploymentEdition.CLOUD, None, CloudPlan.PROFESSIONAL, 15), + (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.SANDBOX, 15), + (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.PROFESSIONAL, 50), + (DeploymentEdition.CLOUD, "tenant-1", CloudPlan.TEAM, 50), ], ) def test_get_knowledge_file_size_limit( @@ -25,7 +24,6 @@ def test_get_knowledge_file_size_limit( config_overrides: Callable[..., None], deployment_edition: DeploymentEdition, tenant_id: str | None, - billing_feature_enabled: bool, plan: CloudPlan, expected: int, ) -> None: @@ -36,7 +34,6 @@ def test_get_knowledge_file_size_limit( ) get_info = Mock( return_value={ - "enabled": billing_feature_enabled, "subscription": {"plan": plan}, } ) @@ -62,7 +59,6 @@ def test_paid_knowledge_file_size_limit_never_reduces_default( feature_service_module.BillingService, "get_info", lambda *_args, **_kwargs: { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL}, }, ) diff --git a/api/tests/unit_tests/services/test_feature_service_vector_space.py b/api/tests/unit_tests/services/test_feature_service_vector_space.py index 5a7cfc1dda4..b61091fd730 100644 --- a/api/tests/unit_tests/services/test_feature_service_vector_space.py +++ b/api/tests/unit_tests/services/test_feature_service_vector_space.py @@ -10,7 +10,6 @@ from services.feature_service import FeatureService def test_get_features_exclude_vector_space_sets_vector_space_to_none(config_overrides): tenant_id = "tenant-id" billing_info = { - "enabled": True, "subscription": {"plan": CloudPlan.PROFESSIONAL, "interval": "monthly", "education": False}, "members": {"size": 1, "limit": 10}, "apps": {"size": 2, "limit": 20}, diff --git a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py index f89fb17f6e0..2b73d643387 100644 --- a/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py +++ b/api/tests/unit_tests/services/test_rag_pipeline_task_proxy.py @@ -7,22 +7,22 @@ import pytest from core.app.entities.rag_pipeline_invoke_entities import RagPipelineInvokeEntity from core.rag.pipeline.queue import TenantIsolatedTaskQueue -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from extensions.storage.storage_type import StorageType from models.enums import CreatorUserRole from models.model import UploadFile from services.rag_pipeline.rag_pipeline_task_proxy import RagPipelineTaskProxy +from tests.unit_tests.config_override import config_overrides_context class RagPipelineTaskProxyTestDataFactory: """Factory class for creating test data and mock objects for RagPipelineTaskProxy tests.""" @staticmethod - def create_mock_features(billing_enabled: bool = False, plan: CloudPlan = CloudPlan.SANDBOX) -> Mock: + def create_mock_features(plan: CloudPlan = CloudPlan.SANDBOX) -> Mock: """Create mock features with billing configuration.""" features = Mock() features.billing = Mock() - features.billing.enabled = billing_enabled features.billing.subscription = Mock() features.billing.subscription.plan = plan return features @@ -330,17 +330,16 @@ class TestRagPipelineTaskProxy: # Assert proxy._send_to_direct_queue.assert_called_once_with(upload_file_id, mock_task) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_with_billing_enabled_sandbox_plan( + def test_dispatch_with_cloud_sandbox_plan( self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock ): - """Test _dispatch method when billing is enabled with sandbox plan.""" + """Test _dispatch method in Cloud with Sandbox plan.""" # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.SANDBOX - ) + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) mock_feature_service.get_features.return_value = mock_features proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._send_to_default_tenant_queue = Mock() @@ -356,17 +355,14 @@ class TestRagPipelineTaskProxy: # If billing is enabled with sandbox plan, should send to default tenant queue proxy._send_to_default_tenant_queue.assert_called_once_with("file-123") + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_with_billing_enabled_non_sandbox_plan( - self, mock_db, mock_file_service_class, mock_feature_service - ): - """Test _dispatch method when billing is enabled with non-sandbox plan.""" + def test_dispatch_with_cloud_paid_plan(self, mock_db, mock_file_service_class, mock_feature_service): + """Test _dispatch method in Cloud with a paid plan.""" # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.TEAM - ) + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.TEAM) mock_feature_service.get_features.return_value = mock_features proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -382,15 +378,16 @@ class TestRagPipelineTaskProxy: # If billing is enabled with non-sandbox plan, should send to priority tenant queue proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") - def test_dispatch_with_billing_disabled( + def test_dispatch_outside_cloud( self, mock_db: MagicMock, mock_file_service_class: MagicMock, mock_feature_service: MagicMock ): - """Test _dispatch method when billing is disabled.""" + """Test _dispatch method outside Cloud.""" # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(billing_enabled=False) + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features() mock_feature_service.get_features.return_value = mock_features proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._send_to_priority_direct_queue = Mock() @@ -422,6 +419,7 @@ class TestRagPipelineTaskProxy: with pytest.raises(ValueError, match="upload_file_id is empty"): proxy._dispatch() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") @@ -430,7 +428,7 @@ class TestRagPipelineTaskProxy: ): """Test _dispatch method with empty plan string.""" # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(billing_enabled=True, plan="") + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan="") mock_feature_service.get_features.return_value = mock_features proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -446,6 +444,7 @@ class TestRagPipelineTaskProxy: # Assert proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") @@ -454,7 +453,7 @@ class TestRagPipelineTaskProxy: ): """Test _dispatch method with None plan.""" # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(billing_enabled=True, plan=None) + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=None) mock_feature_service.get_features.return_value = mock_features proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._send_to_priority_tenant_queue = Mock() @@ -470,6 +469,7 @@ class TestRagPipelineTaskProxy: # Assert proxy._send_to_priority_tenant_queue.assert_called_once_with("file-123") + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @patch("services.rag_pipeline.rag_pipeline_task_proxy.FeatureService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService") @patch("services.rag_pipeline.rag_pipeline_task_proxy.db") @@ -478,9 +478,7 @@ class TestRagPipelineTaskProxy: ): """Test delay method integration.""" # Arrange - mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features( - billing_enabled=True, plan=CloudPlan.SANDBOX - ) + mock_features = RagPipelineTaskProxyTestDataFactory.create_mock_features(plan=CloudPlan.SANDBOX) mock_feature_service.get_features.return_value = mock_features proxy = RagPipelineTaskProxyTestDataFactory.create_rag_pipeline_task_proxy() proxy._dispatch = Mock() diff --git a/api/tests/unit_tests/services/test_vector_space_admission_service.py b/api/tests/unit_tests/services/test_vector_space_admission_service.py index b29eb45f1be..267d35f7757 100644 --- a/api/tests/unit_tests/services/test_vector_space_admission_service.py +++ b/api/tests/unit_tests/services/test_vector_space_admission_service.py @@ -549,9 +549,22 @@ def test_billing_plan_lookup_excludes_vector_space_and_is_cached() -> None: service = VectorSpaceAdmissionService() with patch( "services.vector_space_admission_service.BillingService.get_info", - return_value={"enabled": True, "subscription": {"plan": "professional"}}, + return_value={"subscription": {"plan": "professional"}}, ) as get_info: assert service._get_plan("tenant-1") == CloudPlan.PROFESSIONAL assert service._get_plan("tenant-1") == CloudPlan.PROFESSIONAL get_info.assert_called_once_with("tenant-1", exclude_vector_space=True) + + +def test_unknown_billing_plan_is_not_logged(caplog: pytest.LogCaptureFixture) -> None: + service = VectorSpaceAdmissionService() + upstream_plan = "unexpected-private-plan" + with patch( + "services.vector_space_admission_service.BillingService.get_info", + return_value={"subscription": {"plan": upstream_plan}}, + ): + assert service._get_plan("tenant-1") is None + + assert "unknown plan tenant_id=tenant-1" in caplog.text + assert upstream_plan not in caplog.text diff --git a/api/tests/unit_tests/services/test_web_app_runtime_query_service.py b/api/tests/unit_tests/services/test_web_app_runtime_query_service.py index 5464c33ec91..487d0f80b67 100644 --- a/api/tests/unit_tests/services/test_web_app_runtime_query_service.py +++ b/api/tests/unit_tests/services/test_web_app_runtime_query_service.py @@ -2,6 +2,7 @@ from unittest.mock import MagicMock, create_autospec import pytest +from enums import DeploymentEdition from services.app_definition_query_service import AppSiteConfiguration from services.entities.feature_entities import FeatureModel from services.file_service import FileService @@ -62,6 +63,7 @@ def _runtime_record( def _service( runtime: MagicMock, *, + deployment_edition: DeploymentEdition = DeploymentEdition.COMMUNITY, file_service: MagicMock | None = None, workspace_features: MagicMock | None = None, ) -> WebAppRuntimeQueryService: @@ -75,6 +77,7 @@ def _service( file_service=file_service, workspace_features=workspace_features, files_url=_FILES_URL, + deployment_edition=deployment_edition, ) @@ -87,13 +90,25 @@ def test_get_bootstrap_rejects_unavailable_runtime(record: WebAppRuntimeRecord | _service(runtime).get_bootstrap("app-1") +@pytest.mark.parametrize( + ("deployment_edition", "copyright_enabled", "expected_copyright", "expected_placeholder"), + [ + (DeploymentEdition.CLOUD, False, None, None), + (DeploymentEdition.CLOUD, True, "Copyright", "Ask anything"), + (DeploymentEdition.COMMUNITY, False, "Copyright", "Ask anything"), + (DeploymentEdition.ENTERPRISE, False, "Copyright", "Ask anything"), + ], +) def test_get_bootstrap_applies_feature_and_branding_policy_after_record_load( workspace_features: MagicMock, + deployment_edition: DeploymentEdition, + copyright_enabled: bool, + expected_copyright: str | None, + expected_placeholder: str | None, ) -> None: runtime: MagicMock = create_autospec(WebAppRuntimeQuery, instance=True, spec_set=True) record = _runtime_record() - features = FeatureModel(can_replace_logo=True, webapp_copyright_enabled=False) - features.billing.enabled = True + features = FeatureModel(can_replace_logo=True, webapp_copyright_enabled=copyright_enabled) events: list[str] = [] runtime.get_runtime_record.side_effect = lambda _app_id: events.append("record") or record workspace_features.side_effect = lambda _tenant_id, **_kwargs: events.append("features") or features @@ -104,6 +119,7 @@ def test_get_bootstrap_applies_feature_and_branding_policy_after_record_load( runtime, file_service=file_service, workspace_features=workspace_features, + deployment_edition=deployment_edition, ).get_bootstrap("app-1") assert result == WebAppBootstrap( @@ -112,8 +128,8 @@ def test_get_bootstrap_applies_feature_and_branding_policy_after_record_load( enable_site=True, site={ **record.site._asdict(), - "copyright": None, - "input_placeholder": None, + "copyright": expected_copyright, + "input_placeholder": expected_placeholder, "icon_url": "https://icon", }, plan="pro", diff --git a/api/tests/unit_tests/services/test_workspace_credit_pool.py b/api/tests/unit_tests/services/test_workspace_credit_pool.py index 8d88bd0b28d..aef101dcdc4 100644 --- a/api/tests/unit_tests/services/test_workspace_credit_pool.py +++ b/api/tests/unit_tests/services/test_workspace_credit_pool.py @@ -27,7 +27,6 @@ def test_get_effective_credit_pool_prefers_available_paid_pool( quota_used=quota_used, ) billing_info = { - "enabled": True, "subscription": {"plan": CloudPlan.TEAM}, "next_credit_reset_date": 1775001600, } @@ -59,7 +58,6 @@ def test_get_effective_credit_pool_exposes_exhausted_trial_pool(unbound_session: exhausted_at=1772323200, ) billing_info = { - "enabled": True, "subscription": {"plan": CloudPlan.SANDBOX}, } config = SimpleNamespace(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) diff --git a/api/tests/unit_tests/services/test_workspace_service.py b/api/tests/unit_tests/services/test_workspace_service.py index 9d01cb7b847..3a63cfc985b 100644 --- a/api/tests/unit_tests/services/test_workspace_service.py +++ b/api/tests/unit_tests/services/test_workspace_service.py @@ -19,7 +19,6 @@ def test_get_current_workspace_summary_sandbox_uses_trial_only() -> None: quota_used=20, ) billing_info = { - "enabled": True, "subscription": {"plan": CloudPlan.SANDBOX}, } config = SimpleNamespace(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @@ -61,7 +60,6 @@ def test_get_current_workspace_summary_falls_back_from_exhausted_paid_pool() -> quota_used=40, ) billing_info = { - "enabled": True, "subscription": {"plan": CloudPlan.TEAM}, } config = SimpleNamespace(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) diff --git a/api/tests/unit_tests/tasks/test_dataset_indexing_task.py b/api/tests/unit_tests/tasks/test_dataset_indexing_task.py index e12e8b1b243..7ef6aaac1bf 100644 --- a/api/tests/unit_tests/tasks/test_dataset_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_dataset_indexing_task.py @@ -15,7 +15,7 @@ from sqlalchemy.orm import Session from core.indexing_runner import DocumentIsPausedError from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from enums import CloudPlan +from enums import CloudPlan, DeploymentEdition from extensions.ext_redis import redis_client from models.dataset import Dataset, Document from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus @@ -27,7 +27,7 @@ from tasks.document_indexing_task import ( normal_document_indexing_task, priority_document_indexing_task, ) -from tests.unit_tests.config_override import apply_config_overrides +from tests.unit_tests.config_override import apply_config_overrides, config_overrides_context @pytest.fixture @@ -68,13 +68,12 @@ def indexing_runner(monkeypatch: pytest.MonkeyPatch) -> MagicMock: def _features( *, - billing_enabled: bool = False, plan: CloudPlan = CloudPlan.PROFESSIONAL, vector_limit: int = 1000, vector_size: int = 0, ) -> SimpleNamespace: return SimpleNamespace( - billing=SimpleNamespace(enabled=billing_enabled, subscription=SimpleNamespace(plan=plan)), + billing=SimpleNamespace(subscription=SimpleNamespace(plan=plan)), vector_space=SimpleNamespace(limit=vector_limit, size=vector_size), ) @@ -137,6 +136,7 @@ def _persisted_documents(session: Session, document_ids: list[str]) -> list[Docu class TestTaskEnqueuing: + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_self_hosted_dispatches_directly_to_priority_task( self, tenant_id: str, dataset_id: str, document_ids: list[str], mock_redis: MagicMock ) -> None: @@ -144,7 +144,6 @@ class TestTaskEnqueuing: patch.object(DocumentIndexingTaskProxy, "features") as features, patch.object(DocumentIndexingTaskProxy, "PRIORITY_TASK_FUNC", Mock()) as task, ): - features.billing.enabled = False DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids).delay() task.delay.assert_called_once_with( @@ -153,6 +152,7 @@ class TestTaskEnqueuing: document_ids=document_ids, ) + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @pytest.mark.parametrize( ("plan", "task_attribute"), [ @@ -173,13 +173,13 @@ class TestTaskEnqueuing: patch.object(DocumentIndexingTaskProxy, "features") as features, patch.object(DocumentIndexingTaskProxy, task_attribute, Mock()) as task, ): - features.billing.enabled = True features.billing.subscription.plan = plan DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids).delay() mock_redis.setex.assert_called() task.delay.assert_called_once() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) def test_running_tenant_task_queues_followup_work( self, tenant_id: str, dataset_id: str, document_ids: list[str], mock_redis: MagicMock ) -> None: @@ -188,7 +188,6 @@ class TestTaskEnqueuing: patch.object(DocumentIndexingTaskProxy, "features") as features, patch.object(DocumentIndexingTaskProxy, "PRIORITY_TASK_FUNC", Mock()) as task, ): - features.billing.enabled = True features.billing.subscription.plan = CloudPlan.PROFESSIONAL DocumentIndexingTaskProxy(tenant_id, dataset_id, document_ids).delay() @@ -285,12 +284,13 @@ class TestDocumentIndexing: get_features.assert_not_called() runner_class.assert_not_called() + @config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD) @pytest.mark.parametrize( ("features", "batch_limit", "message"), [ - (_features(billing_enabled=True), 1, "batch upload limit"), - (_features(billing_enabled=True, plan=CloudPlan.SANDBOX), 100, "does not support batch upload"), - (_features(billing_enabled=True, vector_limit=100, vector_size=100), 100, "over the limit"), + (_features(), 1, "batch upload limit"), + (_features(plan=CloudPlan.SANDBOX), 100, "does not support batch upload"), + (_features(vector_limit=100, vector_size=100), 100, "over the limit"), ], ) def test_validation_failure_marks_every_scoped_document_error( diff --git a/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py b/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py index e5e1961112d..47d02381da9 100644 --- a/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_retry_document_indexing_task.py @@ -4,13 +4,16 @@ from uuid import uuid4 from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType +from enums import DeploymentEdition from models import Account, Tenant, TenantAccountJoin from models.account import TenantAccountRole from models.dataset import Dataset, Document from models.enums import DatasetRuntimeMode, DataSourceType, DocumentCreatedFrom, IndexingStatus from tasks.retry_document_indexing_task import retry_document_indexing_task +from tests.unit_tests.config_override import config_overrides_context +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_retry_enforces_vector_space_admission(sqlite_session: Session) -> None: tenant = Tenant(name="Retry tenant") user = Account(name="Retry user", email=f"retry-{uuid4()}@example.com") @@ -49,7 +52,6 @@ def test_retry_enforces_vector_space_admission(sqlite_session: Session) -> None: sqlite_session.add_all([tenant, user, membership, dataset, document]) sqlite_session.commit() features = MagicMock() - features.billing.enabled = False with ( patch("tasks.retry_document_indexing_task.FeatureService.get_features", return_value=features), diff --git a/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py b/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py index b5f88f1a706..041e9ba55e5 100644 --- a/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py +++ b/api/tests/unit_tests/tasks/test_sync_website_document_indexing_task.py @@ -5,9 +5,11 @@ from sqlalchemy import select from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType +from enums import DeploymentEdition from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom from tasks.sync_website_document_indexing_task import sync_website_document_indexing_task +from tests.unit_tests.config_override import config_overrides_context def _dataset(tenant_id: str) -> Dataset: @@ -66,6 +68,7 @@ def test_rejects_document_outside_dataset_before_side_effects(sqlite_session: Se processor_factory.assert_not_called() +@config_overrides_context(DEPLOYMENT_EDITION=DeploymentEdition.COMMUNITY) def test_cleanup_is_owner_scoped_and_skips_empty_vector_ids(sqlite_session: Session) -> None: tenant_id = str(uuid.uuid4()) dataset = _dataset(tenant_id) @@ -83,7 +86,6 @@ def test_cleanup_is_owner_scoped_and_skips_empty_vector_ids(sqlite_session: Sess sqlite_session.commit() features = MagicMock() - features.billing.enabled = False with ( patch("tasks.sync_website_document_indexing_task.FeatureService.get_features", return_value=features), patch("tasks.sync_website_document_indexing_task.IndexProcessorFactory") as processor_factory, diff --git a/packages/contracts/generated/api/console/features/types.gen.ts b/packages/contracts/generated/api/console/features/types.gen.ts index 02c727aa04d..ecdf566866c 100644 --- a/packages/contracts/generated/api/console/features/types.gen.ts +++ b/packages/contracts/generated/api/console/features/types.gen.ts @@ -46,7 +46,6 @@ export type Quota = { } export type BillingModel = { - enabled: boolean subscription: SubscriptionModel } diff --git a/packages/contracts/generated/api/console/features/zod.gen.ts b/packages/contracts/generated/api/console/features/zod.gen.ts index 8978ebf9cae..a4aa76550e4 100644 --- a/packages/contracts/generated/api/console/features/zod.gen.ts +++ b/packages/contracts/generated/api/console/features/zod.gen.ts @@ -79,7 +79,6 @@ export const zSubscriptionModel = z.object({ * BillingModel */ export const zBillingModel = z.object({ - enabled: z.boolean().default(false), subscription: zSubscriptionModel.default({ interval: '', plan: 'sandbox' }), }) @@ -94,10 +93,7 @@ export const zFeatureModel = z.object({ usage: 0, }), apps: zLimitationModel.default({ limit: 10, size: 0 }), - billing: zBillingModel.default({ - enabled: false, - subscription: { interval: '', plan: 'sandbox' }, - }), + billing: zBillingModel.default({ subscription: { interval: '', plan: 'sandbox' } }), can_replace_logo: z.boolean().default(false), dataset_operator_enabled: z.boolean().default(false), docs_processing: z.string().default('standard'), diff --git a/web/__tests__/billing/pricing-modal-flow.test.tsx b/web/__tests__/billing/pricing-modal-flow.test.tsx index 6f32dc9d044..9a77532723d 100644 --- a/web/__tests__/billing/pricing-modal-flow.test.tsx +++ b/web/__tests__/billing/pricing-modal-flow.test.tsx @@ -32,7 +32,6 @@ const render = (ui: React.ReactElement) => { }) seedFeatures(queryClient, { billing: { - enabled: true, subscription: { interval: 'month', plan: mockCurrentPlan }, }, education: { enabled: mockEducationEnabled }, diff --git a/web/app/components/billing/utils/__tests__/index.spec.ts b/web/app/components/billing/utils/__tests__/index.spec.ts index c33312a9209..2d8bf044d81 100644 --- a/web/app/components/billing/utils/__tests__/index.spec.ts +++ b/web/app/components/billing/utils/__tests__/index.spec.ts @@ -61,7 +61,6 @@ describe('billing utils', () => { limit: 5, }, billing: { - enabled: true, subscription: { interval: '', plan: 'sandbox', @@ -144,7 +143,6 @@ describe('billing utils', () => { it('should derive vector space total from plan config', () => { const data = createMockPlanData({ billing: { - enabled: true, subscription: { interval: '', plan: 'professional',