mirror of
https://github.com/langgenius/dify.git
synced 2026-09-08 02:43:49 +08:00
refactor(api): remove legacy billing enabled state (#41917)
This commit is contained in:
parent
fc0136647f
commit
eba589db89
@ -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)
|
||||
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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 {})
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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.")
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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 = (
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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()
|
||||
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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():
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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"}
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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]
|
||||
|
||||
@ -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
|
||||
):
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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},
|
||||
}
|
||||
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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"):
|
||||
|
||||
@ -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")
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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"},
|
||||
}
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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},
|
||||
},
|
||||
)
|
||||
|
||||
@ -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},
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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),
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -46,7 +46,6 @@ export type Quota = {
|
||||
}
|
||||
|
||||
export type BillingModel = {
|
||||
enabled: boolean
|
||||
subscription: SubscriptionModel
|
||||
}
|
||||
|
||||
|
||||
@ -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'),
|
||||
|
||||
@ -32,7 +32,6 @@ const render = (ui: React.ReactElement) => {
|
||||
})
|
||||
seedFeatures(queryClient, {
|
||||
billing: {
|
||||
enabled: true,
|
||||
subscription: { interval: 'month', plan: mockCurrentPlan },
|
||||
},
|
||||
education: { enabled: mockEducationEnabled },
|
||||
|
||||
@ -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',
|
||||
|
||||
Loading…
Reference in New Issue
Block a user