from unittest.mock import MagicMock, patch import pytest from faker import Faker from sqlalchemy import select 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 import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, DatasetAutoDisableLog, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from tasks.add_document_to_index_task import add_document_to_index_task class TestAddDocumentToIndexTask: """Integration tests for add_document_to_index_task using testcontainers.""" @pytest.fixture def mock_external_service_dependencies(self): """Mock setup for external service dependencies.""" with ( patch( "tasks.add_document_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): """ Helper method to create test document segments. Args: db_session_with_containers: Database session from testcontainers infrastructure document: Document instance dataset: Dataset instance Returns: list: List of created DocumentSegment instances """ fake = Faker() segments = [] for i in range(3): segment = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=i, content=fake.text(max_nb_chars=200), word_count=len(fake.text(max_nb_chars=200).split()), tokens=len(fake.text(max_nb_chars=200).split()) * 2, index_node_id=f"node_{i}", index_node_hash=f"hash_{i}", enabled=False, status=SegmentStatus.COMPLETED, created_by=document.created_by, ) db_session_with_containers.add(segment) segments.append(segment) db_session_with_containers.commit() return segments def test_add_document_to_index_success( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test successful document indexing with paragraph index type. This test verifies: - Proper document retrieval from database - Correct segment processing and document creation - Index processor integration - Database state updates - Segment status changes - Redis cache key deletion """ # 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 key to simulate indexing in progress indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # 5 minutes expiry # Verify cache key exists assert redis_client.exists(indexing_cache_key) == 1 # Act: Execute the task add_document_to_index_task(document.id) # Assert: Verify the expected outcomes # Verify index processor was called correctly mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.PARAGRAPH_INDEX ) mock_external_service_dependencies["index_processor"].load.assert_called_once() # Verify database state changes db_session_with_containers.refresh(document) for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is True assert segment.disabled_at is None assert segment.disabled_by is None # Verify Redis cache key was deleted assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_with_different_index_type( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test document 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 key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Act: Execute the task add_document_to_index_task(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 database state changes db_session_with_containers.refresh(document) for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is True assert segment.disabled_at is None assert segment.disabled_by is None # Verify Redis cache key was deleted assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_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 - Redis cache key not affected (since it was never created) """ # Arrange: Use non-existent document ID fake = Faker() non_existent_id = str(fake.uuid4()) # Act: Execute the task with non-existent document add_document_to_index_task(non_existent_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() # Note: redis_client.delete is not called when document is not found # because indexing_cache_key is not defined in that case def test_add_document_to_index_invalid_indexing_status( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test handling of document with invalid indexing status. This test verifies: - Early return when indexing_status is not "completed" - No index processing for documents not ready for indexing - Proper database session cleanup - No unnecessary external service calls - Redis cache key not affected """ # Arrange: Create test data with invalid indexing status _, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Set invalid indexing status document.indexing_status = IndexingStatus.INDEXING db_session_with_containers.commit() # Act: Execute the task add_document_to_index_task(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_add_document_to_index_dataset_not_found( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test handling when document's dataset doesn't exist. This test verifies: - Proper error handling when dataset is missing - Document status is set to error - Document is disabled - Error information is recorded - Redis cache is cleared despite error """ # Arrange: Create test data dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Set up Redis cache key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Delete the dataset to simulate dataset not found scenario db_session_with_containers.delete(dataset) db_session_with_containers.commit() # Act: Execute the task add_document_to_index_task(document.id) # Assert: Verify error handling db_session_with_containers.refresh(document) assert document.enabled is False assert document.indexing_status == IndexingStatus.ERROR assert document.error is not None assert "doesn't exist" in document.error assert document.disabled_at is not None # Verify no index processing occurred mock_external_service_dependencies["index_processor_factory"].assert_not_called() mock_external_service_dependencies["index_processor"].load.assert_not_called() # Verify redis cache was cleared despite error assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_with_parent_child_structure( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test document 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 key indexing_cache_key = f"document_{document.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 add_document_to_index_task(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 database state changes db_session_with_containers.refresh(document) for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is True assert segment.disabled_at is None assert segment.disabled_by is None # Verify redis cache was cleared assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_with_already_enabled_segments( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test document indexing when segments are already enabled. This test verifies: - Segments with status=SegmentStatus.COMPLETED are processed regardless of enabled status - Index processing occurs with all completed segments - Auto disable log deletion still occurs - Redis cache is cleared """ # Arrange: Create test data dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Create segments that are already enabled fake = Faker() segments = [] for i in range(3): segment = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=i, content=fake.text(max_nb_chars=200), word_count=len(fake.text(max_nb_chars=200).split()), tokens=len(fake.text(max_nb_chars=200).split()) * 2, index_node_id=f"node_{i}", index_node_hash=f"hash_{i}", enabled=True, # Already enabled status=SegmentStatus.COMPLETED, created_by=document.created_by, ) db_session_with_containers.add(segment) segments.append(segment) db_session_with_containers.commit() # Set up Redis cache key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Act: Execute the task add_document_to_index_task(document.id) # Assert: Verify index processing occurred but with empty documents list mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.PARAGRAPH_INDEX ) mock_external_service_dependencies["index_processor"].load.assert_called_once() # Verify the load method was called with all completed segments # (implementation doesn't filter by enabled status, only by status=SegmentStatus.COMPLETED) 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 # All completed segments are processed # Verify redis cache was cleared assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_auto_disable_log_deletion( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test that auto disable logs are properly deleted during indexing. This test verifies: - Auto disable log entries are deleted for the document - Database state is properly managed - Index processing continues normally - Redis cache key deletion """ # 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) # Create some auto disable log entries fake = Faker() auto_disable_logs = [] for _ in range(2): log_entry = DatasetAutoDisableLog( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, ) log_entry.id = str(fake.uuid4()) db_session_with_containers.add(log_entry) auto_disable_logs.append(log_entry) db_session_with_containers.commit() # Set up Redis cache key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Verify logs exist before processing existing_logs = db_session_with_containers.scalars( select(DatasetAutoDisableLog).where(DatasetAutoDisableLog.document_id == document.id) ).all() assert len(existing_logs) == 2 # Act: Execute the task add_document_to_index_task(document.id) # Assert: Verify auto disable logs were deleted remaining_logs = db_session_with_containers.scalars( select(DatasetAutoDisableLog).where(DatasetAutoDisableLog.document_id == document.id) ).all() assert len(remaining_logs) == 0 # Verify index processing occurred normally mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.PARAGRAPH_INDEX ) mock_external_service_dependencies["index_processor"].load.assert_called_once() # Verify segments were enabled for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is True # Verify redis cache was cleared assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_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 - Document status is set to error - Document is 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 key indexing_cache_key = f"document_{document.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 add_document_to_index_task(document.id) # Assert: Verify error handling db_session_with_containers.refresh(document) assert document.enabled is False assert document.indexing_status == IndexingStatus.ERROR assert document.error is not None assert "Index processing failed" in document.error assert document.disabled_at is not None # Verify segments were not enabled due to error for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is False # Should remain disabled due to error # Verify redis cache was still cleared despite error assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_segment_filtering_edge_cases( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test segment filtering with various edge cases. This test verifies: - Only segments with status=SegmentStatus.COMPLETED are processed (regardless of enabled status) - Segments with status!="completed" are NOT processed - Segments are ordered by position correctly - Mixed segment states are handled properly - All segments are updated to enabled=True after processing - Redis cache key deletion """ # Arrange: Create test data dataset, document = self._create_test_dataset_and_document( db_session_with_containers, mock_external_service_dependencies ) # Create segments with mixed states fake = Faker() segments = [] # Segment 1: Should be processed (enabled=False, status=SegmentStatus.COMPLETED) segment1 = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=0, content=fake.text(max_nb_chars=200), word_count=len(fake.text(max_nb_chars=200).split()), tokens=len(fake.text(max_nb_chars=200).split()) * 2, index_node_id="node_0", index_node_hash="hash_0", enabled=False, status=SegmentStatus.COMPLETED, created_by=document.created_by, ) db_session_with_containers.add(segment1) segments.append(segment1) # Segment 2: Should be processed (enabled=True, status=SegmentStatus.COMPLETED) # Note: Implementation doesn't filter by enabled status, only by status=SegmentStatus.COMPLETED segment2 = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=1, content=fake.text(max_nb_chars=200), word_count=len(fake.text(max_nb_chars=200).split()), tokens=len(fake.text(max_nb_chars=200).split()) * 2, index_node_id="node_1", index_node_hash="hash_1", enabled=True, # Already enabled, but will still be processed status=SegmentStatus.COMPLETED, created_by=document.created_by, ) db_session_with_containers.add(segment2) segments.append(segment2) # Segment 3: Should NOT be processed (enabled=False, status="processing") segment3 = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=2, content=fake.text(max_nb_chars=200), word_count=len(fake.text(max_nb_chars=200).split()), tokens=len(fake.text(max_nb_chars=200).split()) * 2, index_node_id="node_2", index_node_hash="hash_2", enabled=False, status=SegmentStatus.INDEXING, # Not completed created_by=document.created_by, ) db_session_with_containers.add(segment3) segments.append(segment3) # Segment 4: Should be processed (enabled=False, status=SegmentStatus.COMPLETED) segment4 = DocumentSegment( tenant_id=document.tenant_id, dataset_id=dataset.id, document_id=document.id, position=3, content=fake.text(max_nb_chars=200), word_count=len(fake.text(max_nb_chars=200).split()), tokens=len(fake.text(max_nb_chars=200).split()) * 2, index_node_id="node_3", index_node_hash="hash_3", enabled=False, status=SegmentStatus.COMPLETED, created_by=document.created_by, ) db_session_with_containers.add(segment4) segments.append(segment4) db_session_with_containers.commit() # Set up Redis cache key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Act: Execute the task add_document_to_index_task(document.id) # Assert: Verify only eligible segments were processed mock_external_service_dependencies["index_processor_factory"].assert_called_once_with( IndexStructureType.PARAGRAPH_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 with status=SegmentStatus.COMPLETED should be processed # Verify correct segments were processed (by position order) # Segments 1, 2, 4 should be processed (positions 0, 1, 3) # Segment 3 is skipped (position 2, status="processing") assert documents[0].metadata["doc_id"] == "node_0" # segment1, position 0 assert documents[1].metadata["doc_id"] == "node_1" # segment2, position 1 assert documents[2].metadata["doc_id"] == "node_3" # segment4, position 3 # Verify database state changes db_session_with_containers.refresh(document) db_session_with_containers.refresh(segment1) db_session_with_containers.refresh(segment2) db_session_with_containers.refresh(segment3) db_session_with_containers.refresh(segment4) # All segments should be enabled because the task updates ALL segments for the document assert segment1.enabled is True assert segment2.enabled is True # Was already enabled, stays True assert segment3.enabled is True # Was not processed but still updated to True assert segment4.enabled is True # Verify redis cache was cleared assert redis_client.exists(indexing_cache_key) == 0 def test_add_document_to_index_comprehensive_error_scenarios( self, db_session_with_containers: Session, mock_external_service_dependencies ): """ Test comprehensive error scenarios and recovery. This test verifies: - Multiple types of exceptions are handled properly - Error state is consistently managed - Resource cleanup occurs in all error cases - Database session management is robust - Redis cache key deletion in all scenarios """ # 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) # Test different exception types test_exceptions = [ ("Database connection error", Exception("Database connection failed")), ("Index processor error", RuntimeError("Index processor initialization failed")), ("Memory error", MemoryError("Out of memory")), ("Value error", ValueError("Invalid index type")), ] for error_name, exception in test_exceptions: # Reset mocks for each test mock_external_service_dependencies["index_processor"].load.side_effect = exception # Reset document state document.enabled = True document.indexing_status = "completed" document.error = None document.disabled_at = None db_session_with_containers.commit() # Set up Redis cache key indexing_cache_key = f"document_{document.id}_indexing" redis_client.set(indexing_cache_key, "processing", ex=300) # Act: Execute the task add_document_to_index_task(document.id) # Assert: Verify consistent error handling db_session_with_containers.refresh(document) assert document.enabled is False, f"Document should be disabled for {error_name}" assert document.indexing_status == IndexingStatus.ERROR, f"Document status should be error for {error_name}" assert document.error is not None, f"Error should be recorded for {error_name}" assert str(exception) in document.error, f"Error message should contain exception for {error_name}" assert document.disabled_at is not None, f"Disabled timestamp should be set for {error_name}" # Verify segments remain disabled due to error for segment in segments: db_session_with_containers.refresh(segment) assert segment.enabled is False, f"Segments should remain disabled for {error_name}" # Verify redis cache was still cleared despite error assert redis_client.exists(indexing_cache_key) == 0, f"Redis cache should be cleared for {error_name}"