from unittest.mock import MagicMock, patch import pytest from faker import Faker from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from extensions.ext_redis import redis_client from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from tasks.enable_segments_to_index_task import enable_segments_to_index_task class TestEnableSegmentsToIndexTask: """Integration tests for enable_segments_to_index_task using testcontainers.""" @pytest.fixture def mock_external_service_dependencies(self): """Mock setup for external service dependencies.""" with ( patch( "tasks.enable_segments_to_index_task.IndexProcessorFactory", autospec=True ) as mock_index_processor_factory, ): # Setup mock index processor mock_processor = MagicMock() mock_index_processor_factory.return_value.init_index_processor.return_value = mock_processor yield { "index_processor_factory": mock_index_processor_factory, "index_processor": mock_processor, } def _create_test_dataset_and_document( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Helper method to create a test dataset and document for testing. Args: db_session_with_containers: Database session from testcontainers infrastructure mock_external_service_dependencies: Mock dependencies Returns: tuple: (dataset, document) - Created dataset and document instances """ fake = Faker() # Create account and tenant account = Account( email=fake.email(), name=fake.name(), interface_language="en-US", status="active", ) db_session_with_containers.add(account) db_session_with_containers.commit() tenant = Tenant( name=fake.company(), status="normal", ) db_session_with_containers.add(tenant) db_session_with_containers.commit() # Create tenant-account join join = TenantAccountJoin( tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.OWNER, current=True, ) db_session_with_containers.add(join) db_session_with_containers.commit() # Create dataset dataset = Dataset( id=fake.uuid4(), tenant_id=tenant.id, name=fake.company(), description=fake.text(max_nb_chars=100), data_source_type=DataSourceType.UPLOAD_FILE, indexing_technique=IndexTechniqueType.HIGH_QUALITY, created_by=account.id, ) db_session_with_containers.add(dataset) db_session_with_containers.commit() # Create document document = Document( id=fake.uuid4(), tenant_id=tenant.id, dataset_id=dataset.id, position=1, data_source_type=DataSourceType.UPLOAD_FILE, batch="test_batch", name=fake.file_name(), created_from=DocumentCreatedFrom.WEB, created_by=account.id, indexing_status=IndexingStatus.COMPLETED, enabled=True, doc_form=IndexStructureType.PARAGRAPH_INDEX, ) db_session_with_containers.add(document) db_session_with_containers.commit() # Refresh dataset to ensure doc_form property works correctly db_session_with_containers.refresh(dataset) return dataset, document def _create_test_segments( self, db_session_with_containers: Session, document, dataset, count=3, enabled=False, status=SegmentStatus.COMPLETED, ): """ Helper method to create test document segments. Args: db_session_with_containers: Database session from testcontainers infrastructure document: Document instance dataset: Dataset instance count: Number of segments to create enabled: Whether segments should be enabled status: Status of the segments Returns: list: List of created DocumentSegment instances """ fake = Faker() segments = [] for i in range(count): text = fake.text(max_nb_chars=200) segment = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=i, content=text, word_count=len(text.split()), tokens=len(text.split()) * 2, index_node_id=f"node_{i}", index_node_hash=f"hash_{i}", enabled=enabled, status=status, created_by=document.created_by, ) db_session_with_containers.add(segment) segments.append(segment) db_session_with_containers.commit() return segments def test_enable_segments_to_index_with_different_index_type( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test segments indexing with different index types. This test verifies: - Proper handling of different index types - Index processor factory integration - Document processing with various configurations - Redis cache key deletion """ # Arrange: Create test data with different index type dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Update document to use different index type document.doc_form = IndexStructureType.QA_INDEX db_session_with_containers.commit() # Refresh dataset to ensure doc_form property reflects the updated document db_session_with_containers.refresh(dataset) # Create segments segments = self._create_test_segments(db_session_with_containers, document, dataset) # Set up Redis cache keys segment_ids = [segment.id for segment in segments] for segment in segments: indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Act: Execute the task enable_segments_to_index_task(segment_ids, dataset.id, document.id) # Assert: Verify different index type handling mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.QA_INDEX ) mock_external_service_dependencies["index_processor"].load.assert_called_once() # Verify the load method was called with correct parameters call_args = mock_external_service_dependencies["index_processor"].load.call_args assert call_args is not None documents = call_args[0][1] # Second argument should be documents list assert len(documents) == 3 # Verify Redis cache keys were deleted for segment in segments: indexing_cache_key = f"segment_{segment.id}_indexing" assert redis_client.exists(indexing_cache_key) == 0 def test_enable_segments_to_index_dataset_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test handling of non-existent dataset. This test verifies: - Proper error handling for missing datasets - Early return without processing - Database session cleanup - No unnecessary index processor calls """ # Arrange: Use non-existent dataset ID fake = Faker() non_existent_dataset_id = fake.uuid4() non_existent_document_id = fake.uuid4() segment_ids = [fake.uuid4()] # Act: Execute the task with non-existent dataset enable_segments_to_index_task(segment_ids, non_existent_dataset_id, non_existent_document_id) # Assert: Verify no processing occurred mock_external_service_dependencies["index_processor_factory"].assert_not_called() mock_external_service_dependencies["index_processor"].load.assert_not_called() def test_enable_segments_to_index_document_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test handling of non-existent document. This test verifies: - Proper error handling for missing documents - Early return without processing - Database session cleanup - No unnecessary index processor calls """ # Arrange: Create dataset but use non-existent document ID dataset, _ = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) fake = Faker() non_existent_document_id = fake.uuid4() segment_ids = [fake.uuid4()] # Act: Execute the task with non-existent document enable_segments_to_index_task(segment_ids, dataset.id, non_existent_document_id) # Assert: Verify no processing occurred mock_external_service_dependencies["index_processor_factory"].assert_not_called() mock_external_service_dependencies["index_processor"].load.assert_not_called() def test_enable_segments_to_index_invalid_document_status( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test handling of document with invalid status. This test verifies: - Early return when document is disabled, archived, or not completed - No index processing for documents not ready for indexing - Proper database session cleanup - No unnecessary external service calls """ # Arrange: Create test data with invalid document status dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Test different invalid statuses invalid_statuses = [ ("disabled", {"enabled": False}), ("archived", {"archived": True}), ("not_completed", {"indexing_status": IndexingStatus.INDEXING}), ] for _, status_attrs in invalid_statuses: # Reset document status document.enabled = True document.archived = False document.indexing_status = "completed" db_session_with_containers.commit() # Set invalid status for attr, value in status_attrs.items(): setattr(document, attr, value) db_session_with_containers.commit() # Create segments segments = self._create_test_segments(db_session_with_containers, document, dataset) segment_ids = [segment.id for segment in segments] # Act: Execute the task enable_segments_to_index_task(segment_ids, dataset.id, document.id) # Assert: Verify no processing occurred mock_external_service_dependencies["index_processor_factory"].assert_not_called() mock_external_service_dependencies["index_processor"].load.assert_not_called() # Clean up segments for next iteration for segment in segments: db_session_with_containers.delete(segment) db_session_with_containers.commit() def test_enable_segments_to_index_segments_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test handling when no segments are found. This test verifies: - Proper handling when segments don't exist - Early return without processing - Database session cleanup - Index processor is created but load is not called """ # Arrange: Create test data dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Use non-existent segment IDs fake = Faker() non_existent_segment_ids = [fake.uuid4() for _ in range(3)] # Act: Execute the task with non-existent segments enable_segments_to_index_task(non_existent_segment_ids, dataset.id, document.id) # Assert: Verify index processor was created but load was not called mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.PARAGRAPH_INDEX ) mock_external_service_dependencies["index_processor"].load.assert_not_called() def test_enable_segments_to_index_with_parent_child_structure( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test segments indexing with parent-child structure. This test verifies: - Proper handling of PARENT_CHILD_INDEX type - Child document creation from segments - Correct document structure for parent-child indexing - Index processor receives properly structured documents - Redis cache key deletion """ # Arrange: Create test data with parent-child index type dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Update document to use parent-child index type document.doc_form = IndexStructureType.PARENT_CHILD_INDEX db_session_with_containers.commit() # Refresh dataset to ensure doc_form property reflects the updated document db_session_with_containers.refresh(dataset) # Create segments with mock child chunks segments = self._create_test_segments(db_session_with_containers, document, dataset) # Set up Redis cache keys segment_ids = [segment.id for segment in segments] for segment in segments: indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Mock the get_child_chunks method for each segment with patch.object(DocumentSegment, "get_child_chunks", autospec=True) as mock_get_child_chunks: # Setup mock to return child chunks for each segment mock_child_chunks = [] for i in range(2): # Each segment has 2 child chunks mock_child = MagicMock() mock_child.content = f"child_content_{i}" mock_child.index_node_id = f"child_node_{i}" mock_child.index_node_hash = f"child_hash_{i}" mock_child_chunks.append(mock_child) mock_get_child_chunks.return_value = mock_child_chunks # Act: Execute the task enable_segments_to_index_task(segment_ids, dataset.id, document.id) # Assert: Verify parent-child index processing mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.PARENT_CHILD_INDEX ) mock_external_service_dependencies["index_processor"].load.assert_called_once() # Verify the load method was called with correct parameters call_args = mock_external_service_dependencies["index_processor"].load.call_args assert call_args is not None documents = call_args[0][1] # Second argument should be documents list assert len(documents) == 3 # 3 segments # Verify each document has children for doc in documents: assert hasattr(doc, "children") assert len(doc.children) == 2 # Each document has 2 children # Verify Redis cache keys were deleted for segment in segments: indexing_cache_key = f"segment_{segment.id}_indexing" assert redis_client.exists(indexing_cache_key) == 0 def test_enable_segments_to_index_general_exception_handling( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test general exception handling during indexing process. This test verifies: - Exceptions are properly caught and handled - Segment status is set to error - Segments are disabled - Error information is recorded - Redis cache is still cleared - Database session is properly closed """ # Arrange: Create test data dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) segments = self._create_test_segments(db_session_with_containers, document, dataset) # Set up Redis cache keys segment_ids = [segment.id for segment in segments] for segment in segments: indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Mock the index processor to raise an exception mock_external_service_dependencies["index_processor"].load.side_effect = Exception("Index processing failed") # Act: Execute the task enable_segments_to_index_task(segment_ids, dataset.id, document.id) # Assert: Verify error handling for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is False assert segment.status == SegmentStatus.ERROR assert segment.error is not None assert "Index processing failed" in segment.error assert segment.disabled_at is not None # Verify Redis cache keys were still cleared despite error for segment in segments: indexing_cache_key = f"segment_{segment.id}_indexing" assert redis_client.exists(indexing_cache_key) == 0