dify/api/tests/unit_tests/controllers/console/datasets/test_datasets.py

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_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="ds-abc",
            last_used_at=None,
            created_at=None,
        )
        mock_key_2 = ApiToken(
            id="key-2",
            type="dataset",
            token="ds-def",
            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
        assert response["data"][0]["id"] == "key-1"
        assert response["data"][0]["token"] == "ds-abc"
        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"] == "ds-def"
        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",
        }


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)