"""Integration tests for SQL-oriented DatasetService scenarios.

This suite migrates SQL-backed behaviors from the old unit suite to real
container-backed integration tests. The tests exercise real ORM persistence and
only patch non-DB collaborators when needed.
"""

from unittest.mock import Mock, patch
from uuid import uuid4

import pytest
from sqlalchemy.orm import Session

from core.rag.retrieval.retrieval_methods import RetrievalMethod
from dify_graph.model_runtime.entities.model_entities import ModelType
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, DatasetPermissionEnum, Document, ExternalKnowledgeBindings, Pipeline
from services.dataset_service import DatasetService
from services.entities.knowledge_entities.knowledge_entities import RerankingModel, RetrievalModel
from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, RagPipelineDatasetCreateEntity
from services.errors.dataset import DatasetNameDuplicateError


class DatasetServiceIntegrationDataFactory:
    """Factory for creating real database entities used by integration tests."""

    @staticmethod
    def create_account_with_tenant(
        db_session_with_containers: Session, role: TenantAccountRole = TenantAccountRole.OWNER
    ) -> tuple[Account, Tenant]:
        """Create an account and tenant, then bind the account as current tenant member."""
        account = Account(
            email=f"{uuid4()}@example.com",
            name=f"user-{uuid4()}",
            interface_language="en-US",
            status="active",
        )
        tenant = Tenant(name=f"tenant-{uuid4()}", status="normal")
        db_session_with_containers.add_all([account, tenant])
        db_session_with_containers.flush()

        join = TenantAccountJoin(
            tenant_id=tenant.id,
            account_id=account.id,
            role=role,
            current=True,
        )
        db_session_with_containers.add(join)
        db_session_with_containers.flush()

        # Keep tenant context on the in-memory user without opening a separate session.
        account.role = role
        account._current_tenant = tenant
        return account, tenant

    @staticmethod
    def create_dataset(
        db_session_with_containers: Session,
        tenant_id: str,
        created_by: str,
        name: str = "Test Dataset",
        description: str | None = "Test description",
        provider: str = "vendor",
        indexing_technique: str | None = "high_quality",
        permission: str = DatasetPermissionEnum.ONLY_ME,
        retrieval_model: dict | None = None,
        embedding_model_provider: str | None = None,
        embedding_model: str | None = None,
        collection_binding_id: str | None = None,
        chunk_structure: str | None = None,
    ) -> Dataset:
        """Create a dataset record with configurable SQL fields."""
        dataset = Dataset(
            tenant_id=tenant_id,
            name=name,
            description=description,
            data_source_type="upload_file",
            indexing_technique=indexing_technique,
            created_by=created_by,
            provider=provider,
            permission=permission,
            retrieval_model=retrieval_model,
            embedding_model_provider=embedding_model_provider,
            embedding_model=embedding_model,
            collection_binding_id=collection_binding_id,
            chunk_structure=chunk_structure,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()
        return dataset

    @staticmethod
    def create_document(
        db_session_with_containers: Session, dataset: Dataset, created_by: str, name: str = "doc.txt"
    ) -> Document:
        """Create a document row belonging to the given dataset."""
        document = Document(
            tenant_id=dataset.tenant_id,
            dataset_id=dataset.id,
            position=1,
            data_source_type="upload_file",
            data_source_info='{"upload_file_id": "upload-file-id"}',
            batch=str(uuid4()),
            name=name,
            created_from="web",
            created_by=created_by,
            indexing_status="completed",
            doc_form="text_model",
        )
        db_session_with_containers.add(document)
        db_session_with_containers.flush()
        return document

    @staticmethod
    def create_embedding_model(provider: str = "openai", model_name: str = "text-embedding-ada-002") -> Mock:
        """Create a fake embedding model object for external provider boundary patching."""
        embedding_model = Mock()
        embedding_model.provider = provider
        embedding_model.model_name = model_name
        return embedding_model


class TestDatasetServiceCreateDataset:
    """Integration coverage for DatasetService.create_empty_dataset."""

    def test_create_internal_dataset_basic_success(self, db_session_with_containers: Session):
        """Create a basic internal dataset with minimal configuration."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)

        # Act
        result = DatasetService.create_empty_dataset(
            tenant_id=tenant.id,
            name="Basic Internal Dataset",
            description="Test description",
            indexing_technique=None,
            account=account,
        )

        # Assert
        created_dataset = db_session_with_containers.get(Dataset, result.id)
        assert created_dataset is not None
        assert created_dataset.provider == "vendor"
        assert created_dataset.permission == DatasetPermissionEnum.ONLY_ME
        assert created_dataset.embedding_model_provider is None
        assert created_dataset.embedding_model is None

    def test_create_internal_dataset_with_economy_indexing(self, db_session_with_containers: Session):
        """Create an internal dataset with economy indexing and no embedding model."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)

        # Act
        result = DatasetService.create_empty_dataset(
            tenant_id=tenant.id,
            name="Economy Dataset",
            description=None,
            indexing_technique="economy",
            account=account,
        )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.indexing_technique == "economy"
        assert result.embedding_model_provider is None
        assert result.embedding_model is None

    def test_create_internal_dataset_with_high_quality_indexing(self, db_session_with_containers: Session):
        """Create a high-quality dataset and persist embedding model settings."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        embedding_model = DatasetServiceIntegrationDataFactory.create_embedding_model()

        # Act
        with patch("services.dataset_service.ModelManager") as mock_model_manager:
            mock_model_manager.return_value.get_default_model_instance.return_value = embedding_model

            result = DatasetService.create_empty_dataset(
                tenant_id=tenant.id,
                name="High Quality Dataset",
                description=None,
                indexing_technique="high_quality",
                account=account,
            )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.indexing_technique == "high_quality"
        assert result.embedding_model_provider == embedding_model.provider
        assert result.embedding_model == embedding_model.model_name
        mock_model_manager.return_value.get_default_model_instance.assert_called_once_with(
            tenant_id=tenant.id,
            model_type=ModelType.TEXT_EMBEDDING,
        )

    def test_create_dataset_duplicate_name_error(self, db_session_with_containers: Session):
        """Raise duplicate-name error when the same tenant already has the name."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            name="Duplicate Dataset",
            indexing_technique=None,
        )

        # Act / Assert
        with pytest.raises(DatasetNameDuplicateError):
            DatasetService.create_empty_dataset(
                tenant_id=tenant.id,
                name="Duplicate Dataset",
                description=None,
                indexing_technique=None,
                account=account,
            )

    def test_create_external_dataset_success(self, db_session_with_containers: Session):
        """Create an external dataset and persist external knowledge binding."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        external_knowledge_api_id = str(uuid4())
        external_knowledge_id = "knowledge-123"

        # Act
        with patch("services.dataset_service.ExternalDatasetService.get_external_knowledge_api") as mock_get_api:
            mock_get_api.return_value = Mock(id=external_knowledge_api_id)
            result = DatasetService.create_empty_dataset(
                tenant_id=tenant.id,
                name="External Dataset",
                description=None,
                indexing_technique=None,
                account=account,
                provider="external",
                external_knowledge_api_id=external_knowledge_api_id,
                external_knowledge_id=external_knowledge_id,
            )

        # Assert
        binding = db_session_with_containers.query(ExternalKnowledgeBindings).filter_by(dataset_id=result.id).first()
        assert result.provider == "external"
        assert binding is not None
        assert binding.external_knowledge_id == external_knowledge_id
        assert binding.external_knowledge_api_id == external_knowledge_api_id

    def test_create_dataset_with_retrieval_model_and_reranking(self, db_session_with_containers: Session):
        """Create a high-quality dataset with retrieval/reranking settings."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        embedding_model = DatasetServiceIntegrationDataFactory.create_embedding_model()
        retrieval_model = RetrievalModel(
            search_method=RetrievalMethod.SEMANTIC_SEARCH,
            reranking_enable=True,
            reranking_model=RerankingModel(
                reranking_provider_name="cohere",
                reranking_model_name="rerank-english-v2.0",
            ),
            top_k=3,
            score_threshold_enabled=True,
            score_threshold=0.6,
        )

        # Act
        with (
            patch("services.dataset_service.ModelManager") as mock_model_manager,
            patch("services.dataset_service.DatasetService.check_reranking_model_setting") as mock_check_reranking,
        ):
            mock_model_manager.return_value.get_default_model_instance.return_value = embedding_model

            result = DatasetService.create_empty_dataset(
                tenant_id=tenant.id,
                name="Dataset With Reranking",
                description=None,
                indexing_technique="high_quality",
                account=account,
                retrieval_model=retrieval_model,
            )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.retrieval_model == retrieval_model.model_dump()
        mock_check_reranking.assert_called_once_with(tenant.id, "cohere", "rerank-english-v2.0")

    def test_create_internal_dataset_with_high_quality_indexing_custom_embedding(
        self, db_session_with_containers: Session
    ):
        """Create high-quality dataset with explicitly configured embedding model."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        embedding_provider = "openai"
        embedding_model_name = "text-embedding-3-small"
        embedding_model = DatasetServiceIntegrationDataFactory.create_embedding_model(
            provider=embedding_provider, model_name=embedding_model_name
        )

        # Act
        with (
            patch("services.dataset_service.ModelManager") as mock_model_manager,
            patch("services.dataset_service.DatasetService.check_embedding_model_setting") as mock_check_embedding,
        ):
            mock_model_manager.return_value.get_model_instance.return_value = embedding_model

            result = DatasetService.create_empty_dataset(
                tenant_id=tenant.id,
                name="Custom Embedding Dataset",
                description=None,
                indexing_technique="high_quality",
                account=account,
                embedding_model_provider=embedding_provider,
                embedding_model_name=embedding_model_name,
            )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.indexing_technique == "high_quality"
        assert result.embedding_model_provider == embedding_provider
        assert result.embedding_model == embedding_model_name
        mock_check_embedding.assert_called_once_with(tenant.id, embedding_provider, embedding_model_name)
        mock_model_manager.return_value.get_model_instance.assert_called_once_with(
            tenant_id=tenant.id,
            provider=embedding_provider,
            model_type=ModelType.TEXT_EMBEDDING,
            model=embedding_model_name,
        )

    def test_create_internal_dataset_with_retrieval_model(self, db_session_with_containers: Session):
        """Persist retrieval model settings when creating an internal dataset."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        retrieval_model = RetrievalModel(
            search_method=RetrievalMethod.SEMANTIC_SEARCH,
            reranking_enable=False,
            top_k=2,
            score_threshold_enabled=True,
            score_threshold=0.0,
        )

        # Act
        result = DatasetService.create_empty_dataset(
            tenant_id=tenant.id,
            name="Retrieval Model Dataset",
            description=None,
            indexing_technique=None,
            account=account,
            retrieval_model=retrieval_model,
        )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.retrieval_model == retrieval_model.model_dump()

    def test_create_internal_dataset_with_custom_permission(self, db_session_with_containers: Session):
        """Persist canonical custom permission when creating an internal dataset."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)

        # Act
        result = DatasetService.create_empty_dataset(
            tenant_id=tenant.id,
            name="Custom Permission Dataset",
            description=None,
            indexing_technique=None,
            account=account,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.permission == DatasetPermissionEnum.ALL_TEAM

    def test_create_external_dataset_missing_api_id_error(self, db_session_with_containers: Session):
        """Raise error when external API template does not exist."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        external_knowledge_api_id = str(uuid4())

        # Act / Assert
        with patch("services.dataset_service.ExternalDatasetService.get_external_knowledge_api") as mock_get_api:
            mock_get_api.return_value = None
            with pytest.raises(ValueError, match=r"External API template not found\.?"):
                DatasetService.create_empty_dataset(
                    tenant_id=tenant.id,
                    name="External Missing API Dataset",
                    description=None,
                    indexing_technique=None,
                    account=account,
                    provider="external",
                    external_knowledge_api_id=external_knowledge_api_id,
                    external_knowledge_id="knowledge-123",
                )

    def test_create_external_dataset_missing_knowledge_id_error(self, db_session_with_containers: Session):
        """Raise error when external knowledge id is missing for external dataset creation."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        external_knowledge_api_id = str(uuid4())

        # Act / Assert
        with patch("services.dataset_service.ExternalDatasetService.get_external_knowledge_api") as mock_get_api:
            mock_get_api.return_value = Mock(id=external_knowledge_api_id)
            with pytest.raises(ValueError, match="external_knowledge_id is required"):
                DatasetService.create_empty_dataset(
                    tenant_id=tenant.id,
                    name="External Missing Knowledge Dataset",
                    description=None,
                    indexing_technique=None,
                    account=account,
                    provider="external",
                    external_knowledge_api_id=external_knowledge_api_id,
                    external_knowledge_id=None,
                )


class TestDatasetServiceCreateRagPipelineDataset:
    """Integration coverage for DatasetService.create_empty_rag_pipeline_dataset."""

    def test_create_rag_pipeline_dataset_with_name_success(self, db_session_with_containers: Session):
        """Create rag-pipeline dataset and pipeline rows when a name is provided."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        icon_info = IconInfo(icon="📙", icon_background="#FFF4ED", icon_type="emoji")
        entity = RagPipelineDatasetCreateEntity(
            name="RAG Pipeline Dataset",
            description="RAG Pipeline Description",
            icon_info=icon_info,
            permission=DatasetPermissionEnum.ONLY_ME,
        )

        # Act
        with patch("services.dataset_service.current_user", account):
            result = DatasetService.create_empty_rag_pipeline_dataset(
                tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity
            )

        # Assert
        created_dataset = db_session_with_containers.get(Dataset, result.id)
        created_pipeline = db_session_with_containers.get(Pipeline, result.pipeline_id)
        assert created_dataset is not None
        assert created_dataset.name == entity.name
        assert created_dataset.runtime_mode == "rag_pipeline"
        assert created_dataset.created_by == account.id
        assert created_dataset.permission == DatasetPermissionEnum.ONLY_ME
        assert created_pipeline is not None
        assert created_pipeline.name == entity.name
        assert created_pipeline.created_by == account.id

    def test_create_rag_pipeline_dataset_with_auto_generated_name(self, db_session_with_containers: Session):
        """Create rag-pipeline dataset with generated incremental name when input name is empty."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        generated_name = "Untitled 1"
        icon_info = IconInfo(icon="📙", icon_background="#FFF4ED", icon_type="emoji")
        entity = RagPipelineDatasetCreateEntity(
            name="",
            description="",
            icon_info=icon_info,
            permission=DatasetPermissionEnum.ONLY_ME,
        )

        # Act
        with (
            patch("services.dataset_service.current_user", account),
            patch("services.dataset_service.generate_incremental_name") as mock_generate_name,
        ):
            mock_generate_name.return_value = generated_name
            result = DatasetService.create_empty_rag_pipeline_dataset(
                tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity
            )

        # Assert
        db_session_with_containers.refresh(result)
        created_pipeline = db_session_with_containers.get(Pipeline, result.pipeline_id)
        assert result.name == generated_name
        assert created_pipeline is not None
        assert created_pipeline.name == generated_name
        mock_generate_name.assert_called_once()

    def test_create_rag_pipeline_dataset_duplicate_name_error(self, db_session_with_containers: Session):
        """Raise duplicate-name error when rag-pipeline dataset name already exists."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        duplicate_name = "Duplicate RAG Dataset"
        DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            name=duplicate_name,
            indexing_technique=None,
        )
        db_session_with_containers.commit()
        icon_info = IconInfo(icon="📙", icon_background="#FFF4ED", icon_type="emoji")
        entity = RagPipelineDatasetCreateEntity(
            name=duplicate_name,
            description="",
            icon_info=icon_info,
            permission=DatasetPermissionEnum.ONLY_ME,
        )

        # Act / Assert
        with (
            patch("services.dataset_service.current_user", account),
            pytest.raises(DatasetNameDuplicateError, match=f"Dataset with name {duplicate_name} already exists"),
        ):
            DatasetService.create_empty_rag_pipeline_dataset(
                tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity
            )

    def test_create_rag_pipeline_dataset_with_custom_permission(self, db_session_with_containers: Session):
        """Persist canonical custom permission for rag-pipeline dataset creation."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        icon_info = IconInfo(icon="📙", icon_background="#FFF4ED", icon_type="emoji")
        entity = RagPipelineDatasetCreateEntity(
            name="Custom Permission RAG Dataset",
            description="",
            icon_info=icon_info,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )

        # Act
        with patch("services.dataset_service.current_user", account):
            result = DatasetService.create_empty_rag_pipeline_dataset(
                tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity
            )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.permission == DatasetPermissionEnum.ALL_TEAM

    def test_create_rag_pipeline_dataset_with_icon_info(self, db_session_with_containers: Session):
        """Persist icon metadata when creating rag-pipeline dataset."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        icon_info = IconInfo(
            icon="📚",
            icon_background="#E8F5E9",
            icon_type="emoji",
            icon_url="https://example.com/icon.png",
        )
        entity = RagPipelineDatasetCreateEntity(
            name="Icon Info RAG Dataset",
            description="",
            icon_info=icon_info,
            permission=DatasetPermissionEnum.ONLY_ME,
        )

        # Act
        with patch("services.dataset_service.current_user", account):
            result = DatasetService.create_empty_rag_pipeline_dataset(
                tenant_id=tenant.id, rag_pipeline_dataset_create_entity=entity
            )

        # Assert
        db_session_with_containers.refresh(result)
        assert result.icon_info == icon_info.model_dump()


class TestDatasetServiceUpdateAndDeleteDataset:
    """Integration coverage for SQL-backed update and delete behavior."""

    def test_update_dataset_duplicate_name_error(self, db_session_with_containers: Session):
        """Reject update when target name already exists within the same tenant."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        source_dataset = DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            name="Source Dataset",
        )
        DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            name="Existing Dataset",
        )

        # Act / Assert
        with pytest.raises(ValueError, match="Dataset name already exists"):
            DatasetService.update_dataset(source_dataset.id, {"name": "Existing Dataset"}, account)

    def test_delete_dataset_with_documents_success(self, db_session_with_containers: Session):
        """Delete a dataset that already has documents."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            indexing_technique="high_quality",
            chunk_structure="text_model",
        )
        DatasetServiceIntegrationDataFactory.create_document(
            db_session_with_containers, dataset=dataset, created_by=account.id
        )

        # Act
        with patch("services.dataset_service.dataset_was_deleted") as dataset_deleted_signal:
            result = DatasetService.delete_dataset(dataset.id, account)

        # Assert
        assert result is True
        assert db_session_with_containers.get(Dataset, dataset.id) is None
        dataset_deleted_signal.send.assert_called_once_with(dataset)

    def test_delete_empty_dataset_success(self, db_session_with_containers: Session):
        """Delete a dataset that has no documents and no indexing technique."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            indexing_technique=None,
            chunk_structure=None,
        )

        # Act
        with patch("services.dataset_service.dataset_was_deleted") as dataset_deleted_signal:
            result = DatasetService.delete_dataset(dataset.id, account)

        # Assert
        assert result is True
        assert db_session_with_containers.get(Dataset, dataset.id) is None
        dataset_deleted_signal.send.assert_called_once_with(dataset)

    def test_delete_dataset_with_partial_none_values(self, db_session_with_containers: Session):
        """Delete dataset when indexing_technique is None but doc_form path still exists."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            indexing_technique=None,
            chunk_structure="text_model",
        )

        # Act
        with patch("services.dataset_service.dataset_was_deleted") as dataset_deleted_signal:
            result = DatasetService.delete_dataset(dataset.id, account)

        # Assert
        assert result is True
        assert db_session_with_containers.get(Dataset, dataset.id) is None
        dataset_deleted_signal.send.assert_called_once_with(dataset)


class TestDatasetServiceRetrievalConfiguration:
    """Integration coverage for retrieval configuration persistence."""

    def test_get_dataset_retrieval_configuration(self, db_session_with_containers: Session):
        """Return retrieval configuration that is persisted in SQL."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        retrieval_model = {
            "search_method": "semantic_search",
            "top_k": 5,
            "score_threshold": 0.5,
            "reranking_enable": True,
        }
        dataset = DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            retrieval_model=retrieval_model,
        )

        # Act
        result = DatasetService.get_dataset(dataset.id)

        # Assert
        assert result is not None
        assert result.retrieval_model == retrieval_model
        assert result.retrieval_model["search_method"] == "semantic_search"
        assert result.retrieval_model["top_k"] == 5

    def test_update_dataset_retrieval_configuration(self, db_session_with_containers: Session):
        """Persist retrieval configuration updates through DatasetService.update_dataset."""
        # Arrange
        account, tenant = DatasetServiceIntegrationDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetServiceIntegrationDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            indexing_technique="high_quality",
            retrieval_model={"search_method": "semantic_search", "top_k": 2, "score_threshold": 0.0},
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id=str(uuid4()),
        )
        update_data = {
            "indexing_technique": "high_quality",
            "retrieval_model": {
                "search_method": "full_text_search",
                "top_k": 10,
                "score_threshold": 0.7,
            },
        }

        # Act
        result = DatasetService.update_dataset(dataset.id, update_data, account)

        # Assert
        db_session_with_containers.refresh(dataset)
        assert result.id == dataset.id
        assert dataset.retrieval_model == update_data["retrieval_model"]
