""" Integration tests for disable_segment_from_index_task using TestContainers. This module provides comprehensive integration tests for the disable_segment_from_index_task using real database and Redis containers to ensure the task works correctly with actual data and external dependencies. """ import logging import time from datetime import UTC, datetime from unittest.mock import 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 import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from tasks.disable_segment_from_index_task import disable_segment_from_index_task logger = logging.getLogger(__name__) class TestDisableSegmentFromIndexTask: """Integration tests for disable_segment_from_index_task using testcontainers.""" @pytest.fixture def mock_index_processor(self): """Mock IndexProcessorFactory and its clean method.""" with patch("tasks.disable_segment_from_index_task.IndexProcessorFactory") as mock_factory: mock_processor = mock_factory.return_value.init_index_processor.return_value mock_processor.clean.return_value = None yield mock_processor def _create_test_account_and_tenant(self, db_session_with_containers: Session) -> tuple[Account, Tenant]: """ Helper method to create a test account and tenant for testing. Args: db_session_with_containers: Database session from testcontainers infrastructure Returns: tuple: (account, tenant) - Created account and tenant instances """ fake = Faker() # Create account 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() # Create tenant tenant = Tenant( name=fake.company(), status="normal", plan="basic", ) db_session_with_containers.add(tenant) db_session_with_containers.commit() # Create tenant-account join with owner role 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() # Set current tenant for account account.current_tenant = tenant return account, tenant def _create_test_dataset(self, db_session_with_containers: Session, tenant: Tenant, account: Account) -> Dataset: """ Helper method to create a test dataset. Args: tenant: Tenant instance account: Account instance Returns: Dataset: Created dataset instance """ fake = Faker() dataset = Dataset( tenant_id=tenant.id, name=fake.sentence(nb_words=3), description=fake.text(max_nb_chars=200), 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() return dataset def _create_test_document( self, db_session_with_containers: Session, dataset: Dataset, tenant: Tenant, account: Account, doc_form: str = IndexStructureType.PARAGRAPH_INDEX, ) -> Document: """ Helper method to create a test document. Args: dataset: Dataset instance tenant: Tenant instance account: Account instance doc_form: Document form type Returns: Document: Created document instance """ fake = Faker() document = Document( tenant_id=tenant.id, dataset_id=dataset.id, position=1, data_source_type=DataSourceType.UPLOAD_FILE, batch=fake.uuid4(), name=fake.file_name(), created_from=DocumentCreatedFrom.API, created_by=account.id, indexing_status=IndexingStatus.COMPLETED, enabled=True, archived=False, doc_form=doc_form, word_count=1000, tokens=500, completed_at=datetime.now(UTC), ) db_session_with_containers.add(document) db_session_with_containers.commit() return document def _create_test_segment( self, db_session_with_containers: Session, document: Document, dataset: Dataset, tenant: Tenant, account: Account, status: SegmentStatus = SegmentStatus.COMPLETED, enabled: bool = True, ) -> DocumentSegment: """ Helper method to create a test document segment. Args: document: Document instance dataset: Dataset instance tenant: Tenant instance account: Account instance status: Segment status enabled: Whether segment is enabled Returns: DocumentSegment: Created segment instance """ fake = Faker() segment = DocumentSegment( tenant_id=tenant.id, dataset_id=dataset.id, document_id=document.id, position=1, content=fake.text(max_nb_chars=500), word_count=100, tokens=50, index_node_id=fake.uuid4(), index_node_hash=fake.sha256(), status=status, enabled=enabled, created_by=account.id, completed_at=datetime.now(UTC) if status == SegmentStatus.COMPLETED else None, ) db_session_with_containers.add(segment) db_session_with_containers.commit() return segment def test_disable_segment_success(self, db_session_with_containers: Session, mock_index_processor): """ Test successful segment disabling from index. This test verifies: - Segment is found and validated - Index processor clean method is called with correct parameters - Redis cache is cleared - Task completes successfully """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Set up Redis cache indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.setex(indexing_cache_key, 600, 1) # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task completed successfully assert result is None # Task returns None on success # Verify index processor was called correctly mock_index_processor.clean.assert_called_once() call_args = mock_index_processor.clean.call_args assert call_args[0][0].id == dataset.id # Check dataset ID assert call_args[0][1] == [segment.index_node_id] # Check index node IDs # Verify Redis cache was cleared assert redis_client.get(indexing_cache_key) is None # Verify segment is still in database db_session_with_containers.refresh(segment) assert segment.id is not None def test_disable_segment_not_found(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when segment is not found. This test verifies: - Task handles non-existent segment gracefully - No index processor operations are performed - Task returns early without errors """ # Arrange: Use a non-existent segment ID fake = Faker() non_existent_segment_id = fake.uuid4() # Act: Execute the task with non-existent segment result = disable_segment_from_index_task(non_existent_segment_id) # Assert: Verify the task handled the error gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_not_completed(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when segment is not in completed status. This test verifies: - Task rejects segments that are not completed - No index processor operations are performed - Task returns early without errors """ # Arrange: Create test data with non-completed segment account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment( db_session_with_containers, document, dataset, tenant, account, status=SegmentStatus.INDEXING, enabled=True ) # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the invalid status gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_no_dataset(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when segment has no associated dataset. This test verifies: - Task handles segments without dataset gracefully - No index processor operations are performed - Task returns early without errors """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Manually remove dataset association segment.dataset_id = "00000000-0000-0000-0000-000000000000" db_session_with_containers.commit() # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the missing dataset gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_no_document(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when segment has no associated document. This test verifies: - Task handles segments without document gracefully - No index processor operations are performed - Task returns early without errors """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Manually remove document association segment.document_id = "00000000-0000-0000-0000-000000000000" db_session_with_containers.commit() # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the missing document gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_document_disabled(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when document is disabled. This test verifies: - Task handles disabled documents gracefully - No index processor operations are performed - Task returns early without errors """ # Arrange: Create test data with disabled document account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) document.enabled = False db_session_with_containers.commit() segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the disabled document gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_document_archived(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when document is archived. This test verifies: - Task handles archived documents gracefully - No index processor operations are performed - Task returns early without errors """ # Arrange: Create test data with archived document account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) document.archived = True db_session_with_containers.commit() segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the archived document gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_document_indexing_not_completed( self, db_session_with_containers: Session, mock_index_processor ): """ Test handling when document indexing is not completed. This test verifies: - Task handles documents with incomplete indexing gracefully - No index processor operations are performed - Task returns early without errors """ # Arrange: Create test data with incomplete indexing account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) document.indexing_status = "indexing" db_session_with_containers.commit() segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the incomplete indexing gracefully assert result is None # Verify index processor was not called mock_index_processor.clean.assert_not_called() def test_disable_segment_index_processor_exception(self, db_session_with_containers: Session, mock_index_processor): """ Test handling when index processor raises an exception. This test verifies: - Task handles index processor exceptions gracefully - Segment is re-enabled on failure - Redis cache is still cleared - Database changes are committed """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Set up Redis cache indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.setex(indexing_cache_key, 600, 1) # Configure mock to raise exception mock_index_processor.clean.side_effect = Exception("Index processor error") # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task handled the exception gracefully assert result is None # Verify index processor was called mock_index_processor.clean.assert_called_once() call_args = mock_index_processor.clean.call_args # Check that the call was made with the correct parameters assert len(call_args[0]) == 2 # Check two arguments were passed assert call_args[0][1] == [segment.index_node_id] # Check index node IDs # Verify segment was re-enabled db_session_with_containers.refresh(segment) assert segment.enabled is True # Verify Redis cache was still cleared assert redis_client.get(indexing_cache_key) is None def test_disable_segment_different_doc_forms(self, db_session_with_containers: Session, mock_index_processor): """ Test disabling segments with different document forms. This test verifies: - Task works with different document form types - Correct index processor is initialized for each form - Index processor clean method is called correctly """ # Test different document forms doc_forms = [ IndexStructureType.PARAGRAPH_INDEX, IndexStructureType.QA_INDEX, IndexStructureType.PARENT_CHILD_INDEX, ] for doc_form in doc_forms: # Arrange: Create test data for each form account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document( db_session_with_containers, dataset, tenant, account, doc_form=doc_form ) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Reset mock for each iteration mock_index_processor.reset_mock() # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify the task completed successfully assert result is None # Verify correct index processor was initialized mock_index_processor.clean.assert_called_once() call_args = mock_index_processor.clean.call_args assert call_args[0][0].id == dataset.id # Check dataset ID assert call_args[0][1] == [segment.index_node_id] # Check index node IDs def test_disable_segment_redis_cache_handling(self, db_session_with_containers: Session, mock_index_processor): """ Test Redis cache handling during segment disabling. This test verifies: - Redis cache is properly set before task execution - Cache is cleared after task completion - Cache handling works with different scenarios """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Test with cache present indexing_cache_key = f"segment_{segment.id}_indexing" redis_client.setex(indexing_cache_key, 600, 1) assert redis_client.get(indexing_cache_key) is not None # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify cache was cleared assert result is None assert redis_client.get(indexing_cache_key) is None # Test with no cache present segment2 = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) result2 = disable_segment_from_index_task(segment2.id) # Assert: Verify task still works without cache assert result2 is None def test_disable_segment_performance_timing(self, db_session_with_containers: Session, mock_index_processor): """ Test performance timing of segment disabling task. This test verifies: - Task execution time is reasonable - Performance logging works correctly - Task completes within expected time bounds """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Act: Execute the task and measure time start_time = time.perf_counter() result = disable_segment_from_index_task(segment.id) end_time = time.perf_counter() # Assert: Verify task completed successfully and timing is reasonable assert result is None execution_time = end_time - start_time assert execution_time < 5.0 # Should complete within 5 seconds def test_disable_segment_database_session_management( self, db_session_with_containers: Session, mock_index_processor ): """ Test database session management during task execution. This test verifies: - Database sessions are properly managed - Sessions are closed after task completion - No session leaks occur """ # Arrange: Create test data account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) # Act: Execute the task result = disable_segment_from_index_task(segment.id) # Assert: Verify task completed and session management worked assert result is None # Verify segment is still accessible (session was properly managed) db_session_with_containers.refresh(segment) assert segment.id is not None def test_disable_segment_concurrent_execution(self, db_session_with_containers: Session, mock_index_processor): """ Test concurrent execution of segment disabling tasks. This test verifies: - Multiple tasks can run concurrently - Each task processes its own segment correctly - No interference between concurrent tasks """ # Arrange: Create multiple test segments account, tenant = self._create_test_account_and_tenant(db_session_with_containers) dataset = self._create_test_dataset(db_session_with_containers, tenant, account) document = self._create_test_document(db_session_with_containers, dataset, tenant, account) segments = [] for i in range(3): segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account) segments.append(segment) # Act: Execute tasks concurrently (simulated) results = [] for segment in segments: result = disable_segment_from_index_task(segment.id) results.append(result) # Assert: Verify all tasks completed successfully assert all(result is None for result in results) # Verify all segments were processed assert mock_index_processor.clean.call_count == len(segments) # Verify each segment was processed with correct parameters for segment in segments: # Check that clean was called with this segment's dataset and index_node_id found = False for call in mock_index_processor.clean.call_args_list: if call[0][0].id == dataset.id and call[0][1] == [segment.index_node_id]: found = True break assert found, f"Segment {segment.id} was not processed correctly"