refactor(api): remove legacy billing enabled state (#41917)

This commit is contained in:
yyh 2026-09-07 08:29:57 +00:00 committed by GitHub
parent fc0136647f
commit eba589db89
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
66 changed files with 485 additions and 518 deletions

View File

@ -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)

View File

@ -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()

View File

@ -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)

View File

@ -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 {})

View File

@ -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)

View File

@ -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

View File

@ -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),

View File

@ -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

View File

@ -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.")

View File

@ -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

View File

@ -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 = (

View File

@ -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()

View File

@ -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()

View File

@ -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)

View File

@ -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)

View File

@ -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

View File

@ -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

View File

@ -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(

View File

@ -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():

View File

@ -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(

View File

@ -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(

View File

@ -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:

View File

@ -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

View File

@ -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"}

View File

@ -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):

View File

@ -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",

View File

@ -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

View File

@ -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):

View File

@ -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

View File

@ -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]

View File

@ -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
):

View File

@ -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)

View File

@ -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},
}

View File

@ -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,

View File

@ -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

View File

@ -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()

View File

@ -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)

View File

@ -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)

View File

@ -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()

View File

@ -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

View File

@ -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

View File

@ -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),

View File

@ -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"):

View File

@ -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")

View File

@ -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:

View File

@ -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)

View File

@ -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()

View File

@ -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(

View File

@ -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()

View File

@ -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()

View File

@ -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"},
}

View File

@ -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(

View File

@ -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},
},
)

View File

@ -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},

View File

@ -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()

View File

@ -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

View File

@ -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",

View File

@ -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)

View File

@ -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)

View File

@ -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(

View File

@ -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),

View File

@ -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,

View File

@ -46,7 +46,6 @@ export type Quota = {
}
export type BillingModel = {
enabled: boolean
subscription: SubscriptionModel
}

View File

@ -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'),

View File

@ -32,7 +32,6 @@ const render = (ui: React.ReactElement) => {
})
seedFeatures(queryClient, {
billing: {
enabled: true,
subscription: { interval: 'month', plan: mockCurrentPlan },
},
education: { enabled: mockEducationEnabled },

View File

@ -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',