from unittest.mock import MagicMock, patch import pytest from faker import Faker from sqlalchemy import func, select from sqlalchemy.orm import Session from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from models import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole, TenantStatus from models.dataset import Dataset, Document, DocumentSegment from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus from tasks.document_indexing_update_task import document_indexing_update_task class TestDocumentIndexingUpdateTask: @pytest.fixture def mock_external_dependencies(self): """Patch external collaborators used by the update task. - IndexProcessorFactory.init_index_processor().clean(...) - IndexingRunner.run([...]) """ with ( patch("tasks.document_indexing_update_task.IndexProcessorFactory", autospec=True) as mock_factory, patch("tasks.document_indexing_update_task.IndexingRunner", autospec=True) as mock_runner, ): processor_instance = MagicMock() mock_factory.return_value.init_index_processor.return_value = processor_instance runner_instance = mock_runner.return_value yield { "factory": mock_factory, "processor": processor_instance, "runner": mock_runner, "runner_instance": runner_instance, } def _create_dataset_document_with_segments(self, db_session_with_containers: Session, *, segment_count: int = 2): fake = Faker() # Account and tenant account = Account( email=fake.email(), name=fake.name(), interface_language="en-US", status=AccountStatus.ACTIVE, ) db_session_with_containers.add(account) db_session_with_containers.commit() tenant = Tenant(name=fake.company(), status=TenantStatus.NORMAL) db_session_with_containers.add(tenant) db_session_with_containers.commit() 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() # Dataset and document dataset = Dataset( tenant_id=tenant.id, name=fake.company(), description=fake.text(max_nb_chars=64), 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() document = Document( tenant_id=tenant.id, dataset_id=dataset.id, position=0, data_source_type=DataSourceType.UPLOAD_FILE, batch="test_batch", name=fake.file_name(), created_from=DocumentCreatedFrom.WEB, created_by=account.id, indexing_status=IndexingStatus.WAITING, enabled=True, doc_form=IndexStructureType.PARAGRAPH_INDEX, ) db_session_with_containers.add(document) db_session_with_containers.commit() # Segments node_ids = [] for i in range(segment_count): node_id = f"node-{i + 1}" seg = DocumentSegment( tenant_id=tenant.id, dataset_id=dataset.id, document_id=document.id, position=i, content=fake.text(max_nb_chars=32), answer=None, word_count=10, tokens=5, index_node_id=node_id, status=SegmentStatus.COMPLETED, created_by=account.id, ) db_session_with_containers.add(seg) node_ids.append(node_id) db_session_with_containers.commit() # Refresh to ensure ORM state db_session_with_containers.refresh(dataset) db_session_with_containers.refresh(document) return dataset, document, node_ids def test_cleans_segments_and_reindexes(self, db_session_with_containers: Session, mock_external_dependencies): dataset, document, node_ids = self._create_dataset_document_with_segments(db_session_with_containers) # Act document_indexing_update_task(dataset.id, document.id) # Ensure we see committed changes from another session db_session_with_containers.expire_all() # Assert document status updated before reindex updated = db_session_with_containers.scalar(select(Document).where(Document.id == document.id).limit(1)) assert updated.indexing_status == IndexingStatus.PARSING assert updated.processing_started_at is not None # Segments should be deleted remaining = db_session_with_containers.scalar( select(func.count()).select_from(DocumentSegment).where(DocumentSegment.document_id == document.id) ) assert remaining == 0 # Assert index processor clean was called with expected args clean_call = mock_external_dependencies["processor"].clean.call_args assert clean_call is not None args, kwargs = clean_call # args[0] is a Dataset instance (from another session) — validate by id assert getattr(args[0], "id", None) == dataset.id # args[1] should contain our node_ids assert set(args[1]) == set(node_ids) assert kwargs.get("with_keywords") is True assert kwargs.get("delete_child_chunks") is True # Assert indexing runner invoked with the updated document run_call = mock_external_dependencies["runner_instance"].run.call_args assert run_call is not None run_docs = run_call[0][0] assert len(run_docs) == 1 first = run_docs[0] assert getattr(first, "id", None) == document.id def test_clean_error_is_logged_and_indexing_continues( self, db_session_with_containers: Session, mock_external_dependencies ): dataset, document, node_ids = self._create_dataset_document_with_segments(db_session_with_containers) # Force clean to raise; task should continue to indexing mock_external_dependencies["processor"].clean.side_effect = Exception("boom") document_indexing_update_task(dataset.id, document.id) # Ensure we see committed changes from another session db_session_with_containers.expire_all() # Indexing should still be triggered mock_external_dependencies["runner_instance"].run.assert_called_once() # Segments should remain (since clean failed before DB delete) remaining = db_session_with_containers.scalar( select(func.count()).select_from(DocumentSegment).where(DocumentSegment.document_id == document.id) ) assert remaining > 0 def test_document_not_found_noop(self, db_session_with_containers: Session, mock_external_dependencies): fake = Faker() # Act with non-existent document id document_indexing_update_task(dataset_id=fake.uuid4(), document_id=fake.uuid4()) # Neither processor nor runner should be called mock_external_dependencies["processor"].clean.assert_not_called() mock_external_dependencies["runner_instance"].run.assert_not_called()