import datetime import json from collections.abc import Callable from contextlib import ExitStack from inspect import unwrap from types import SimpleNamespace from unittest.mock import ANY, MagicMock, PropertyMock, call, patch import pytest from flask import Flask from sqlalchemy.orm import Session from werkzeug.exceptions import BadRequest, Forbidden, NotFound import services from controllers.console import console_ns from controllers.console.app.error import ProviderNotInitializeError from controllers.console.datasets.datasets import ( DatasetApi, DatasetApiBaseUrlApi, DatasetApiDeleteApi, DatasetApiKeyApi, DatasetAutoDisableLogApi, DatasetCreatePayload, DatasetEnableApiApi, DatasetErrorDocs, DatasetIndexingEstimateApi, DatasetIndexingStatusApi, DatasetKnowledgeFSUpgradeApi, DatasetKnowledgeFSUpgradeJobApi, DatasetKnowledgeFSUpgradeJobsApi, DatasetListApi, DatasetPermissionUserListApi, DatasetQueryApi, DatasetRelatedAppListApi, DatasetRetrievalSettingApi, DatasetRetrievalSettingMockApi, DatasetUpdatePayload, DatasetUseCheckApi, IndexingEstimatePayload, _get_retrieval_methods_by_vector_type, ) from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError from core.entities.knowledge_entities import IndexingEstimate from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.provider_manager import ProviderManager from core.rag.datasource.vdb.vector_type import VectorType from core.rag.index_processor.constant.index_type import IndexStructureType from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.storage.storage_type import StorageType from models.account import Account, TenantAccountRole from models.dataset import AppDatasetJoin, Dataset, DatasetPermission, DatasetQuery, Document, DocumentSegment from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus from models.knowledge_fs import KnowledgeFSUpgradeJobStatus, KnowledgeFSUpgradeStage from models.model import ApiToken, App, AppMode, IconType, UploadFile from services.dataset_ref_service import DatasetRef from services.dataset_service import DatasetPermissionService, DatasetService from services.enterprise import rbac_service as enterprise_rbac_service from tests.unit_tests.config_override import apply_config_overrides @pytest.fixture(autouse=True) def dataset_model_property_defaults(): getter_values: dict[str, object] = { "get_app_count": 0, "get_document_count": 0, "get_word_count": 0, "get_author_name": None, "get_tags": [], "get_doc_form": None, "get_external_knowledge_info": None, "get_doc_metadata": [], "get_is_published": False, "get_total_documents": 0, "get_total_available_documents": 0, } getters = {} with ExitStack() as stack: for name, value in getter_values.items(): getters[name] = stack.enter_context(patch.object(Dataset, name, autospec=True, return_value=value)) stack.enter_context( patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=enterprise_rbac_service.MyPermissionsResponse(), ) ) stack.enter_context( patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetPermissions.batch_get", return_value={}, ) ) yield getters def make_dataset(**overrides) -> Dataset: base = { "id": "ds-1", "tenant_id": "tenant-1", "name": "Dataset", "description": "desc", "provider": "vendor", "permission": "only_me", "data_source_type": None, "indexing_technique": "economy", "created_by": "account-1", "created_at": datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.UTC), "updated_by": None, "updated_at": datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.UTC), "embedding_model": None, "embedding_model_provider": None, "retrieval_model": None, "summary_index_setting": None, "built_in_field_enabled": False, "pipeline_id": None, "runtime_mode": "general", "chunk_structure": None, "icon_info": None, "enable_api": False, "is_multimodal": False, } base.update(overrides) return Dataset(**base) def make_account(role: TenantAccountRole = TenantAccountRole.EDITOR) -> Account: account = Account(name="Test User", email="user@example.com") account.id = "account-1" account.role = role return account def make_related_app(**overrides) -> App: base = { "id": "app-1", "tenant_id": "tenant-1", "name": "App", "description": "desc", "mode": AppMode.CHAT, "icon_type": IconType.EMOJI, "icon": "🤖", "icon_background": "#fff", "app_model_config_id": None, "workflow_id": None, "enable_site": False, "enable_api": False, "created_by": "account-1", } base.update(overrides) return App(**base) def make_document_status(**overrides) -> Document: base = { "id": "doc-1", "tenant_id": "tenant-1", "dataset_id": "dataset-1", "position": 1, "data_source_type": DataSourceType.UPLOAD_FILE, "batch": "batch-1", "name": "doc.txt", "created_from": DocumentCreatedFrom.WEB, "created_by": "account-1", "indexing_status": IndexingStatus.COMPLETED, "enabled": True, "archived": False, "processing_started_at": None, "parsing_completed_at": None, "cleaning_completed_at": None, "splitting_completed_at": None, "completed_at": None, "paused_at": None, "error": None, "stopped_at": None, } base.update(overrides) return Document(**base) def make_document_segment(*, position: int, completed: bool) -> DocumentSegment: return DocumentSegment( tenant_id="tenant-1", dataset_id="dataset-1", document_id="doc-1", position=position, content=f"segment {position}", word_count=2, tokens=2, created_by="account-1", completed_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) if completed else None, ) class _UsesSQLiteSession: session: Session @pytest.fixture(autouse=True) def _inject_sqlite_session(self, sqlite_session: Session) -> None: self.session = sqlite_session class TestDatasetList(_UsesSQLiteSession): def _mock_user(self): user = make_account() return user def test_get_success_basic(self, app: Flask, monkeypatch: pytest.MonkeyPatch): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset(icon_info={"icon": "📙", "icon_type": "emoji"})] apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=True) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 assert resp["total"] == 1 assert resp["data"][0]["embedding_available"] is True assert resp["data"][0]["icon_info"] == { "icon": "📙", "icon_background": None, "icon_type": "emoji", "icon_url": None, } assert resp["data"][0]["knowledge_fs_upgrade"] == { "job": None, "can_upgrade": True, "can_retry": False, "block_reason": None, } def test_get_attaches_upgrade_summaries_with_one_batch_lookup(self, app: Flask, monkeypatch: pytest.MonkeyPatch): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset(id="dataset-1"), make_dataset(id="dataset-2")] succeeded = SimpleNamespace( id="upgrade-job-1", old_dataset_id="dataset-1", new_control_space_id="space-1", status=KnowledgeFSUpgradeJobStatus.SUCCEEDED, stage=KnowledgeFSUpgradeStage.COMPLETED, snapshot_at=datetime.datetime(2026, 8, 18, tzinfo=datetime.UTC), total_documents=2, completed_documents=2, total_sources=1, completed_sources=1, last_error_code=None, last_error_message=None, completed_at=datetime.datetime(2026, 8, 18, tzinfo=datetime.UTC), ) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=True) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 2)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), patch( "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService.get_latest_by_dataset_ids", return_value={"dataset-1": succeeded}, ) as get_latest, ): resp, status = method(api, MagicMock(), "tenant-1", current_user) assert status == 200 get_latest.assert_called_once() assert resp["data"][0]["knowledge_fs_upgrade"]["block_reason"] is None assert resp["data"][0]["knowledge_fs_upgrade"]["can_upgrade"] is True assert resp["data"][0]["knowledge_fs_upgrade"]["job"]["new_control_space_id"] == "space-1" assert resp["data"][1]["knowledge_fs_upgrade"]["can_upgrade"] is True def test_get_serializes_database_fields_with_caller_session(self, app: Flask, dataset_model_property_defaults): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() dataset = make_dataset() session = self.session with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([dataset], 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): method(api, session, "tenant-1", current_user) for getter in dataset_model_property_defaults.values(): getter.assert_called_once_with(dataset, session=session) def test_get_with_ids_filter(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset()] with app.test_request_context("/datasets?ids=1&ids=2"): with ( patch.object(DatasetService, "get_datasets_by_ids", return_value=(datasets, 2)) as by_ids_mock, patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, self.session, "tenant-1", current_user) by_ids_mock.assert_called_once() assert status == 200 assert resp["total"] == 2 def test_get_with_creator_ids_filter(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset()] with app.test_request_context("/datasets?creator_ids=account-1&creator_ids=account-2"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)) as get_datasets, patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, MagicMock(), "tenant-1", current_user) assert get_datasets.call_args.kwargs["creator_ids"] == ["account-1", "account-2"] assert status == 200 assert resp["total"] == 1 def test_get_attaches_current_user_permission_keys(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() dataset = make_dataset(id="dataset-1") permissions = enterprise_rbac_service.MyPermissionsResponse( dataset=enterprise_rbac_service.ResourcePermissionSnapshot( default_permission_keys=["dataset.acl.readonly"], overrides=[ enterprise_rbac_service.ResourcePermissionKeys( resource_id="dataset-1", permission_keys=["dataset.acl.readonly", "dataset.acl.edit"] ) ], ) ) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([dataset], 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=permissions, ) as get_permissions, ): resp, status = method(api, self.session, "tenant-1", current_user) get_permissions.assert_called_once_with("tenant-1", current_user.id, session=ANY) assert status == 200 assert resp["data"][0]["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] def test_get_restricted_whitelist_blocks_own_dataset_fallback( self, app: Flask, config_overrides: Callable[..., None] ): config_overrides(RBAC_ENABLED=True) api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() permissions = enterprise_rbac_service.MyPermissionsResponse( workspace=enterprise_rbac_service.WorkspacePermissionSnapshot( permission_keys=["dataset.create_and_management"] ) ) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=permissions, ), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", return_value=SimpleNamespace(resource_ids=[]), ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): method(api, self.session, "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [] assert get_datasets.call_args.kwargs["include_own_datasets"] is False def test_get_default_read_is_unrestricted_when_whitelist_unrestricted( self, app: Flask, config_overrides: Callable[..., None] ): config_overrides(RBAC_ENABLED=True) api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() permissions = enterprise_rbac_service.MyPermissionsResponse( dataset=enterprise_rbac_service.ResourcePermissionSnapshot(default_permission_keys=["dataset.preview"]) ) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=permissions, ), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", return_value=SimpleNamespace(unrestricted=True, resource_ids=[]), ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): method(api, self.session, "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] is None def test_get_restricted_whitelist_overrides_default_read_permission( self, app: Flask, config_overrides: Callable[..., None] ): config_overrides(RBAC_ENABLED=True) api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() permissions = enterprise_rbac_service.MyPermissionsResponse( dataset=enterprise_rbac_service.ResourcePermissionSnapshot(default_permission_keys=["dataset.preview"]) ) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=permissions, ), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", return_value=SimpleNamespace(resource_ids=["dataset-whitelist-only"]), ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): method(api, MagicMock(), "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == ["dataset-whitelist-only"] assert get_datasets.call_args.kwargs["include_own_datasets"] is False def test_get_restricted_whitelist_ignores_dataset_read_overrides( self, app: Flask, config_overrides: Callable[..., None] ): config_overrides(RBAC_ENABLED=True) api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() permissions = enterprise_rbac_service.MyPermissionsResponse( dataset=enterprise_rbac_service.ResourcePermissionSnapshot( overrides=[ enterprise_rbac_service.ResourcePermissionKeys( resource_id="dataset-acl-shared", permission_keys=["dataset.acl.preview"] ), enterprise_rbac_service.ResourcePermissionKeys( resource_id="dataset-full", permission_keys=["dataset.full_access"] ), enterprise_rbac_service.ResourcePermissionKeys( resource_id="dataset-shared", permission_keys=["dataset.preview"] ), enterprise_rbac_service.ResourcePermissionKeys(resource_id="dataset-hidden", permission_keys=[]), ] ) ) with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=([], 0)) as get_datasets, patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=permissions, ), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", return_value=SimpleNamespace(resource_ids=["dataset-whitelist-only"]), ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): method(api, self.session, "tenant-1", current_user) assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [ "dataset-whitelist-only", ] assert get_datasets.call_args.kwargs["include_own_datasets"] is False def test_get_with_ids_applies_dataset_visibility(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(RBAC_ENABLED=True) api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() permissions = enterprise_rbac_service.MyPermissionsResponse() with app.test_request_context("/datasets?ids=dataset-1"): with ( patch.object(DatasetService, "get_datasets_by_ids", return_value=([], 0)) as get_datasets_by_ids, patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=permissions, ), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources", return_value=SimpleNamespace(resource_ids=[]), ), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): method(api, self.session, "tenant-1", current_user) session = get_datasets_by_ids.call_args.kwargs["session"] assert session is self.session assert get_datasets_by_ids.call_args.args == (["dataset-1"], "tenant-1") assert get_datasets_by_ids.call_args.kwargs == { "user": current_user, "accessible_dataset_ids": [], "include_own_datasets": False, "session": session, } def test_get_with_tag_ids(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset()] with app.test_request_context("/datasets?tag_ids=tag1"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 def test_get_allows_legacy_weighted_score_without_weight_type(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [ make_dataset( retrieval_model={ "search_method": "hybrid_search", "reranking_enable": True, "reranking_mode": "weighted_score", "reranking_model": None, "weights": { "vector_setting": { "vector_weight": 0.7, "embedding_model_name": "text-embedding", "embedding_provider_name": "openai", }, "keyword_setting": {"keyword_weight": 0.3}, }, "top_k": 3, "score_threshold_enabled": False, "score_threshold": 0.0, } ) ] with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 assert resp["data"][0]["retrieval_model_dict"]["weights"]["weight_type"] is None def test_get_merges_partial_retrieval_model_defaults(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset(retrieval_model={"top_k": 4, "score_threshold_enabled": False})] with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, self.session, "tenant-1", current_user) assert status == 200 retrieval_model = resp["data"][0]["retrieval_model_dict"] assert retrieval_model["search_method"] == "semantic_search" assert retrieval_model["reranking_enable"] is False assert retrieval_model["top_k"] == 4 assert retrieval_model["score_threshold_enabled"] is False def test_embedding_available_false(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [ make_dataset( indexing_technique="high_quality", embedding_model="text-embed", embedding_model_provider="openai" ) ] config = MagicMock() config.get_models.return_value = [] with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=config), ): resp, status = method(api, self.session, "tenant-1", current_user) assert resp["data"][0]["embedding_available"] is False def test_partial_members_permission(self, app: Flask): api = DatasetListApi() method = unwrap(api.get) current_user = self._mock_user() datasets = [make_dataset(permission="partial_members")] session = self.session session.add(DatasetPermission(dataset_id="ds-1", account_id="u1", tenant_id="tenant-1")) session.flush() with app.test_request_context("/datasets"): with ( patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)), patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])), ): resp, status = method(api, session, "tenant-1", current_user) assert resp["data"][0]["partial_member_list"] == ["u1"] class TestDatasetListApiPost(_UsesSQLiteSession): def test_post_success(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) payload = {"name": "My Dataset", "description": "desc", "indexing_technique": "economy", "provider": "vendor"} user = make_account() dataset = make_dataset(name=payload["name"], description=payload["description"]) with ( app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "create_empty_dataset", return_value=dataset), ): _, status = method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) assert status == 201 def test_post_forbidden(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) payload = {"name": "test"} user = make_account(TenantAccountRole.NORMAL) with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(Forbidden): method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) def test_post_duplicate_name(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) payload = {"name": "duplicate"} user = make_account() with ( app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload), patch.object( DatasetService, "create_empty_dataset", side_effect=services.errors.dataset.DatasetNameDuplicateError() ), ): with pytest.raises(DatasetNameDuplicateError): method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user) def test_post_invalid_payload_missing_name(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) with app.test_request_context("/datasets", json={}), patch.object(type(console_ns), "payload", {}): with pytest.raises(ValueError): method(api, DatasetCreatePayload(), self.session, "tenant-1", make_account()) def test_post_invalid_indexing_technique(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) payload = {"name": "bad", "indexing_technique": "invalid-tech"} with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(ValueError, match="Invalid indexing technique"): method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account()) def test_post_invalid_provider(self, app: Flask): api = DatasetListApi() method = unwrap(api.post) payload = {"name": "bad", "provider": "unknown"} with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload): with pytest.raises(ValueError, match="Invalid provider"): method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account()) class TestDatasetApiGet(_UsesSQLiteSession): def test_get_success_basic(self, app: Flask): api = DatasetApi() method = unwrap(api.get) dataset_id = "123e4567-e89b-12d3-a456-426614174000" user = make_account() tenant_id = "tenant-1" dataset = make_dataset(id=dataset_id) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] data, status = method(api, self.session, tenant_id, user, dataset_id) assert status == 200 assert data["embedding_available"] is True def test_get_attaches_permission_keys_when_rbac_enabled(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(RBAC_ENABLED=True) api = DatasetApi() method = unwrap(api.get) dataset_id = "123e4567-e89b-12d3-a456-426614174000" user = make_account() tenant_id = "tenant-1" dataset = make_dataset(id=dataset_id) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch( "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.MyPermissions.get", return_value=enterprise_rbac_service.MyPermissionsResponse( dataset=enterprise_rbac_service.ResourcePermissionSnapshot( overrides=[ enterprise_rbac_service.ResourcePermissionKeys( resource_id=dataset_id, permission_keys=["dataset.acl.readonly", "dataset.acl.edit"] ) ] ) ), ) as get_permissions, patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] data, status = method(api, self.session, tenant_id, user, dataset_id) get_permissions.assert_called_once_with(tenant_id, user.id, dataset_id=dataset_id, session=ANY) assert status == 200 assert data["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"] def test_get_uses_default_external_retrieval_model(self, app: Flask): api = DatasetApi() method = unwrap(api.get) dataset_id = "dataset-id" dataset = make_dataset(id=dataset_id, retrieval_model=None) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] data, status = method(api, self.session, "tenant", make_account(), dataset_id) assert status == 200 assert data["external_retrieval_model"] == {"top_k": 2, "score_threshold": 0.0, "score_threshold_enabled": None} def test_get_dataset_not_found(self, app: Flask): api = DatasetApi() method = unwrap(api.get) dataset_id = "missing-id" with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): method(api, self.session, "tenant", make_account(), dataset_id) def test_get_permission_denied(self, app: Flask): api = DatasetApi() method = unwrap(api.get) dataset_id = "dataset-id" dataset = make_dataset(id=dataset_id) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object( DatasetService, "check_dataset_permission", side_effect=services.errors.account.NoPermissionError("no access"), ), ): with pytest.raises(Forbidden, match="no access"): method(api, self.session, "tenant", make_account(), dataset_id) def test_get_high_quality_embedding_unavailable(self, app: Flask): api = DatasetApi() method = unwrap(api.get) dataset_id = "dataset-id" user = make_account() tenant_id = "tenant-1" dataset = make_dataset( id=dataset_id, indexing_technique="high_quality", embedding_model="text-embedding", embedding_model_provider="openai", ) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] data, _ = method(api, self.session, tenant_id, user, dataset_id) assert data["embedding_available"] is False def test_get_partial_members_permission(self, app: Flask): api = DatasetApi() method = unwrap(api.get) dataset_id = "dataset-id" dataset = make_dataset(id=dataset_id, permission="partial_members") partial_members = ["u1", "u2"] with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=partial_members), patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock, ): provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = [] data, _ = method(api, self.session, "tenant", make_account(), dataset_id) assert data["partial_member_list"] == partial_members class TestDatasetApiPatch(_UsesSQLiteSession): def test_patch_success_basic(self, app: Flask): api = DatasetApi() method = unwrap(api.patch) dataset_id = "dataset-id" payload = {"name": "updated-name", "description": "updated description"} user = make_account() tenant_id = "tenant-1" dataset = make_dataset(id=dataset_id, tenant_id=tenant_id) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetPermissionService, "check_permission", return_value=None), patch.object(DatasetService, "update_dataset", return_value=dataset), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), ): result, status = method(api, DatasetUpdatePayload(), self.session, tenant_id, user, dataset_id) assert status == 200 assert result["partial_member_list"] == [] def test_patch_dataset_not_found(self, app: Flask): api = DatasetApi() method = unwrap(api.patch) with ( app.test_request_context("/datasets/missing"), patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): method(api, DatasetUpdatePayload(), self.session, "tenant-1", make_account(), "missing") def test_patch_permission_denied(self, app: Flask): api = DatasetApi() method = unwrap(api.patch) dataset_id = "dataset-id" dataset = make_dataset(id=dataset_id) payload = {"name": "x"} with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetPermissionService, "check_permission", side_effect=Forbidden("no permission")), ): with pytest.raises(Forbidden): method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) def test_patch_partial_members_update(self, app: Flask): api = DatasetApi() method = unwrap(api.patch) dataset_id = "dataset-id" payload = {"permission": "partial_members", "partial_member_list": [{"user_id": "u1"}, {"user_id": "u2"}]} dataset = make_dataset(id=dataset_id, permission="partial_members") with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetPermissionService, "check_permission", return_value=None), patch.object(DatasetService, "update_dataset", return_value=dataset), patch.object(DatasetPermissionService, "update_partial_member_list", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["u1", "u2"]), ): result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) assert result["partial_member_list"] == ["u1", "u2"] def test_patch_clear_partial_members(self, app: Flask): api = DatasetApi() method = unwrap(api.patch) dataset_id = "dataset-id" payload = {"permission": "only_me"} dataset = make_dataset(id=dataset_id) with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(type(console_ns), "payload", payload), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetPermissionService, "check_permission", return_value=None), patch.object(DatasetService, "update_dataset", return_value=dataset), patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]), ): result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id) assert result["partial_member_list"] == [] class TestDatasetApiDelete(_UsesSQLiteSession): def test_delete_success(self, app: Flask): api = DatasetApi() method = unwrap(api.delete) dataset_id = "dataset-id" user = make_account() with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "delete_dataset", return_value=True), patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None), ): result, status = method(api, self.session, user, dataset_id) assert status == 204 assert result == "" def test_delete_forbidden_no_permission(self, app: Flask): api = DatasetApi() method = unwrap(api.delete) dataset_id = "dataset-id" user = make_account(TenantAccountRole.NORMAL) with app.test_request_context(f"/datasets/{dataset_id}"): with pytest.raises(Forbidden): method(api, self.session, user, dataset_id) def test_delete_dataset_not_found(self, app: Flask): api = DatasetApi() method = unwrap(api.delete) dataset_id = "missing-dataset" user = make_account() with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "delete_dataset", return_value=False), ): with pytest.raises(NotFound, match="Dataset not found"): method(api, self.session, user, dataset_id) def test_delete_dataset_in_use(self, app: Flask): api = DatasetApi() method = unwrap(api.delete) dataset_id = "dataset-id" user = make_account() with ( app.test_request_context(f"/datasets/{dataset_id}"), patch.object(DatasetService, "delete_dataset", side_effect=services.errors.dataset.DatasetInUseError()), ): with pytest.raises(DatasetInUseError): method(api, self.session, user, dataset_id) class TestDatasetKnowledgeFSUpgradeApi: @staticmethod def _job(dataset_id: str, **overrides): values = { "id": "upgrade-job-1", "old_dataset_id": dataset_id, "new_control_space_id": None, "status": KnowledgeFSUpgradeJobStatus.QUEUED, "stage": KnowledgeFSUpgradeStage.VALIDATING, "snapshot_at": datetime.datetime(2026, 8, 17, tzinfo=datetime.UTC), "total_documents": 1, "completed_documents": 0, "total_sources": 0, "completed_sources": 0, "last_error_code": None, "last_error_message": None, "completed_at": None, } values.update(overrides) return SimpleNamespace(**values) def test_list_upgrade_jobs_returns_only_accessible_dataset_jobs(self, app: Flask, monkeypatch: pytest.MonkeyPatch): allowed_dataset = make_dataset(id="dataset-1", tenant_id="tenant-1") denied_dataset = make_dataset(id="dataset-2", tenant_id="tenant-1") snapshots = MagicMock() snapshots.list_by_statuses.return_value = [ self._job("dataset-1", id="job-1"), self._job("dataset-2", id="job-2", status=KnowledgeFSUpgradeJobStatus.FAILED), ] api = DatasetKnowledgeFSUpgradeJobsApi() method = unwrap(api.get) def check_permission(dataset, _user, _session): if dataset.id == "dataset-2": raise services.errors.account.NoPermissionError() apply_config_overrides(monkeypatch, RBAC_ENABLED=False) with ( app.test_request_context("/datasets/knowledge-fs-upgrade-jobs"), patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), patch( "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", return_value=snapshots, ), patch.object( DatasetService, "get_datasets_by_ids", return_value=([allowed_dataset, denied_dataset], 2), ) as get_datasets, patch.object(DatasetService, "check_dataset_permission", side_effect=check_permission), ): response = method(api, MagicMock(), "tenant-1", make_account()) assert [job["id"] for job in response["data"]] == ["job-1"] snapshots.list_by_statuses.assert_called_once_with( tenant_id="tenant-1", statuses=( KnowledgeFSUpgradeJobStatus.QUEUED, KnowledgeFSUpgradeJobStatus.RUNNING, KnowledgeFSUpgradeJobStatus.FAILED, ), ) assert get_datasets.call_args.args[:2] == (["dataset-1", "dataset-2"], "tenant-1") def test_create_snapshots_and_enqueues_without_remote_work_in_request( self, app: Flask, monkeypatch: pytest.MonkeyPatch ): dataset_id = "123e4567-e89b-12d3-a456-426614174000" dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") user = make_account() snapshots = MagicMock() snapshots.create.return_value = self._job(dataset_id) session_context = MagicMock() session_context.__enter__.return_value = MagicMock() api = DatasetKnowledgeFSUpgradeApi() method = unwrap(api.post) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=True, RBAC_ENABLED=False) with ( app.test_request_context( f"/datasets/{dataset_id}/knowledge-fs-upgrades", headers={"Idempotency-Key": "upgrade-request-1"}, ), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission"), patch( "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", return_value=snapshots, ), patch("controllers.console.datasets.datasets._enqueue_upgrade_job") as enqueue, ): response, status = method(api, "tenant-1", user, dataset_id) assert status == 202 assert response["id"] == "upgrade-job-1" snapshots.create.assert_called_once_with( tenant_id="tenant-1", dataset_id=dataset_id, requested_by_account_id=user.id, idempotency_key="upgrade-request-1", ) enqueue.assert_called_once_with(snapshots, tenant_id="tenant-1", job_id="upgrade-job-1") def test_create_is_unavailable_when_knowledge_fs_is_disabled(self, app: Flask, monkeypatch: pytest.MonkeyPatch): dataset_id = "123e4567-e89b-12d3-a456-426614174000" api = DatasetKnowledgeFSUpgradeApi() method = unwrap(api.post) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=False) with app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades"), pytest.raises(NotFound): method(api, "tenant-1", make_account(), dataset_id) def test_create_rejects_a_missing_dataset(self, app: Flask, monkeypatch: pytest.MonkeyPatch): dataset_id = "123e4567-e89b-12d3-a456-426614174000" session_context = MagicMock() session_context.__enter__.return_value = MagicMock() api = DatasetKnowledgeFSUpgradeApi() method = unwrap(api.post) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=True) with ( app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades"), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch.object(DatasetService, "get_dataset_for_tenant", return_value=None), pytest.raises(NotFound, match="Dataset not found"), ): method(api, "tenant-1", make_account(), dataset_id) def test_create_requires_idempotency_key(self, app: Flask, monkeypatch: pytest.MonkeyPatch): dataset_id = "123e4567-e89b-12d3-a456-426614174000" dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") session_context = MagicMock() session_context.__enter__.return_value = MagicMock() api = DatasetKnowledgeFSUpgradeApi() method = unwrap(api.post) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=True, RBAC_ENABLED=True) with ( app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades"), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), pytest.raises(BadRequest, match="Idempotency-Key"), ): method(api, "tenant-1", make_account(), dataset_id) def test_discovery_returns_latest_job_and_allowed_action(self, app: Flask, monkeypatch: pytest.MonkeyPatch): dataset_id = "123e4567-e89b-12d3-a456-426614174000" dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") snapshots = MagicMock() snapshots.get_latest.return_value = self._job( dataset_id, status=KnowledgeFSUpgradeJobStatus.FAILED, last_error_code="RuntimeError", last_error_message="source failed", ) session_context = MagicMock() session_context.__enter__.return_value = MagicMock() api = DatasetKnowledgeFSUpgradeApi() method = unwrap(api.get) apply_config_overrides(monkeypatch, KNOWLEDGE_FS_ENABLED=True, RBAC_ENABLED=True) with ( app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades"), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch( "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", return_value=snapshots, ), ): response = method(api, "tenant-1", make_account(), dataset_id) assert response["can_upgrade"] is True assert response["can_retry"] is True assert response["block_reason"] == "retry_required" assert response["job"]["last_error_message"] == "source failed" snapshots.get_latest.assert_called_once_with(tenant_id="tenant-1", dataset_id=dataset_id) def test_status_uses_legacy_dataset_permission_in_community_edition( self, app: Flask, monkeypatch: pytest.MonkeyPatch ): dataset_id = "123e4567-e89b-12d3-a456-426614174000" dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") api = DatasetKnowledgeFSUpgradeJobApi() method = unwrap(api.get) session_context = MagicMock() session_context.__enter__.return_value = MagicMock() apply_config_overrides(monkeypatch, RBAC_ENABLED=False) with ( app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades/job-1"), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object( DatasetService, "check_dataset_permission", side_effect=services.errors.account.NoPermissionError("no access"), ), patch("controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService") as snapshots, pytest.raises(Forbidden, match="no access"), ): method(api, "tenant-1", make_account(), dataset_id, "job-1") snapshots.assert_not_called() def test_status_relies_on_rbac_decorator_when_enterprise_rbac_is_enabled( self, app: Flask, monkeypatch: pytest.MonkeyPatch ): dataset_id = "123e4567-e89b-12d3-a456-426614174000" dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") snapshots = MagicMock() snapshots.get.return_value = self._job(dataset_id) session_context = MagicMock() session_context.__enter__.return_value = MagicMock() api = DatasetKnowledgeFSUpgradeJobApi() method = unwrap(api.get) apply_config_overrides(monkeypatch, RBAC_ENABLED=True) with ( app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades/job-1"), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission") as legacy_permission, patch( "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", return_value=snapshots, ), ): response = method(api, "tenant-1", make_account(), dataset_id, "job-1") assert response["id"] == "upgrade-job-1" legacy_permission.assert_not_called() def test_retry_rejects_a_job_from_another_dataset(self, app: Flask, monkeypatch: pytest.MonkeyPatch): dataset_id = "123e4567-e89b-12d3-a456-426614174000" dataset = make_dataset(id=dataset_id, tenant_id="tenant-1") snapshots = MagicMock() snapshots.retry.return_value = self._job("223e4567-e89b-12d3-a456-426614174000") session_context = MagicMock() session_context.__enter__.return_value = MagicMock() api = DatasetKnowledgeFSUpgradeJobApi() method = unwrap(api.post) apply_config_overrides(monkeypatch, RBAC_ENABLED=False) with ( app.test_request_context(f"/datasets/{dataset_id}/knowledge-fs-upgrades/job-1"), patch("controllers.console.datasets.datasets.session_factory.create_session", return_value=session_context), patch("controllers.console.datasets.datasets.session_factory.get_session_maker", return_value="maker"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission"), patch( "controllers.console.datasets.datasets.KnowledgeFSUpgradeSnapshotService", return_value=snapshots, ), patch("controllers.console.datasets.datasets._enqueue_upgrade_job") as enqueue, pytest.raises(NotFound, match="Upgrade job was not found"), ): method(api, "tenant-1", make_account(), dataset_id, "job-1") enqueue.assert_not_called() class TestDatasetUseCheckApi(_UsesSQLiteSession): @pytest.mark.parametrize("is_using", [True, False]) def test_get_use_check(self, app: Flask, is_using: bool): api = DatasetUseCheckApi() method = unwrap(api.get) dataset_id = "dataset-id" dataset = make_dataset(id=dataset_id) current_user = make_account() session = self.session with ( app.test_request_context(f"/datasets/{dataset_id}/use-check"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, patch.object(DatasetService, "check_dataset_permission") as check_permission, patch.object(DatasetService, "dataset_use_check", return_value=is_using) as dataset_use_check, ): result, status = method(api, session, "tenant-1", current_user, dataset_id) assert status == 200 assert result == {"is_using": is_using} get_dataset.assert_called_once_with(dataset_id, "tenant-1", session=session) check_permission.assert_called_once_with(dataset, current_user, session) dataset_use_check.assert_called_once_with(DatasetRef("tenant-1", dataset_id), session) def test_get_use_check_relies_on_rbac_in_rbac_mode(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(RBAC_ENABLED=True) api = DatasetUseCheckApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-id") session = self.session with ( app.test_request_context("/datasets/dataset-id/use-check"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission") as check_permission, patch.object(DatasetService, "dataset_use_check", return_value=False), ): _, status = method(api, session, "tenant-1", make_account(), "dataset-id") assert status == 200 check_permission.assert_not_called() @pytest.mark.parametrize( "api_cls", [DatasetUseCheckApi, DatasetIndexingStatusApi, DatasetErrorDocs, DatasetAutoDisableLogApi], ) def test_dataset_scoped_read_permission_denied(app: Flask, api_cls, sqlite_session: Session): api = api_cls() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") session = sqlite_session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object( DatasetService, "check_dataset_permission", side_effect=services.errors.account.NoPermissionError("no permission"), ), ): with pytest.raises(Forbidden, match="no permission"): method(api, session, "tenant-1", make_account(), "dataset-1") class TestDatasetQueryApi(_UsesSQLiteSession): def _query_record(self, index: int = 1) -> DatasetQuery: query = DatasetQuery( dataset_id="dataset-id", content=json.dumps([{"content_type": "text_query", "content": f"question {index}", "file_info": None}]), source="hit_testing", source_app_id=None, created_by_role=CreatorUserRole.ACCOUNT, created_by=f"account-{index}", ) query.id = f"query-{index}" query.created_at = datetime.datetime(2024, 1, index, 12, 0, 0, tzinfo=datetime.UTC) return query def test_get_queries_success(self, app: Flask): api = DatasetQueryApi() method = unwrap(api.get) dataset_id = "dataset-id" current_user = make_account() dataset = make_dataset(id=dataset_id) queries = [self._query_record(1), self._query_record(2)] with ( app.test_request_context("/datasets/queries?page=1&limit=20"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 2)), ): response, status = method(api, self.session, current_user, dataset_id) assert status == 200 assert response["total"] == 2 assert response["page"] == 1 assert response["limit"] == 20 assert response["has_more"] is False assert len(response["data"]) == 2 assert response["data"][0] == { "id": "query-1", "queries": [{"content_type": "text_query", "content": "question 1", "file_info": None}], "source": "hit_testing", "source_app_id": None, "created_by_role": "account", "created_by": "account-1", "created_at": 1704110400, } def test_get_image_query_uses_caller_session(self, app: Flask): api = DatasetQueryApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-id") query = self._query_record() query.content = json.dumps([{"content_type": "image_query", "content": "file-1"}]) upload_file = UploadFile( tenant_id="tenant-1", storage_type=StorageType.LOCAL, key="image.png", name="image.png", size=10, extension="png", mime_type="image/png", created_by_role=CreatorUserRole.ACCOUNT, created_by="account-1", created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC), used=False, ) upload_file.id = "file-1" session = self.session session.add(upload_file) session.flush() with ( app.test_request_context("/datasets/queries"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetService, "get_dataset_queries", return_value=([query], 1)), patch("models.dataset.sign_upload_file_preview_url", return_value="signed-url"), ): response, status = method(api, session, make_account(), "dataset-id") assert status == 200 assert response["data"][0]["queries"] == [ { "content_type": "image_query", "content": "file-1", "file_info": { "id": "file-1", "name": "image.png", "size": 10, "extension": "png", "mime_type": "image/png", "source_url": "signed-url", }, } ] assert session.get(UploadFile, "file-1") is upload_file def test_get_queries_dataset_not_found(self, app: Flask): api = DatasetQueryApi() method = unwrap(api.get) dataset_id = "dataset-id" current_user = make_account() with ( app.test_request_context("/datasets/queries"), patch.object(DatasetService, "get_dataset", return_value=None), ): with pytest.raises(NotFound, match="Dataset not found"): method(api, self.session, current_user, dataset_id) def test_get_queries_permission_denied(self, app: Flask): api = DatasetQueryApi() method = unwrap(api.get) dataset_id = "dataset-id" current_user = make_account() dataset = make_dataset(id=dataset_id) with ( app.test_request_context("/datasets/queries"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object( DatasetService, "check_dataset_permission", side_effect=services.errors.account.NoPermissionError("no access"), ), ): with pytest.raises(Forbidden): method(api, self.session, current_user, dataset_id) def test_get_queries_pagination_has_more(self, app: Flask): api = DatasetQueryApi() method = unwrap(api.get) dataset_id = "dataset-id" current_user = make_account() dataset = make_dataset(id=dataset_id) queries = [self._query_record(index) for index in range(1, 21)] with ( app.test_request_context("/datasets/queries?page=1&limit=20"), patch.object(DatasetService, "get_dataset", return_value=dataset), patch.object(DatasetService, "check_dataset_permission", return_value=None), patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 40)), ): response, status = method(api, self.session, current_user, dataset_id) assert status == 200 assert response["has_more"] is True assert len(response["data"]) == 20 class TestDatasetIndexingEstimateApi(_UsesSQLiteSession): def _upload_file(self, *, tenant_id: str = "tenant-1", file_id: str = "file-1") -> UploadFile: upload_file = UploadFile( tenant_id=tenant_id, storage_type=StorageType.LOCAL, key="key", name="name.txt", size=1, extension="txt", mime_type="text/plain", created_by_role=CreatorUserRole.ACCOUNT, created_by="user-1", created_at=datetime.datetime.now(tz=datetime.UTC), used=False, ) upload_file.id = file_id return upload_file def _base_payload(self): return { "info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": ["file-1"]}}, "process_rule": {"chunk_size": 100}, "indexing_technique": "high_quality", "doc_form": IndexStructureType.PARAGRAPH_INDEX, "doc_language": "English", "dataset_id": None, } def test_post_success_upload_file(self, app: Flask): api = DatasetIndexingEstimateApi() method = unwrap(api.post) payload = self._base_payload() mock_file = self._upload_file() session = self.session session.add(mock_file) session.flush() mock_response = IndexingEstimate(total_segments=100, preview=[]) with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), patch("controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", return_value=mock_response), ): response, status = method( api, IndexingEstimatePayload(**payload), session, "tenant-1", ) assert status == 200 assert response == { "tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 100, "preview": [], } def test_post_file_not_found(self, app: Flask): api = DatasetIndexingEstimateApi() method = unwrap(api.post) payload = self._base_payload() session = self.session with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), ): with pytest.raises(NotFound): method( api, IndexingEstimatePayload(**payload), session, "tenant-1", ) def test_post_llm_bad_request_error(self, app: Flask): api = DatasetIndexingEstimateApi() method = unwrap(api.post) mock_file = self._upload_file() payload = self._base_payload() session = self.session session.add(mock_file) session.flush() with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), patch( "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", side_effect=LLMBadRequestError(), ), ): with pytest.raises(ProviderNotInitializeError): method( api, IndexingEstimatePayload(**payload), session, "tenant-1", ) def test_post_provider_token_not_init(self, app: Flask): api = DatasetIndexingEstimateApi() method = unwrap(api.post) mock_file = self._upload_file() payload = self._base_payload() session = self.session session.add(mock_file) session.flush() with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), patch( "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", side_effect=ProviderTokenNotInitError("token missing"), ), ): with pytest.raises(ProviderNotInitializeError): method( api, IndexingEstimatePayload(**payload), session, "tenant-1", ) def test_post_generic_exception(self, app: Flask): api = DatasetIndexingEstimateApi() method = unwrap(api.post) mock_file = self._upload_file() payload = self._base_payload() session = self.session session.add(mock_file) session.flush() with ( app.test_request_context("/"), patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload), patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None), patch( "controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", side_effect=Exception("boom") ), ): with pytest.raises(IndexingEstimateError): method( api, IndexingEstimatePayload(**payload), session, "tenant-1", ) class TestDatasetRelatedAppListApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetRelatedAppListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") app1 = make_related_app(id="app-1", name="App 1") app2 = make_related_app(id="app-2", name="App 2") join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1") join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1") session = self.session with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), patch("controllers.console.datasets.datasets.DatasetService.check_dataset_permission", return_value=None), patch("controllers.console.datasets.datasets.DatasetService.get_related_apps", return_value=[join1, join2]), patch( "controllers.console.datasets.datasets.AppService.get_app_by_id", side_effect=[app1, app2], ) as get_app_by_id, patch.object( App, "mode_compatible_with_agent_with_session", autospec=True, side_effect=lambda app_model, *, session: str(app_model.mode), ) as compatible_mode, ): response, status = method(api, session, make_account(), "dataset-1") assert status == 200 assert response["total"] == 2 assert response["data"] == [ { "id": "app-1", "name": "App 1", "description": "desc", "mode": "chat", "icon_type": "emoji", "icon": "🤖", "icon_background": "#fff", "icon_url": None, }, { "id": "app-2", "name": "App 2", "description": "desc", "mode": "chat", "icon_type": "emoji", "icon": "🤖", "icon_background": "#fff", "icon_url": None, }, ] assert compatible_mode.call_args_list == [call(app1, session=session), call(app2, session=session)] assert get_app_by_id.call_args_list == [call("app-1", session), call("app-2", session)] def test_get_dataset_not_found(self, app: Flask): api = DatasetRelatedAppListApi() method = unwrap(api.get) with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=None), ): with pytest.raises(NotFound): method(api, self.session, make_account(), "dataset-1") def test_get_permission_denied(self, app: Flask): api = DatasetRelatedAppListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), patch( "controllers.console.datasets.datasets.DatasetService.check_dataset_permission", side_effect=services.errors.account.NoPermissionError("no permission"), ), ): with pytest.raises(Forbidden): method(api, self.session, make_account(), "dataset-1") def test_get_filters_none_apps(self, app: Flask): api = DatasetRelatedAppListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") app1 = make_related_app() join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1") join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1") session = self.session with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), patch("controllers.console.datasets.datasets.DatasetService.check_dataset_permission", return_value=None), patch("controllers.console.datasets.datasets.DatasetService.get_related_apps", return_value=[join1, join2]), patch( "controllers.console.datasets.datasets.AppService.get_app_by_id", side_effect=[app1, None], ), ): response, status = method(api, session, make_account(), "dataset-1") assert status == 200 assert response["total"] == 1 assert response["data"] == [ { "id": "app-1", "name": "App", "description": "desc", "mode": "chat", "icon_type": "emoji", "icon": "🤖", "icon_background": "#fff", "icon_url": None, } ] class TestDatasetIndexingStatusApi(_UsesSQLiteSession): def test_get_success_with_documents(self, app: Flask): api = DatasetIndexingStatusApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") current_user = make_account() document = make_document_status() session = self.session session.add(document) session.add_all([make_document_segment(position=position, completed=True) for position in range(1, 4)]) session.flush() with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, patch.object(DatasetService, "check_dataset_permission") as check_permission, ): response, status = method(api, session, "tenant-1", current_user, "dataset-1") assert status == 200 assert "data" in response assert len(response["data"]) == 1 item = response["data"][0] assert item["completed_segments"] == 3 assert item["total_segments"] == 3 get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) check_permission.assert_called_once_with(dataset, current_user, session) assert session.get(Document, "doc-1") is document def test_get_success_no_documents(self, app: Flask): api = DatasetIndexingStatusApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission"), ): response, status = method(api, session, "tenant-1", make_account(), "dataset-1") assert status == 200 assert response == {"data": []} def test_segment_counts_different_values(self, app: Flask): api = DatasetIndexingStatusApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") document = make_document_status(indexing_status=IndexingStatus.INDEXING) session = self.session session.add(document) session.add_all([make_document_segment(position=position, completed=position <= 2) for position in range(1, 6)]) session.flush() with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission"), ): response, status = method(api, session, "tenant-1", make_account(), "dataset-1") assert status == 200 item = response["data"][0] assert item["completed_segments"] == 2 assert item["total_segments"] == 5 class TestDatasetApiKeyApi(_UsesSQLiteSession): def test_get_api_keys_success(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.get) mock_key_1 = ApiToken( id="key-1", type="dataset", token="dataset-aaaa1111bbbb", last_used_at=None, created_at=None, ) mock_key_2 = ApiToken( id="key-2", type="dataset", token="dataset-cccc2222dddd", last_used_at=None, created_at=None, ) session = self.session mock_key_1.tenant_id = "tenant-1" mock_key_2.tenant_id = "tenant-1" session.add_all([mock_key_1, mock_key_2]) session.flush() pending_last_used_at = datetime.datetime(2026, 8, 11, 12, 30, 0, tzinfo=datetime.UTC) with ( app.test_request_context("/"), patch( "controllers.console.datasets.datasets.get_effective_token_last_used_at", side_effect=[pending_last_used_at, None], ) as mock_effective_last_used, ): response = method(api, session, "tenant-1") assert "data" in response assert len(response["data"]) == 2 # reveal-once: the list returns masked tokens, never the full secret assert response["data"][0]["id"] == "key-1" assert response["data"][0]["token"] == "datas...bbbb" assert response["data"][0]["last_used_at"] == int(pending_last_used_at.timestamp()) assert response["data"][1]["id"] == "key-2" assert response["data"][1]["token"] == "datas...dddd" assert response["data"][1]["last_used_at"] is None assert mock_effective_last_used.call_count == 2 def test_post_create_api_key_success(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) session = self.session with ( app.test_request_context("/"), patch.object(ApiToken, "generate_api_key", return_value="dataset-abc123") as generate_api_key, ): response, status = method(api, session, "tenant-1") assert status == 200 assert isinstance(response, dict) assert response["token"] == "dataset-abc123" assert response["type"] == "dataset" assert response["created_at"] is not None generate_api_key.assert_called_once_with("dataset-", 24, session=session) assert session.get(ApiToken, response["id"]).token == "dataset-abc123" def test_post_exceed_max_keys(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) session = self.session session.add_all( [ ApiToken(id=f"key-{index}", tenant_id="tenant-1", type="dataset", token=f"ds-{index}") for index in range(10) ] ) session.flush() with app.test_request_context("/"): with pytest.raises(BadRequest) as exc_info: method(api, session, "tenant-1") assert exc_info.value.code == 400 assert vars(exc_info.value)["data"] == { "message": "Cannot create more than 10 API keys for this resource type.", "custom": "max_keys_exceeded", } def test_get_api_keys_include_bound_dataset_ids(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.get) mock_key = MagicMock(spec=ApiToken) mock_key.id = "key-1" mock_key.type = "dataset" mock_key.token = "dataset-aaaa1111bbbb" mock_key.last_used_at = None mock_key.created_at = None session = MagicMock() session.scalars.return_value.all.return_value = [mock_key] # Binding rows carry (token id, resource type, dataset id, control space id); the # masked list surfaces legacy datasets and KnowledgeFS spaces in their own fields. session.execute.return_value.all.return_value = [ ("key-1", "dataset", "ds-1", None), ("key-1", "dataset", "ds-2", None), ("key-1", "knowledge_fs_space", None, "space-1"), ] with app.test_request_context("/"): response = method(api, session, "tenant-1") assert response["data"][0]["dataset_ids"] == ["ds-1", "ds-2"] assert response["data"][0]["knowledge_space_ids"] == ["space-1"] def test_post_create_scoped_key_persists_bindings(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) mock_token = MagicMock() mock_token.id = "new-key-id" mock_token.last_used_at = None mock_token.created_at = datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC) mock_api_token_cls = MagicMock() mock_api_token_cls.return_value = mock_token mock_api_token_cls.generate_api_key.return_value = "dataset-abc123" session = MagicMock() # Ownership validation returns exactly the requested ids: first the legacy datasets, # then the KnowledgeFS spaces (each kind is checked against its own table). session.scalars.return_value.all.side_effect = [["ds-1", "ds-2"], ["space-1"]] session.scalar.return_value = 1 with ( app.test_request_context( "/", json={"dataset_ids": ["ds-1", "ds-2", "ds-1"], "knowledge_space_ids": ["space-1", "space-1"]}, ), patch("controllers.console.datasets.datasets.ApiToken", mock_api_token_cls), ): response, status = method(api, session, "tenant-1") assert status == 200 # Duplicates are collapsed and returned in the reveal-once response. assert response["dataset_ids"] == ["ds-1", "ds-2"] assert response["knowledge_space_ids"] == ["space-1"] # One binding row is added per unique resource id (plus the token itself), typed by kind. added = [call.args[0] for call in session.add.call_args_list] bindings = [obj for obj in added if obj.__class__.__name__ == "DatasetApiTokenBinding"] assert [(b.resource_type, b.dataset_id, b.control_space_id) for b in bindings] == [ ("dataset", "ds-1", None), ("dataset", "ds-2", None), ("knowledge_fs_space", None, "space-1"), ] def test_post_rejects_dataset_ids_from_another_tenant(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) session = MagicMock() # None of the requested ids belong to this tenant. session.scalars.return_value.all.return_value = [] with app.test_request_context("/", json={"dataset_ids": ["foreign-ds"]}): with pytest.raises(BadRequest) as exc_info: method(api, session, "tenant-1") assert exc_info.value.code == 400 assert "Unknown knowledge base id(s)" in vars(exc_info.value)["data"]["message"] def test_post_rejects_knowledge_space_ids_from_another_tenant(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) session = MagicMock() # None of the requested KnowledgeFS spaces belong to this tenant (or are bindable). session.scalars.return_value.all.return_value = [] with app.test_request_context("/", json={"knowledge_space_ids": ["foreign-space"]}): with pytest.raises(BadRequest) as exc_info: method(api, session, "tenant-1") assert exc_info.value.code == 400 assert "Unknown knowledge space id(s)" in vars(exc_info.value)["data"]["message"] def test_post_rejects_non_list_knowledge_space_ids(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) session = MagicMock() with app.test_request_context("/", json={"knowledge_space_ids": "not-a-list"}): with pytest.raises(BadRequest) as exc_info: method(api, session, "tenant-1") assert exc_info.value.code == 400 assert vars(exc_info.value)["data"]["message"] == "knowledge_space_ids must be a list of strings." def test_post_rejects_non_list_dataset_ids(self, app: Flask): api = DatasetApiKeyApi() method = unwrap(api.post) session = MagicMock() with app.test_request_context("/", json={"dataset_ids": "not-a-list"}): with pytest.raises(BadRequest) as exc_info: method(api, session, "tenant-1") assert exc_info.value.code == 400 assert vars(exc_info.value)["data"]["message"] == "dataset_ids must be a list of strings." class TestDatasetApiDeleteApi(_UsesSQLiteSession): def test_delete_success(self, app: Flask): api = DatasetApiDeleteApi() method = unwrap(api.delete) session = self.session key = ApiToken(id="api-key-id", tenant_id="tenant-1", type="dataset", token="dataset-secret") session.add(key) session.flush() with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.ApiTokenCache.delete") as delete_cache, ): response, status = method(api, session, "tenant-1", "api-key-id") assert status == 204 assert response == "" delete_cache.assert_called_once() session.flush() assert session.get(ApiToken, "api-key-id") is None def test_delete_key_not_found(self, app: Flask): api = DatasetApiDeleteApi() method = unwrap(api.delete) session = self.session with app.test_request_context("/"): with pytest.raises(NotFound): method(api, session, "tenant-1", "api-key-id") class TestDatasetEnableApiApi(_UsesSQLiteSession): @pytest.mark.parametrize(("status_value", "enabled"), [("enable", True), ("disable", False)]) def test_update_api_status(self, app: Flask, status_value: str, enabled: bool): api = DatasetEnableApiApi() method = unwrap(api.post) dataset = make_dataset(id="dataset-1") current_user = make_account() session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, patch.object(DatasetService, "check_dataset_permission") as check_permission, patch.object(DatasetService, "update_dataset_api_status") as update_status, ): response, status = method(api, session, "tenant-1", current_user, "dataset-1", status_value) assert status == 200 assert response["result"] == "success" get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) check_permission.assert_called_once_with(dataset, current_user, session) update_status.assert_called_once_with(dataset, enabled, current_user, session) def test_rejects_non_editor(self, app: Flask): api = DatasetEnableApiApi() method = unwrap(api.post) dataset = make_dataset(id="dataset-1") session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset), patch.object(DatasetService, "check_dataset_permission"), patch.object(DatasetService, "update_dataset_api_status") as update_status, ): with pytest.raises(Forbidden): method( api, session, "tenant-1", make_account(TenantAccountRole.NORMAL), "dataset-1", "enable", ) update_status.assert_not_called() class TestDatasetApiBaseUrlApi: def test_get_api_base_url_from_config(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(SERVICE_API_URL="https://example.com") api = DatasetApiBaseUrlApi() method = unwrap(api.get) with ( app.test_request_context("/"), ): response = method(api) assert response["api_base_url"] == "https://example.com/v1" def test_get_api_base_url_from_request(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(SERVICE_API_URL=None) api = DatasetApiBaseUrlApi() method = unwrap(api.get) with ( app.test_request_context("http://localhost:5000/"), ): response = method(api) assert response["api_base_url"] == "http://localhost:5000/v1" def test_get_api_base_url_no_double_v1(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(SERVICE_API_URL="https://example.com/v1") api = DatasetApiBaseUrlApi() method = unwrap(api.get) with ( app.test_request_context("/"), ): response = method(api) assert response["api_base_url"] == "https://example.com/v1" class TestDatasetRetrievalSettingApi: def test_get_success(self, app: Flask, config_overrides: Callable[..., None]): config_overrides(VECTOR_STORE="qdrant") api = DatasetRetrievalSettingApi() method = unwrap(api.get) with ( app.test_request_context("/"), patch( "controllers.console.datasets.datasets._get_retrieval_methods_by_vector_type", return_value={"retrieval_method": ["semantic", "hybrid"]}, ), ): response = method(api) assert "retrieval_method" in response def test_tidb_vector_returns_semantic_only_when_fulltext_disabled(self, config_overrides: Callable[..., None]): config_overrides(TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH=False) response = _get_retrieval_methods_by_vector_type(VectorType.TIDB_VECTOR) assert response["retrieval_method"] == [RetrievalMethod.SEMANTIC_SEARCH.value] def test_tidb_vector_returns_full_methods_when_fulltext_enabled(self, config_overrides: Callable[..., None]): config_overrides(TIDB_VECTOR_ENABLE_FULLTEXT_SEARCH=True) response = _get_retrieval_methods_by_vector_type(VectorType.TIDB_VECTOR) assert response["retrieval_method"] == [ RetrievalMethod.SEMANTIC_SEARCH.value, RetrievalMethod.FULL_TEXT_SEARCH.value, RetrievalMethod.HYBRID_SEARCH.value, ] class TestDatasetRetrievalSettingMockApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetRetrievalSettingMockApi() method = unwrap(api.get) with ( app.test_request_context("/"), patch( "controllers.console.datasets.datasets._get_retrieval_methods_by_vector_type", return_value={"retrieval_method": ["semantic"]}, ), ): response = method(api, "milvus") assert response["retrieval_method"] == ["semantic"] class TestDatasetErrorDocs(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetErrorDocs() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") error_doc = make_document_status(id="error-doc", indexing_status=IndexingStatus.ERROR, error="failed") current_user = make_account() session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, patch.object(DatasetService, "check_dataset_permission") as check_permission, patch( "controllers.console.datasets.datasets.DocumentService.get_error_documents_by_dataset_ref", return_value=[error_doc], ) as get_error_documents, ): response, status = method(api, session, "tenant-1", current_user, "dataset-1") assert status == 200 assert response["total"] == 1 get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) check_permission.assert_called_once_with(dataset, current_user, session) get_error_documents.assert_called_once_with(DatasetRef("tenant-1", "dataset-1"), session) def test_get_dataset_not_found(self, app: Flask): api = DatasetErrorDocs() method = unwrap(api.get) session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, ): with pytest.raises(NotFound): method(api, session, "tenant-1", make_account(), "dataset-1") get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) class TestDatasetPermissionUserListApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetPermissionUserListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") users = ["u1", "u2"] with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), patch("controllers.console.datasets.datasets.DatasetService.check_dataset_permission", return_value=None), patch( "controllers.console.datasets.datasets.DatasetPermissionService.get_dataset_partial_member_list", return_value=users, ), ): response, status = method(api, self.session, make_account(), "dataset-1") assert status == 200 assert response["data"] == users def test_get_permission_denied(self, app: Flask): api = DatasetPermissionUserListApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") with ( app.test_request_context("/"), patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset), patch( "controllers.console.datasets.datasets.DatasetService.check_dataset_permission", side_effect=services.errors.account.NoPermissionError("no permission"), ), ): with pytest.raises(Forbidden): method(api, self.session, make_account(), "dataset-1") class TestDatasetAutoDisableLogApi(_UsesSQLiteSession): def test_get_success(self, app: Flask): api = DatasetAutoDisableLogApi() method = unwrap(api.get) dataset = make_dataset(id="dataset-1") logs = {"document_ids": ["doc-1"], "count": 1} current_user = make_account() session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset, patch.object(DatasetService, "check_dataset_permission") as check_permission, patch.object(DatasetService, "get_dataset_auto_disable_logs", return_value=logs) as get_logs, ): response, status = method(api, session, "tenant-1", current_user, "dataset-1") assert status == 200 assert response == logs get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session) check_permission.assert_called_once_with(dataset, current_user, session) get_logs.assert_called_once_with(DatasetRef("tenant-1", "dataset-1"), session) def test_get_dataset_not_found(self, app: Flask): api = DatasetAutoDisableLogApi() method = unwrap(api.get) session = self.session with ( app.test_request_context("/"), patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset, ): with pytest.raises(NotFound): method(api, session, "tenant-1", make_account(), "dataset-1") get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session)