dify-1.8.1/api/tests/unit_tests/services/test_dataset_service_update_dataset.py

import datetime
from typing import Any, Optional

# Mock redis_client before importing dataset_service
from unittest.mock import Mock, patch

import pytest

from core.model_runtime.entities.model_entities import ModelType
from models.dataset import Dataset, ExternalKnowledgeBindings
from services.dataset_service import DatasetService
from services.errors.account import NoPermissionError


class DatasetUpdateTestDataFactory:
    """Factory class for creating test data and mock objects for dataset update tests."""

    @staticmethod
    def create_dataset_mock(
        dataset_id: str = "dataset-123",
        provider: str = "vendor",
        name: str = "old_name",
        description: str = "old_description",
        indexing_technique: str = "high_quality",
        retrieval_model: str = "old_model",
        embedding_model_provider: Optional[str] = None,
        embedding_model: Optional[str] = None,
        collection_binding_id: Optional[str] = None,
        **kwargs,
    ) -> Mock:
        """Create a mock dataset with specified attributes."""
        dataset = Mock(spec=Dataset)
        dataset.id = dataset_id
        dataset.provider = provider
        dataset.name = name
        dataset.description = description
        dataset.indexing_technique = indexing_technique
        dataset.retrieval_model = retrieval_model
        dataset.embedding_model_provider = embedding_model_provider
        dataset.embedding_model = embedding_model
        dataset.collection_binding_id = collection_binding_id
        for key, value in kwargs.items():
            setattr(dataset, key, value)
        return dataset

    @staticmethod
    def create_user_mock(user_id: str = "user-789") -> Mock:
        """Create a mock user."""
        user = Mock()
        user.id = user_id
        return user

    @staticmethod
    def create_external_binding_mock(
        external_knowledge_id: str = "old_knowledge_id", external_knowledge_api_id: str = "old_api_id"
    ) -> Mock:
        """Create a mock external knowledge binding."""
        binding = Mock(spec=ExternalKnowledgeBindings)
        binding.external_knowledge_id = external_knowledge_id
        binding.external_knowledge_api_id = external_knowledge_api_id
        return binding

    @staticmethod
    def create_embedding_model_mock(model: str = "text-embedding-ada-002", provider: str = "openai") -> Mock:
        """Create a mock embedding model."""
        embedding_model = Mock()
        embedding_model.model = model
        embedding_model.provider = provider
        return embedding_model

    @staticmethod
    def create_collection_binding_mock(binding_id: str = "binding-456") -> Mock:
        """Create a mock collection binding."""
        binding = Mock()
        binding.id = binding_id
        return binding

    @staticmethod
    def create_current_user_mock(tenant_id: str = "tenant-123") -> Mock:
        """Create a mock current user."""
        current_user = Mock()
        current_user.current_tenant_id = tenant_id
        return current_user


class TestDatasetServiceUpdateDataset:
    """
    Comprehensive unit tests for DatasetService.update_dataset method.

    This test suite covers all supported scenarios including:
    - External dataset updates
    - Internal dataset updates with different indexing techniques
    - Embedding model updates
    - Permission checks
    - Error conditions and edge cases
    """

    @pytest.fixture
    def mock_dataset_service_dependencies(self):
        """Common mock setup for dataset service dependencies."""
        with (
            patch("services.dataset_service.DatasetService.get_dataset") as mock_get_dataset,
            patch("services.dataset_service.DatasetService.check_dataset_permission") as mock_check_perm,
            patch("extensions.ext_database.db.session") as mock_db,
            patch("services.dataset_service.naive_utc_now") as mock_naive_utc_now,
        ):
            current_time = datetime.datetime(2023, 1, 1, 12, 0, 0)
            mock_naive_utc_now.return_value = current_time

            yield {
                "get_dataset": mock_get_dataset,
                "check_permission": mock_check_perm,
                "db_session": mock_db,
                "naive_utc_now": mock_naive_utc_now,
                "current_time": current_time,
            }

    @pytest.fixture
    def mock_external_provider_dependencies(self):
        """Mock setup for external provider tests."""
        with patch("services.dataset_service.Session") as mock_session:
            from extensions.ext_database import db

            with patch.object(db.__class__, "engine", new_callable=Mock):
                session_mock = Mock()
                mock_session.return_value.__enter__.return_value = session_mock
                yield session_mock

    @pytest.fixture
    def mock_internal_provider_dependencies(self):
        """Mock setup for internal provider tests."""
        with (
            patch("services.dataset_service.ModelManager") as mock_model_manager,
            patch(
                "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding"
            ) as mock_get_binding,
            patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task,
            patch("services.dataset_service.current_user") as mock_current_user,
        ):
            mock_current_user.current_tenant_id = "tenant-123"
            yield {
                "model_manager": mock_model_manager,
                "get_binding": mock_get_binding,
                "task": mock_task,
                "current_user": mock_current_user,
            }

    def _assert_database_update_called(self, mock_db, dataset_id: str, expected_updates: dict[str, Any]):
        """Helper method to verify database update calls."""
        mock_db.query.return_value.filter_by.return_value.update.assert_called_once_with(expected_updates)
        mock_db.commit.assert_called_once()

    def _assert_external_dataset_update(self, mock_dataset, mock_binding, update_data: dict[str, Any]):
        """Helper method to verify external dataset updates."""
        assert mock_dataset.name == update_data.get("name", mock_dataset.name)
        assert mock_dataset.description == update_data.get("description", mock_dataset.description)
        assert mock_dataset.retrieval_model == update_data.get("external_retrieval_model", mock_dataset.retrieval_model)

        if "external_knowledge_id" in update_data:
            assert mock_binding.external_knowledge_id == update_data["external_knowledge_id"]
        if "external_knowledge_api_id" in update_data:
            assert mock_binding.external_knowledge_api_id == update_data["external_knowledge_api_id"]

    # ==================== External Dataset Tests ====================

    def test_update_external_dataset_success(
        self, mock_dataset_service_dependencies, mock_external_provider_dependencies
    ):
        """Test successful update of external dataset."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(
            provider="external", name="old_name", description="old_description", retrieval_model="old_model"
        )
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()
        binding = DatasetUpdateTestDataFactory.create_external_binding_mock()

        # Mock external knowledge binding query
        mock_external_provider_dependencies.query.return_value.filter_by.return_value.first.return_value = binding

        update_data = {
            "name": "new_name",
            "description": "new_description",
            "external_retrieval_model": "new_model",
            "permission": "only_me",
            "external_knowledge_id": "new_knowledge_id",
            "external_knowledge_api_id": "new_api_id",
        }

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify permission check was called
        mock_dataset_service_dependencies["check_permission"].assert_called_once_with(dataset, user)

        # Verify dataset and binding updates
        self._assert_external_dataset_update(dataset, binding, update_data)

        # Verify database operations
        mock_db = mock_dataset_service_dependencies["db_session"]
        mock_db.add.assert_any_call(dataset)
        mock_db.add.assert_any_call(binding)
        mock_db.commit.assert_called_once()

        # Verify return value
        assert result == dataset

    def test_update_external_dataset_missing_knowledge_id_error(self, mock_dataset_service_dependencies):
        """Test error when external knowledge id is missing."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="external")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()
        update_data = {"name": "new_name", "external_knowledge_api_id": "api_id"}

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset("dataset-123", update_data, user)

        assert "External knowledge id is required" in str(context.value)

    def test_update_external_dataset_missing_api_id_error(self, mock_dataset_service_dependencies):
        """Test error when external knowledge api id is missing."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="external")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()
        update_data = {"name": "new_name", "external_knowledge_id": "knowledge_id"}

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset("dataset-123", update_data, user)

        assert "External knowledge api id is required" in str(context.value)

    def test_update_external_dataset_binding_not_found_error(
        self, mock_dataset_service_dependencies, mock_external_provider_dependencies
    ):
        """Test error when external knowledge binding is not found."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="external")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        # Mock external knowledge binding query returning None
        mock_external_provider_dependencies.query.return_value.filter_by.return_value.first.return_value = None

        update_data = {
            "name": "new_name",
            "external_knowledge_id": "knowledge_id",
            "external_knowledge_api_id": "api_id",
        }

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset("dataset-123", update_data, user)

        assert "External knowledge binding not found" in str(context.value)

    # ==================== Internal Dataset Basic Tests ====================

    def test_update_internal_dataset_basic_success(self, mock_dataset_service_dependencies):
        """Test successful update of internal dataset with basic fields."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(
            provider="vendor",
            indexing_technique="high_quality",
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id="binding-123",
        )
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        update_data = {
            "name": "new_name",
            "description": "new_description",
            "indexing_technique": "high_quality",
            "retrieval_model": "new_model",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
        }

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify permission check was called
        mock_dataset_service_dependencies["check_permission"].assert_called_once_with(dataset, user)

        # Verify database update was called with correct filtered data
        expected_filtered_data = {
            "name": "new_name",
            "description": "new_description",
            "indexing_technique": "high_quality",
            "retrieval_model": "new_model",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        self._assert_database_update_called(
            mock_dataset_service_dependencies["db_session"], "dataset-123", expected_filtered_data
        )

        # Verify return value
        assert result == dataset

    def test_update_internal_dataset_filter_none_values(self, mock_dataset_service_dependencies):
        """Test that None values are filtered out except for description field."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="vendor", indexing_technique="high_quality")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        update_data = {
            "name": "new_name",
            "description": None,  # Should be included
            "indexing_technique": "high_quality",
            "retrieval_model": "new_model",
            "embedding_model_provider": None,  # Should be filtered out
            "embedding_model": None,  # Should be filtered out
        }

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify database update was called with filtered data
        expected_filtered_data = {
            "name": "new_name",
            "description": None,  # Description should be included even if None
            "indexing_technique": "high_quality",
            "retrieval_model": "new_model",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        actual_call_args = mock_dataset_service_dependencies[
            "db_session"
        ].query.return_value.filter_by.return_value.update.call_args[0][0]
        # Remove timestamp for comparison as it's dynamic
        del actual_call_args["updated_at"]
        del expected_filtered_data["updated_at"]

        assert actual_call_args == expected_filtered_data

        # Verify return value
        assert result == dataset

    # ==================== Indexing Technique Switch Tests ====================

    def test_update_internal_dataset_indexing_technique_to_economy(
        self, mock_dataset_service_dependencies, mock_internal_provider_dependencies
    ):
        """Test updating internal dataset indexing technique to economy."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="vendor", indexing_technique="high_quality")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        update_data = {"indexing_technique": "economy", "retrieval_model": "new_model"}

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify database update was called with embedding model fields cleared
        expected_filtered_data = {
            "indexing_technique": "economy",
            "embedding_model": None,
            "embedding_model_provider": None,
            "collection_binding_id": None,
            "retrieval_model": "new_model",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        self._assert_database_update_called(
            mock_dataset_service_dependencies["db_session"], "dataset-123", expected_filtered_data
        )

        # Verify return value
        assert result == dataset

    def test_update_internal_dataset_indexing_technique_to_high_quality(
        self, mock_dataset_service_dependencies, mock_internal_provider_dependencies
    ):
        """Test updating internal dataset indexing technique to high_quality."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="vendor", indexing_technique="economy")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        # Mock embedding model
        embedding_model = DatasetUpdateTestDataFactory.create_embedding_model_mock()
        mock_internal_provider_dependencies[
            "model_manager"
        ].return_value.get_model_instance.return_value = embedding_model

        # Mock collection binding
        binding = DatasetUpdateTestDataFactory.create_collection_binding_mock()
        mock_internal_provider_dependencies["get_binding"].return_value = binding

        update_data = {
            "indexing_technique": "high_quality",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
            "retrieval_model": "new_model",
        }

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify embedding model was validated
        mock_internal_provider_dependencies["model_manager"].return_value.get_model_instance.assert_called_once_with(
            tenant_id=mock_internal_provider_dependencies["current_user"].current_tenant_id,
            provider="openai",
            model_type=ModelType.TEXT_EMBEDDING,
            model="text-embedding-ada-002",
        )

        # Verify collection binding was retrieved
        mock_internal_provider_dependencies["get_binding"].assert_called_once_with("openai", "text-embedding-ada-002")

        # Verify database update was called with correct data
        expected_filtered_data = {
            "indexing_technique": "high_quality",
            "embedding_model": "text-embedding-ada-002",
            "embedding_model_provider": "openai",
            "collection_binding_id": "binding-456",
            "retrieval_model": "new_model",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        self._assert_database_update_called(
            mock_dataset_service_dependencies["db_session"], "dataset-123", expected_filtered_data
        )

        # Verify vector index task was triggered
        mock_internal_provider_dependencies["task"].delay.assert_called_once_with("dataset-123", "add")

        # Verify return value
        assert result == dataset

    # ==================== Embedding Model Update Tests ====================

    def test_update_internal_dataset_keep_existing_embedding_model(self, mock_dataset_service_dependencies):
        """Test updating internal dataset without changing embedding model."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(
            provider="vendor",
            indexing_technique="high_quality",
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id="binding-123",
        )
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        update_data = {"name": "new_name", "indexing_technique": "high_quality", "retrieval_model": "new_model"}

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify database update was called with existing embedding model preserved
        expected_filtered_data = {
            "name": "new_name",
            "indexing_technique": "high_quality",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
            "collection_binding_id": "binding-123",
            "retrieval_model": "new_model",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        self._assert_database_update_called(
            mock_dataset_service_dependencies["db_session"], "dataset-123", expected_filtered_data
        )

        # Verify return value
        assert result == dataset

    def test_update_internal_dataset_embedding_model_update(
        self, mock_dataset_service_dependencies, mock_internal_provider_dependencies
    ):
        """Test updating internal dataset with new embedding model."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(
            provider="vendor",
            indexing_technique="high_quality",
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
        )
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        # Mock embedding model
        embedding_model = DatasetUpdateTestDataFactory.create_embedding_model_mock("text-embedding-3-small")
        mock_internal_provider_dependencies[
            "model_manager"
        ].return_value.get_model_instance.return_value = embedding_model

        # Mock collection binding
        binding = DatasetUpdateTestDataFactory.create_collection_binding_mock("binding-789")
        mock_internal_provider_dependencies["get_binding"].return_value = binding

        update_data = {
            "indexing_technique": "high_quality",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-3-small",
            "retrieval_model": "new_model",
        }

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify embedding model was validated
        mock_internal_provider_dependencies["model_manager"].return_value.get_model_instance.assert_called_once_with(
            tenant_id=mock_internal_provider_dependencies["current_user"].current_tenant_id,
            provider="openai",
            model_type=ModelType.TEXT_EMBEDDING,
            model="text-embedding-3-small",
        )

        # Verify collection binding was retrieved
        mock_internal_provider_dependencies["get_binding"].assert_called_once_with("openai", "text-embedding-3-small")

        # Verify database update was called with correct data
        expected_filtered_data = {
            "indexing_technique": "high_quality",
            "embedding_model": "text-embedding-3-small",
            "embedding_model_provider": "openai",
            "collection_binding_id": "binding-789",
            "retrieval_model": "new_model",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        self._assert_database_update_called(
            mock_dataset_service_dependencies["db_session"], "dataset-123", expected_filtered_data
        )

        # Verify vector index task was triggered
        mock_internal_provider_dependencies["task"].delay.assert_called_once_with("dataset-123", "update")

        # Verify return value
        assert result == dataset

    def test_update_internal_dataset_no_indexing_technique_change(self, mock_dataset_service_dependencies):
        """Test updating internal dataset without changing indexing technique."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(
            provider="vendor",
            indexing_technique="high_quality",
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id="binding-123",
        )
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        update_data = {
            "name": "new_name",
            "indexing_technique": "high_quality",  # Same as current
            "retrieval_model": "new_model",
        }

        result = DatasetService.update_dataset("dataset-123", update_data, user)

        # Verify database update was called with correct data
        expected_filtered_data = {
            "name": "new_name",
            "indexing_technique": "high_quality",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
            "collection_binding_id": "binding-123",
            "retrieval_model": "new_model",
            "updated_by": user.id,
            "updated_at": mock_dataset_service_dependencies["current_time"],
        }

        self._assert_database_update_called(
            mock_dataset_service_dependencies["db_session"], "dataset-123", expected_filtered_data
        )

        # Verify return value
        assert result == dataset

    # ==================== Error Handling Tests ====================

    def test_update_dataset_not_found_error(self, mock_dataset_service_dependencies):
        """Test error when dataset is not found."""
        mock_dataset_service_dependencies["get_dataset"].return_value = None

        user = DatasetUpdateTestDataFactory.create_user_mock()
        update_data = {"name": "new_name"}

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset("dataset-123", update_data, user)

        assert "Dataset not found" in str(context.value)

    def test_update_dataset_permission_error(self, mock_dataset_service_dependencies):
        """Test error when user doesn't have permission."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock()
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()
        mock_dataset_service_dependencies["check_permission"].side_effect = NoPermissionError("No permission")

        update_data = {"name": "new_name"}

        with pytest.raises(NoPermissionError):
            DatasetService.update_dataset("dataset-123", update_data, user)

    def test_update_internal_dataset_embedding_model_error(
        self, mock_dataset_service_dependencies, mock_internal_provider_dependencies
    ):
        """Test error when embedding model is not available."""
        dataset = DatasetUpdateTestDataFactory.create_dataset_mock(provider="vendor", indexing_technique="economy")
        mock_dataset_service_dependencies["get_dataset"].return_value = dataset

        user = DatasetUpdateTestDataFactory.create_user_mock()

        # Mock model manager to raise error
        mock_internal_provider_dependencies["model_manager"].return_value.get_model_instance.side_effect = Exception(
            "No Embedding Model available"
        )

        update_data = {
            "indexing_technique": "high_quality",
            "embedding_model_provider": "invalid_provider",
            "embedding_model": "invalid_model",
            "retrieval_model": "new_model",
        }

        with pytest.raises(Exception) as context:
            DatasetService.update_dataset("dataset-123", update_data, user)

        assert "No Embedding Model available".lower() in str(context.value).lower()