From 96655436f45507015e76ace54b3e6d73e85eb9bf Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 10 Aug 2026 17:43:18 +0900 Subject: [PATCH] test: migrate RAG index sessions to SQLite (#40080) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../rag/docstore/test_dataset_docstore.py | 360 +++++------------- .../test_paragraph_index_processor.py | 233 ++++++------ .../test_parent_child_index_processor.py | 128 ++++--- .../processor/test_qa_index_processor.py | 66 ++-- .../core/rag/indexing/test_index_processor.py | 264 ++++++------- 5 files changed, 467 insertions(+), 584 deletions(-) diff --git a/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py b/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py index 366bbcbd501..acc1af08438 100644 --- a/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py +++ b/api/tests/unit_tests/core/rag/docstore/test_dataset_docstore.py @@ -5,7 +5,7 @@ Tests cover all public methods and error paths of the DatasetDocumentStore class which provides document storage and retrieval functionality for datasets in the RAG system. """ -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest from sqlalchemy import func, select @@ -19,6 +19,7 @@ TENANT_ID = "00000000-0000-0000-0000-000000000001" DATASET_ID = "00000000-0000-0000-0000-000000000002" DOCUMENT_ID = "00000000-0000-0000-0000-000000000003" USER_ID = "00000000-0000-0000-0000-000000000004" +ATTACHMENT_ID = "00000000-0000-0000-0000-000000000005" def _dataset() -> Dataset: @@ -69,7 +70,15 @@ def _segment( ) -class TestDatasetDocumentStoreInit: +class _UsesSQLiteSession: + session: Session + + @pytest.fixture(autouse=True) + def _inject_sqlite_session(self, sqlite_session: Session) -> None: + self.session = sqlite_session + + +class TestDatasetDocumentStoreInit(_UsesSQLiteSession): """Tests for DatasetDocumentStore initialization.""" def test_init_with_all_parameters(self): @@ -107,7 +116,7 @@ class TestDatasetDocumentStoreInit: assert store.dataset_id == "test-dataset-id" -class TestDatasetDocumentStoreSerialization: +class TestDatasetDocumentStoreSerialization(_UsesSQLiteSession): """Tests for to_dict and from_dict methods.""" def test_to_dict(self): @@ -142,23 +151,20 @@ class TestDatasetDocumentStoreSerialization: assert store._document_id == "test-doc" -class TestDatasetDocumentStoreDocs: +class TestDatasetDocumentStoreDocs(_UsesSQLiteSession): """Tests for the docs property.""" def test_docs_returns_document_dict(self): """Test that docs property returns a dictionary of documents.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_segment = _segment(index_node_id="node-1") - mock_session = MagicMock() - mock_session.scalars.return_value.all.return_value = [mock_segment] + mock_session = self.session + mock_session.add(mock_segment) + mock_session.flush() store = DatasetDocumentStore( - dataset=mock_dataset, + dataset=_dataset(), user_id="test-user-id", ) @@ -174,8 +180,7 @@ class TestDatasetDocumentStoreDocs: id="test-dataset-id", ) - mock_session = MagicMock() - mock_session.scalars.return_value.all.return_value = [] + mock_session = self.session store = DatasetDocumentStore( dataset=mock_dataset, @@ -187,12 +192,7 @@ class TestDatasetDocumentStoreDocs: assert result == {} -@pytest.mark.parametrize( - "sqlite_session", - [(DocumentSegment, ChildChunk, SegmentAttachmentBinding)], - indirect=True, -) -class TestDatasetDocumentStoreAddDocuments: +class TestDatasetDocumentStoreAddDocuments(_UsesSQLiteSession): """Tests for add_documents method.""" def test_add_documents_new_document_with_token_count(self, sqlite_session: Session): @@ -336,268 +336,147 @@ class TestDatasetDocumentStoreAddDocuments: assert sqlite_session.scalar(select(func.count()).select_from(DocumentSegment)) == 0 -class TestDatasetDocumentStoreExists: +class TestDatasetDocumentStoreExists(_UsesSQLiteSession): """Tests for document_exists method.""" def test_document_exists_returns_true(self): """Test document_exists returns True when segment exists.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) + _persist_segment(self.session) + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - mock_segment = MagicMock() - mock_session = MagicMock() - - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - result = store.document_exists("doc-1", session=mock_session) - - assert result is True + assert store.document_exists("doc-1", session=self.session) is True def test_document_exists_returns_false(self): """Test document_exists returns False when segment doesn't exist.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - result = store.document_exists("doc-1", session=mock_session) - - assert result is False + assert store.document_exists("doc-1", session=self.session) is False -class TestDatasetDocumentStoreGetDocument: +class TestDatasetDocumentStoreGetDocument(_UsesSQLiteSession): """Tests for get_document method.""" def test_get_document_success(self): """Test getting a document successfully.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) + _persist_segment(self.session, index_node_id="node-1") + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - mock_segment = _segment(index_node_id="node-1") - mock_session = MagicMock() + result = store.get_document("node-1", session=self.session, raise_error=False) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - result = store.get_document("node-1", session=mock_session, raise_error=False) - - assert isinstance(result, Document) - assert result.page_content == "Test content" + assert isinstance(result, Document) + assert result.page_content == "Test content" def test_get_document_returns_none_when_not_found(self): """Test get_document returns None when not found and raise_error=False.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) + result = store.get_document("nonexistent", session=self.session, raise_error=False) - result = store.get_document("nonexistent", session=mock_session, raise_error=False) - - assert result is None + assert result is None def test_get_document_raises_when_not_found(self): """Test get_document raises ValueError when not found and raise_error=True.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - with pytest.raises(ValueError, match="not found"): - store.get_document("nonexistent", session=mock_session, raise_error=True) + with pytest.raises(ValueError, match="not found"): + store.get_document("nonexistent", session=self.session, raise_error=True) -class TestDatasetDocumentStoreDeleteDocument: +class TestDatasetDocumentStoreDeleteDocument(_UsesSQLiteSession): """Tests for delete_document method.""" def test_delete_document_success(self): """Test deleting a document successfully.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) + segment = _persist_segment(self.session) + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - mock_segment = MagicMock() - mock_session = MagicMock() + store.delete_document("doc-1", session=self.session) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - store.delete_document("doc-1", session=mock_session) - - mock_session.delete.assert_called_with(mock_segment) - mock_session.flush.assert_called() + assert self.session.get(DocumentSegment, segment.id) is None def test_delete_document_returns_none_when_not_found(self): """Test delete_document returns None when not found and raise_error=False.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) + result = store.delete_document("nonexistent", session=self.session, raise_error=False) - result = store.delete_document("nonexistent", session=mock_session, raise_error=False) - - assert result is None + assert result is None def test_delete_document_raises_when_not_found(self): """Test delete_document raises ValueError when not found and raise_error=True.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - with pytest.raises(ValueError, match="not found"): - store.delete_document("nonexistent", session=mock_session, raise_error=True) + with pytest.raises(ValueError, match="not found"): + store.delete_document("nonexistent", session=self.session, raise_error=True) -class TestDatasetDocumentStoreHashOperations: +class TestDatasetDocumentStoreHashOperations(_UsesSQLiteSession): """Tests for set_document_hash and get_document_hash methods.""" def test_set_document_hash_success(self): """Test setting document hash successfully.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) + segment = _persist_segment(self.session, index_node_hash="old-hash") + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - mock_segment = MagicMock() - mock_segment.index_node_hash = "old-hash" - mock_session = MagicMock() + store.set_document_hash("doc-1", "new-hash", session=self.session) + self.session.expire_all() - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - store.set_document_hash("doc-1", "new-hash", session=mock_session) - - assert mock_segment.index_node_hash == "new-hash" - mock_session.flush.assert_called() + updated_segment = self.session.get(DocumentSegment, segment.id) + assert updated_segment is not None + assert updated_segment.index_node_hash == "new-hash" def test_set_document_hash_returns_none_when_not_found(self): """Test set_document_hash returns None when segment not found.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) + result = store.set_document_hash("nonexistent", "new-hash", session=self.session) - result = store.set_document_hash("nonexistent", "new-hash", session=mock_session) - - assert result is None + assert result is None def test_get_document_hash_success(self): """Test getting document hash successfully.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) + _persist_segment(self.session, index_node_hash="test-hash") + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - mock_segment = MagicMock() - mock_segment.index_node_hash = "test-hash" - mock_session = MagicMock() + result = store.get_document_hash("doc-1", session=self.session) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=mock_segment): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) - - result = store.get_document_hash("doc-1", session=mock_session) - - assert result == "test-hash" + assert result == "test-hash" def test_get_document_hash_returns_none_when_not_found(self): """Test get_document_hash returns None when segment not found.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_session = MagicMock() + store = DatasetDocumentStore(dataset=_dataset(), user_id=USER_ID) - with patch.object(DatasetDocumentStore, "get_document_segment", return_value=None): - store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - ) + result = store.get_document_hash("nonexistent", session=self.session) - result = store.get_document_hash("nonexistent", session=mock_session) - - assert result is None + assert result is None -class TestDatasetDocumentStoreSegment: +class TestDatasetDocumentStoreSegment(_UsesSQLiteSession): """Tests for get_document_segment method.""" def test_get_document_segment_returns_segment(self): """Test getting a document segment.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - mock_segment = _segment() - mock_session = MagicMock() - mock_session.scalar.return_value = mock_segment + mock_session = self.session + mock_session.add(mock_segment) + mock_session.flush() store = DatasetDocumentStore( - dataset=mock_dataset, + dataset=_dataset(), user_id="test-user-id", ) @@ -608,15 +487,10 @@ class TestDatasetDocumentStoreSegment: def test_get_document_segment_returns_none(self): """Test getting a non-existent document segment.""" - mock_dataset = Dataset( - id="test-dataset-id", - ) - - mock_session = MagicMock() - mock_session.scalar.return_value = None + mock_session = self.session store = DatasetDocumentStore( - dataset=mock_dataset, + dataset=_dataset(), user_id="test-user-id", ) @@ -625,102 +499,79 @@ class TestDatasetDocumentStoreSegment: assert result is None -class TestDatasetDocumentStoreMultimodelBinding: +class TestDatasetDocumentStoreMultimodelBinding(_UsesSQLiteSession): """Tests for add_multimodel_documents_binding method.""" def test_add_multimodel_documents_binding_with_attachments(self): """Test adding multimodel document bindings.""" - mock_dataset = Dataset( - id="test-dataset-id", - tenant_id="tenant-1", - ) + attachment = AttachmentDocument(page_content="attachment", metadata={"doc_id": ATTACHMENT_ID}) - mock_attachment = MagicMock(spec=AttachmentDocument) - mock_attachment.metadata = {"doc_id": "attachment-1"} - - mock_session = MagicMock() + mock_session = self.session store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - document_id="test-doc-id", + dataset=_dataset(), + user_id=USER_ID, + document_id=DOCUMENT_ID, ) - store.add_multimodel_documents_binding("seg-1", [mock_attachment], session=mock_session) + store.add_multimodel_documents_binding("seg-1", [attachment], session=mock_session) + mock_session.flush() - mock_session.add.assert_called() + binding = mock_session.scalar(select(SegmentAttachmentBinding)) + assert binding is not None + assert binding.segment_id == "seg-1" + assert binding.attachment_id == ATTACHMENT_ID def test_add_multimodel_documents_binding_without_attachments(self): """Test adding bindings with None attachments.""" - mock_dataset = Dataset( - id="test-dataset-id", - tenant_id="tenant-1", - ) - - mock_session = MagicMock() + mock_session = self.session store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - document_id="test-doc-id", + dataset=_dataset(), + user_id=USER_ID, + document_id=DOCUMENT_ID, ) store.add_multimodel_documents_binding("seg-1", None, session=mock_session) - mock_session.add.assert_not_called() + assert mock_session.scalar(select(func.count()).select_from(SegmentAttachmentBinding)) == 0 def test_add_multimodel_documents_binding_with_empty_list(self): """Test adding bindings with empty list.""" - mock_dataset = Dataset( - id="test-dataset-id", - tenant_id="tenant-1", - ) - - mock_session = MagicMock() + mock_session = self.session store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", - document_id="test-doc-id", + dataset=_dataset(), + user_id=USER_ID, + document_id=DOCUMENT_ID, ) store.add_multimodel_documents_binding("seg-1", [], session=mock_session) - mock_session.add.assert_not_called() + assert mock_session.scalar(select(func.count()).select_from(SegmentAttachmentBinding)) == 0 def test_add_multimodel_documents_binding_with_none_document_id(self): """Test that no bindings are added when document_id is None.""" - mock_dataset = Dataset( - id="test-dataset-id", - tenant_id="tenant-1", - ) + attachment = AttachmentDocument(page_content="attachment", metadata={"doc_id": ATTACHMENT_ID}) - mock_attachment = MagicMock(spec=AttachmentDocument) - mock_attachment.metadata = {"doc_id": "attachment-1"} - - mock_session = MagicMock() + mock_session = self.session store = DatasetDocumentStore( - dataset=mock_dataset, - user_id="test-user-id", + dataset=_dataset(), + user_id=USER_ID, document_id=None, ) - store.add_multimodel_documents_binding("seg-1", [mock_attachment], session=mock_session) + store.add_multimodel_documents_binding("seg-1", [attachment], session=mock_session) - mock_session.add.assert_not_called() + assert mock_session.scalar(select(func.count()).select_from(SegmentAttachmentBinding)) == 0 -@pytest.mark.parametrize( - "sqlite_session", - [(DocumentSegment, ChildChunk, SegmentAttachmentBinding)], - indirect=True, -) -class TestDatasetDocumentStoreAddDocumentsUpdateChild: +class TestDatasetDocumentStoreAddDocumentsUpdateChild(_UsesSQLiteSession): """Tests for add_documents when updating existing documents with children.""" def test_add_documents_update_existing_with_children(self, sqlite_session: Session): @@ -768,12 +619,7 @@ class TestDatasetDocumentStoreAddDocumentsUpdateChild: assert children[0].index_node_id == "child-1" -@pytest.mark.parametrize( - "sqlite_session", - [(DocumentSegment, ChildChunk, SegmentAttachmentBinding)], - indirect=True, -) -class TestDatasetDocumentStoreAddDocumentsUpdateAnswer: +class TestDatasetDocumentStoreAddDocumentsUpdateAnswer(_UsesSQLiteSession): """Tests for add_documents when updating existing documents with answer metadata.""" def test_add_documents_update_existing_with_answer(self, sqlite_session: Session): diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py index 6f9b0d0cc74..228aaf1a531 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py @@ -1,21 +1,36 @@ import logging from contextlib import nullcontext +from datetime import datetime from types import SimpleNamespace from typing import Any from unittest.mock import Mock, patch import pytest +from sqlalchemy import event +from sqlalchemy.orm import Session, sessionmaker from core.entities.knowledge_entities import PreviewDetail from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.processor.paragraph_index_processor import ParagraphIndexProcessor from core.rag.models.document import AttachmentDocument, Document +from extensions.storage.storage_type import StorageType from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, ImagePromptMessageContent from graphon.model_runtime.entities.model_entities import ModelFeature +from models.dataset import DocumentSegment, SegmentAttachmentBinding +from models.enums import CreatorUserRole +from models.model import UploadFile class TestParagraphIndexProcessor: + session: Session + session_factory: sessionmaker[Session] + + @pytest.fixture(autouse=True) + def _inject_sqlite_sessions(self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> None: + self.session = sqlite_session + self.session_factory = sqlite_session_factory + @pytest.fixture def processor(self) -> ParagraphIndexProcessor: return ParagraphIndexProcessor() @@ -54,9 +69,34 @@ class TestParagraphIndexProcessor: usage=LLMUsage.empty_usage(), ) + def _upload_file( + self, + *, + file_id: str, + tenant_id: str = "tenant-1", + name: str = "image.png", + extension: str = "png", + mime_type: str = "image/png", + ) -> UploadFile: + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key=f"key-{file_id}", + name=name, + size=1, + extension=extension, + mime_type=mime_type, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + created_at=datetime(2025, 1, 1), + used=True, + ) + upload_file.id = file_id + return upload_file + def test_extract_forwards_automatic_flag(self, processor: ParagraphIndexProcessor) -> None: extract_setting = Mock() - session = Mock() + session = self.session expected_docs = [Document(page_content="chunk", metadata={})] with patch( @@ -69,7 +109,7 @@ class TestParagraphIndexProcessor: mock_extract.assert_called_once_with(extract_setting=extract_setting, is_automatic=True, session=session) def test_transform_validates_process_rule(self, processor: ParagraphIndexProcessor) -> None: - session = Mock() + session = self.session with pytest.raises(ValueError, match="No process rule found"): processor.transform([Document(page_content="text", metadata={})], process_rule=None, session=session) @@ -82,7 +122,7 @@ class TestParagraphIndexProcessor: self, processor: ParagraphIndexProcessor, process_rule: dict[str, Any] ) -> None: rules_without_segmentation = SimpleNamespace(segmentation=None) - session = Mock() + session = self.session with patch( "core.rag.index_processor.processor.paragraph_index_processor.Rule.model_validate", @@ -99,7 +139,7 @@ class TestParagraphIndexProcessor: self, processor: ParagraphIndexProcessor, process_rule: dict[str, Any] ) -> None: source_document = Document(page_content="source", metadata={"dataset_id": "dataset-1", "document_id": "doc-1"}) - session = Mock() + session = self.session splitter = Mock() splitter.split_documents.return_value = [ Document(page_content=".first", metadata={}), @@ -138,7 +178,7 @@ class TestParagraphIndexProcessor: def test_transform_automatic_mode_uses_default_rules(self, processor: ParagraphIndexProcessor) -> None: splitter = Mock() splitter.split_documents.return_value = [Document(page_content="text", metadata={})] - session = Mock() + session = self.session with ( patch( @@ -173,7 +213,7 @@ class TestParagraphIndexProcessor: ) -> None: docs = [Document(page_content="chunk", metadata={})] multimodal_docs = [AttachmentDocument(page_content="image", metadata={})] - session = Mock() + session = self.session with ( patch("core.rag.index_processor.processor.paragraph_index_processor.Vector") as mock_vector_cls, @@ -191,7 +231,7 @@ class TestParagraphIndexProcessor: ) -> None: dataset.indexing_technique = IndexTechniqueType.ECONOMY docs = [Document(page_content="chunk", metadata={})] - session = Mock() + session = self.session keywords_list = [["k1"], ["k2"]] with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls: @@ -204,7 +244,7 @@ class TestParagraphIndexProcessor: ) -> None: dataset.indexing_technique = IndexTechniqueType.ECONOMY docs = [Document(page_content="chunk", metadata={})] - session = Mock() + session = self.session with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls: processor.load(dataset, docs, session=session) @@ -212,10 +252,20 @@ class TestParagraphIndexProcessor: mock_keyword_cls.return_value.add_texts.assert_called_once_with(docs, session) def test_clean_deletes_summaries_and_vector(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None: - scalars_result = Mock() - scalars_result.all.return_value = [SimpleNamespace(id="seg-1")] - session = Mock() - session.scalars.return_value = scalars_result + session = self.session + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc-1", + position=1, + content="segment", + word_count=1, + tokens=1, + created_by="user-1", + index_node_id="node-1", + ) + session.add(segment) + session.flush() with ( patch( @@ -226,14 +276,14 @@ class TestParagraphIndexProcessor: vector = mock_vector_cls.return_value processor.clean(dataset, ["node-1"], delete_summaries=True, session=session) - mock_summary.assert_called_once_with(dataset, ["seg-1"], session=session) + mock_summary.assert_called_once_with(dataset, [segment.id], session=session) vector.delete_by_ids.assert_called_once_with(["node-1"]) def test_clean_economy_deletes_summaries_and_keywords( self, processor: ParagraphIndexProcessor, dataset: Mock ) -> None: dataset.indexing_technique = IndexTechniqueType.ECONOMY - session = Mock() + session = self.session with ( patch( @@ -248,7 +298,7 @@ class TestParagraphIndexProcessor: def test_clean_deletes_keywords_by_ids(self, processor: ParagraphIndexProcessor, dataset: Mock) -> None: dataset.indexing_technique = IndexTechniqueType.ECONOMY - session = Mock() + session = self.session with patch("core.rag.index_processor.processor.paragraph_index_processor.Keyword") as mock_keyword_cls: processor.clean(dataset, ["node-2"], with_keywords=True, session=session) @@ -257,9 +307,9 @@ class TestParagraphIndexProcessor: def test_index_list_chunks_high_quality( self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock ) -> None: - session = Mock() + session = self.session phase_events: list[str] = [] - session.commit.side_effect = lambda: phase_events.append("commit") + event.listen(session, "after_commit", lambda _session: phase_events.append("commit")) with ( patch( "core.rag.index_processor.processor.paragraph_index_processor.helper.generate_text_hash", @@ -299,9 +349,9 @@ class TestParagraphIndexProcessor: self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock ) -> None: dataset.indexing_technique = IndexTechniqueType.ECONOMY - session = Mock() + session = self.session phase_events: list[str] = [] - session.commit.side_effect = lambda: phase_events.append("commit") + event.listen(session, "after_commit", lambda _session: phase_events.append("commit")) with ( patch( "core.rag.index_processor.processor.paragraph_index_processor.helper.generate_text_hash", @@ -334,8 +384,8 @@ class TestParagraphIndexProcessor: ) chunk_without_files = SimpleNamespace(content="content-2", files=None) structure = SimpleNamespace(general_chunks=[chunk_with_files, chunk_without_files]) - session = Mock() - account_session = Mock() + session = self.session + account_session = self.session_factory() with ( patch( @@ -374,7 +424,7 @@ class TestParagraphIndexProcessor: self, processor: ParagraphIndexProcessor, dataset: Mock, dataset_document: Mock ) -> None: structure = SimpleNamespace(general_chunks=[SimpleNamespace(content="content", files=None)]) - session = Mock() + session = self.session with ( patch( @@ -403,8 +453,8 @@ class TestParagraphIndexProcessor: def test_generate_summary_preview_success_and_failure(self, processor: ParagraphIndexProcessor) -> None: preview_items = [PreviewDetail(content="chunk-1"), PreviewDetail(content="chunk-2")] - session = Mock() - worker_sessions = [Mock(), Mock()] + session = self.session + worker_sessions = [self.session_factory(), self.session_factory()] with ( patch( @@ -440,7 +490,9 @@ class TestParagraphIndexProcessor: patch("flask.current_app", fake_current_app), patch.object(processor, "generate_summary", return_value=("summary", LLMUsage.empty_usage())), ): - result = processor.generate_summary_preview("tenant-1", preview_items, {"enable": True}, session=Mock()) + result = processor.generate_summary_preview( + "tenant-1", preview_items, {"enable": True}, session=self.session + ) assert result[0].summary == "summary" @@ -456,16 +508,16 @@ class TestParagraphIndexProcessor: patch("concurrent.futures.wait", side_effect=[(set(), {future}), (set(), set())]), ): with pytest.raises(ValueError, match="timeout"): - processor.generate_summary_preview("tenant-1", preview_items, {"enable": True}, session=Mock()) + processor.generate_summary_preview("tenant-1", preview_items, {"enable": True}, session=self.session) future.cancel.assert_called_once() def test_generate_summary_validates_input(self) -> None: with pytest.raises(ValueError, match="must be enabled"): - ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": False}, session=Mock()) + ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": False}, session=self.session) with pytest.raises(ValueError, match="model_name and model_provider_name"): - ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": True}, session=Mock()) + ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": True}, session=self.session) def test_generate_summary_text_only_flow(self, caplog: pytest.LogCaptureFixture) -> None: model_instance = Mock() @@ -495,7 +547,7 @@ class TestParagraphIndexProcessor: "text content", {"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"}, document_language="English", - session=Mock(), + session=self.session, ) assert summary == "text summary" @@ -537,7 +589,7 @@ class TestParagraphIndexProcessor: "text content", {"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"}, segment_id="seg-1", - session=Mock(), + session=self.session, ) assert summary == "vision summary" @@ -580,7 +632,7 @@ class TestParagraphIndexProcessor: "tenant-1", "text content", {"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"}, - session=Mock(), + session=self.session, ) assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1 @@ -594,30 +646,19 @@ class TestParagraphIndexProcessor: "![img2](/files/22222222-2222-2222-2222-222222222222/file-preview) " "![tool](/files/tools/33333333-3333-3333-3333-333333333333.png)" ) - image_upload = SimpleNamespace( - id="11111111-1111-1111-1111-111111111111", - tenant_id="tenant-1", - name="image.png", - mime_type="image/png", - extension="png", - source_url="", - size=1, - key="key", + session = self.session + session.add_all( + [ + self._upload_file(file_id="11111111-1111-1111-1111-111111111111"), + self._upload_file( + file_id="22222222-2222-2222-2222-222222222222", + name="file.txt", + extension="txt", + mime_type="text/plain", + ), + ] ) - non_image_upload = SimpleNamespace( - id="22222222-2222-2222-2222-222222222222", - tenant_id="tenant-1", - name="file.txt", - mime_type="text/plain", - extension="txt", - source_url="", - size=1, - key="key", - ) - scalars_result = Mock() - scalars_result.all.return_value = [image_upload, non_image_upload] - session = Mock() - session.scalars.return_value = scalars_result + session.flush() with ( patch( @@ -633,28 +674,14 @@ class TestParagraphIndexProcessor: assert not any(record.levelno == logging.WARNING for record in caplog.records) def test_extract_images_from_text_returns_empty_when_no_matches(self) -> None: - scalars_result = Mock() - scalars_result.all.return_value = [] - session = Mock() - session.scalars.return_value = scalars_result + session = self.session assert ParagraphIndexProcessor._extract_images_from_text("tenant-1", "no images here", session) == [] def test_extract_images_from_text_logs_when_build_fails(self, caplog: pytest.LogCaptureFixture) -> None: text = "![img](/files/11111111-1111-1111-1111-111111111111/image-preview)" - image_upload = SimpleNamespace( - id="11111111-1111-1111-1111-111111111111", - tenant_id="tenant-1", - name="image.png", - mime_type="image/png", - extension="png", - source_url="", - size=1, - key="key", - ) - scalars_result = Mock() - scalars_result.all.return_value = [image_upload] - session = Mock() - session.scalars.return_value = scalars_result + session = self.session + session.add(self._upload_file(file_id="11111111-1111-1111-1111-111111111111")) + session.flush() with ( patch( @@ -669,49 +696,39 @@ class TestParagraphIndexProcessor: assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1 def test_extract_images_from_segment_attachments(self, caplog: pytest.LogCaptureFixture) -> None: - image_upload = SimpleNamespace( - id="file-1", - name="image", - extension="png", - mime_type="image/png", - source_url="", - size=1, - key="k1", - ) - bad_upload = SimpleNamespace( - id="file-2", - name="broken", - extension=None, - mime_type="image/png", - source_url="", - size=1, - key="k2", - ) - non_image_upload = SimpleNamespace( - id="file-3", - name="text", - extension="txt", - mime_type="text/plain", - source_url="", - size=1, - key="k3", - ) - execute_result = Mock() - execute_result.all.return_value = [(None, image_upload), (None, bad_upload), (None, non_image_upload)] - session = Mock() - session.execute.return_value = execute_result + session = self.session + uploads = [ + self._upload_file(file_id="file-1", name="image"), + self._upload_file(file_id="file-2", name="broken"), + self._upload_file(file_id="file-3", name="text", extension="txt", mime_type="text/plain"), + ] + bindings = [ + SegmentAttachmentBinding( + tenant_id="tenant-1", + dataset_id="dataset-1", + document_id="doc-1", + segment_id="seg-1", + attachment_id=upload.id, + ) + for upload in uploads + ] + session.add_all([*uploads, *bindings]) + session.flush() - with caplog.at_level(logging.WARNING, logger="core.rag.index_processor.processor.paragraph_index_processor"): + with ( + patch( + "core.rag.index_processor.processor.paragraph_index_processor.File", + side_effect=[SimpleNamespace(id="file-1"), RuntimeError("bad file")], + ), + caplog.at_level(logging.WARNING, logger="core.rag.index_processor.processor.paragraph_index_processor"), + ): files = ParagraphIndexProcessor._extract_images_from_segment_attachments("tenant-1", "seg-1", session) assert len(files) == 1 assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1 def test_extract_images_from_segment_attachments_empty(self) -> None: - execute_result = Mock() - execute_result.all.return_value = [] - session = Mock() - session.execute.return_value = execute_result + session = self.session empty_files = ParagraphIndexProcessor._extract_images_from_segment_attachments("tenant-1", "seg-1", session) diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py index a2a5c67972d..1a7e54e35c1 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_parent_child_index_processor.py @@ -3,15 +3,26 @@ from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch import pytest +from sqlalchemy import event +from sqlalchemy.orm import Session, sessionmaker from core.entities.knowledge_entities import PreviewDetail from core.rag.entities import ParentMode, Rule, Segmentation from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.processor.parent_child_index_processor import ParentChildIndexProcessor from core.rag.models.document import AttachmentDocument, ChildDocument, Document +from models.dataset import ChildChunk, DatasetProcessRule, DocumentSegment class TestParentChildIndexProcessor: + session: Session + session_factory: sessionmaker[Session] + + @pytest.fixture(autouse=True) + def _inject_sqlite_sessions(self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> None: + self.session = sqlite_session + self.session_factory = sqlite_session_factory + @pytest.fixture def processor(self) -> ParentChildIndexProcessor: return ParentChildIndexProcessor() @@ -50,7 +61,7 @@ class TestParentChildIndexProcessor: def test_extract_forwards_automatic_flag(self, processor: ParentChildIndexProcessor) -> None: extract_setting = Mock() - session = Mock() + session = self.session expected = [Document(page_content="chunk", metadata={})] with patch( @@ -63,7 +74,7 @@ class TestParentChildIndexProcessor: mock_extract.assert_called_once_with(extract_setting=extract_setting, is_automatic=True, session=session) def test_transform_validates_process_rule(self, processor: ParentChildIndexProcessor) -> None: - session = MagicMock() + session = self.session with pytest.raises(ValueError, match="No process rule found"): processor.transform([Document(page_content="text", metadata={})], process_rule=None, session=session) @@ -82,7 +93,7 @@ class TestParentChildIndexProcessor: processor.transform( [Document(page_content="text", metadata={})], process_rule={"mode": "custom", "rules": {"enabled": True}}, - session=MagicMock(), + session=self.session, ) def test_transform_paragraph_builds_parent_and_child_docs(self, processor: ParentChildIndexProcessor) -> None: @@ -117,7 +128,7 @@ class TestParentChildIndexProcessor: [parent_document], process_rule={"mode": "custom", "rules": {"enabled": True}}, preview=False, - session=MagicMock(), + session=self.session, ) assert len(result) == 1 @@ -154,7 +165,7 @@ class TestParentChildIndexProcessor: documents, process_rule={"mode": "custom", "rules": {"enabled": True}}, preview=True, - session=MagicMock(), + session=self.session, ) assert len(result) == 10 @@ -188,7 +199,7 @@ class TestParentChildIndexProcessor: docs, process_rule={"mode": "hierarchical", "rules": {"enabled": True}}, preview=True, - session=MagicMock(), + session=self.session, ) assert len(result) == 1 @@ -205,7 +216,7 @@ class TestParentChildIndexProcessor: ], ) multimodal_docs = [AttachmentDocument(page_content="image", metadata={})] - session = MagicMock() + session = self.session with patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls: vector = mock_vector_cls.return_value @@ -219,7 +230,7 @@ class TestParentChildIndexProcessor: vector.create_multimodal.assert_called_once_with(multimodal_docs) def test_clean_with_precomputed_child_ids(self, processor: ParentChildIndexProcessor, dataset: Mock) -> None: - session = MagicMock() + session = self.session with ( patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls, @@ -234,16 +245,42 @@ class TestParentChildIndexProcessor: ) vector.delete_by_ids.assert_called_once_with(["child-1", "child-2"]) - session.execute.assert_called() - session.flush.assert_called_once() + assert session.query(ChildChunk).count() == 0 def test_clean_queries_child_ids_when_not_precomputed( self, processor: ParentChildIndexProcessor, dataset: Mock ) -> None: - execute_result = Mock() - execute_result.all.return_value = [("child-1",), (None,), ("child-2",)] - session = MagicMock() - session.execute.return_value = execute_result + session = self.session + parent = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc-1", + position=1, + content="parent", + word_count=1, + tokens=1, + created_by="user-1", + index_node_id="node-1", + ) + session.add(parent) + session.flush() + session.add_all( + [ + ChildChunk( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc-1", + segment_id=parent.id, + position=index, + content=f"child-{index}", + word_count=1, + created_by="user-1", + index_node_id=f"child-{index}", + ) + for index in (1, 2) + ] + ) + session.flush() with ( patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls, @@ -254,7 +291,7 @@ class TestParentChildIndexProcessor: vector.delete_by_ids.assert_called_once_with(["child-1", "child-2"]) def test_clean_dataset_wide_cleanup(self, processor: ParentChildIndexProcessor, dataset: Mock) -> None: - session = MagicMock() + session = self.session with ( patch("core.rag.index_processor.processor.parent_child_index_processor.Vector") as mock_vector_cls, @@ -263,17 +300,23 @@ class TestParentChildIndexProcessor: processor.clean(dataset, None, delete_child_chunks=True, session=session) vector.delete.assert_called_once() - session.execute.assert_called() - session.flush.assert_called_once() + assert session.query(ChildChunk).count() == 0 def test_clean_deletes_summaries_when_requested(self, processor: ParentChildIndexProcessor, dataset: Mock) -> None: - scalars_result = Mock() - scalars_result.all.return_value = [SimpleNamespace(id="seg-1")] - session = MagicMock() - session.scalars.return_value = scalars_result - session_ctx = MagicMock() - session_ctx.__enter__.return_value = session - session_ctx.__exit__.return_value = False + session = self.session + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc-1", + position=1, + content="parent", + word_count=1, + tokens=1, + created_by="user-1", + index_node_id="node-1", + ) + session.add(segment) + session.flush() with ( patch( @@ -283,7 +326,7 @@ class TestParentChildIndexProcessor: ): processor.clean(dataset, ["node-1"], delete_summaries=True, precomputed_child_node_ids=[], session=session) - mock_summary.assert_called_once_with(dataset, ["seg-1"], session=session) + mock_summary.assert_called_once_with(dataset, [segment.id], session=session) def test_clean_deletes_all_summaries_when_node_ids_missing( self, processor: ParentChildIndexProcessor, dataset: Mock @@ -294,7 +337,7 @@ class TestParentChildIndexProcessor: ) as mock_summary, patch("core.rag.index_processor.processor.parent_child_index_processor.Vector"), ): - session = MagicMock() + session = self.session processor.clean(dataset, None, delete_summaries=True, session=session) mock_summary.assert_called_once_with(dataset, None, session=session) @@ -341,20 +384,15 @@ class TestParentChildIndexProcessor: ) ], ) - dataset_rule = SimpleNamespace(id="rule-1") - session = MagicMock() + session = self.session phase_events: list[str] = [] - session.commit.side_effect = lambda: phase_events.append("commit") + event.listen(session, "after_commit", lambda _session: phase_events.append("commit")) with ( patch( "core.rag.index_processor.processor.parent_child_index_processor.ParentChildStructureChunk.model_validate", return_value=parent_childs, ), - patch( - "core.rag.index_processor.processor.parent_child_index_processor.DatasetProcessRule", - return_value=dataset_rule, - ), patch( "core.rag.index_processor.processor.parent_child_index_processor.helper.generate_text_hash", side_effect=lambda text: f"hash-{text}", @@ -373,9 +411,8 @@ class TestParentChildIndexProcessor: processor.index(dataset, dataset_document, {"parent_child_chunks": []}, session) assert phase_events == ["count", "store", "commit", "vector"] - assert dataset_document.dataset_process_rule_id == "rule-1" - session.add.assert_called_once_with(dataset_rule) - session.flush.assert_called_once() + dataset_rule = session.get(DatasetProcessRule, dataset_document.dataset_process_rule_id) + assert dataset_rule is not None documents = mock_token_counter.call_args.kwargs["documents"] assert [document.page_content for document in documents] == ["parent text"] mock_token_counter.assert_called_once_with(dataset=dataset, documents=documents) @@ -396,19 +433,14 @@ class TestParentChildIndexProcessor: parent_mode=ParentMode.PARAGRAPH, parent_child_chunks=[SimpleNamespace(parent_content="parent", child_contents=["child"], files=None)], ) - dataset_rule = SimpleNamespace(id="rule-1") - session = MagicMock() - account_session = MagicMock() + session = self.session + account_session = self.session_factory() with ( patch( "core.rag.index_processor.processor.parent_child_index_processor.ParentChildStructureChunk.model_validate", return_value=parent_childs, ), - patch( - "core.rag.index_processor.processor.parent_child_index_processor.DatasetProcessRule", - return_value=dataset_rule, - ), patch( "core.rag.index_processor.processor.parent_child_index_processor.helper.generate_text_hash", return_value="hash", @@ -480,8 +512,8 @@ class TestParentChildIndexProcessor: def test_generate_summary_preview_sets_summaries(self, processor: ParentChildIndexProcessor) -> None: preview_texts = [PreviewDetail(content="chunk-1"), PreviewDetail(content="chunk-2")] - session = MagicMock() - worker_sessions = [MagicMock(), MagicMock()] + session = self.session + worker_sessions = [self.session_factory(), self.session_factory()] with ( patch( @@ -513,7 +545,7 @@ class TestParentChildIndexProcessor: side_effect=RuntimeError("summary failed"), ): with pytest.raises(ValueError, match="Failed to generate summaries"): - processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=MagicMock()) + processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=self.session) def test_generate_summary_preview_falls_back_without_flask_context( self, processor: ParentChildIndexProcessor @@ -529,7 +561,7 @@ class TestParentChildIndexProcessor: ), ): result = processor.generate_summary_preview( - "tenant-1", preview_texts, {"enable": True}, session=MagicMock() + "tenant-1", preview_texts, {"enable": True}, session=self.session ) assert result[0].summary == "summary" @@ -546,6 +578,6 @@ class TestParentChildIndexProcessor: patch("concurrent.futures.wait", side_effect=[(set(), {future}), (set(), set())]), ): with pytest.raises(ValueError, match="timeout"): - processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=MagicMock()) + processor.generate_summary_preview("tenant-1", preview_texts, {"enable": True}, session=self.session) future.cancel.assert_called_once() diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py index 3d7cd0a98f3..3bcc1c59a17 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_qa_index_processor.py @@ -1,10 +1,11 @@ import logging from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import Mock, patch import pandas as pd import pytest +from sqlalchemy import event from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage @@ -15,8 +16,6 @@ from core.rag.models.document import AttachmentDocument, Document from models.dataset import Dataset, DocumentSegment from models.dataset import Document as DatasetDocument -TABLES = (DocumentSegment,) - class _ImmediateThread: def __init__(self, target, args=(), kwargs=None): @@ -62,9 +61,9 @@ class TestQAIndexProcessor: segmentation = SimpleNamespace(max_tokens=256, chunk_overlap=10, separator="\n") return SimpleNamespace(segmentation=segmentation) - def test_extract_forwards_automatic_flag(self, processor: QAIndexProcessor) -> None: + def test_extract_forwards_automatic_flag(self, processor: QAIndexProcessor, sqlite_session: Session) -> None: extract_setting = Mock() - session = Mock() + session = sqlite_session expected_docs = [Document(page_content="chunk", metadata={})] with patch("core.rag.index_processor.processor.qa_index_processor.ExtractProcessor.extract") as mock_extract: @@ -75,22 +74,22 @@ class TestQAIndexProcessor: assert docs == expected_docs mock_extract.assert_called_once_with(extract_setting=extract_setting, is_automatic=True, session=session) - def test_transform_rejects_none_process_rule(self, processor: QAIndexProcessor) -> None: - session = MagicMock() + def test_transform_rejects_none_process_rule(self, processor: QAIndexProcessor, sqlite_session: Session) -> None: + session = sqlite_session with pytest.raises(ValueError, match="No process rule found"): processor.transform([Document(page_content="text", metadata={})], process_rule=None, session=session) - def test_transform_rejects_missing_rules_key(self, processor: QAIndexProcessor) -> None: - session = MagicMock() + def test_transform_rejects_missing_rules_key(self, processor: QAIndexProcessor, sqlite_session: Session) -> None: + session = sqlite_session with pytest.raises(ValueError, match="No rules found in process rule"): processor.transform( [Document(page_content="text", metadata={})], process_rule={"mode": "custom"}, session=session ) def test_transform_preview_calls_formatter_once( - self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app + self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app, sqlite_session: Session ) -> None: - session = MagicMock() + session = sqlite_session document = Document(page_content="raw text", metadata={"dataset_id": "dataset-1", "document_id": "doc-1"}) split_node = Document(page_content=".question", metadata={}) splitter = Mock() @@ -132,9 +131,9 @@ class TestQAIndexProcessor: mock_format.assert_called_once() def test_transform_non_preview_uses_thread_batches( - self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app + self, processor: QAIndexProcessor, process_rule: dict[str, Any], fake_flask_app, sqlite_session: Session ) -> None: - session = MagicMock() + session = sqlite_session documents = [ Document(page_content="doc-1", metadata={"document_id": "doc-1", "dataset_id": "dataset-1"}), Document(page_content="doc-2", metadata={"document_id": "doc-2", "dataset_id": "dataset-1"}), @@ -216,8 +215,10 @@ class TestQAIndexProcessor: with pytest.raises(ValueError, match="bad csv"): processor.format_by_template(csv_file) - def test_load_creates_vectors_for_high_quality_dataset(self, processor: QAIndexProcessor, dataset: Dataset) -> None: - session = MagicMock() + def test_load_creates_vectors_for_high_quality_dataset( + self, processor: QAIndexProcessor, dataset: Dataset, sqlite_session: Session + ) -> None: + session = sqlite_session docs = [Document(page_content="Q1", metadata={"answer": "A1"})] multimodal_docs = [AttachmentDocument(page_content="image", metadata={})] @@ -229,8 +230,10 @@ class TestQAIndexProcessor: vector.create.assert_called_once_with(docs) vector.create_multimodal.assert_called_once_with(multimodal_docs) - def test_load_skips_vector_for_non_high_quality(self, processor: QAIndexProcessor, dataset: Dataset) -> None: - session = MagicMock() + def test_load_skips_vector_for_non_high_quality( + self, processor: QAIndexProcessor, dataset: Dataset, sqlite_session: Session + ) -> None: + session = sqlite_session dataset.indexing_technique = IndexTechniqueType.ECONOMY docs = [Document(page_content="Q1", metadata={"answer": "A1"})] @@ -239,7 +242,6 @@ class TestQAIndexProcessor: mock_vector_cls.assert_not_called() - @pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) def test_clean_handles_summary_deletion_and_vector_cleanup( self, processor: QAIndexProcessor, @@ -294,8 +296,10 @@ class TestQAIndexProcessor: mock_summary.assert_called_once_with(dataset, [matching_segment.id], session=sqlite_session) vector.delete_by_ids.assert_called_once_with(["node-1"]) - def test_clean_handles_dataset_wide_cleanup(self, processor: QAIndexProcessor, dataset: Dataset) -> None: - session = MagicMock() + def test_clean_handles_dataset_wide_cleanup( + self, processor: QAIndexProcessor, dataset: Dataset, sqlite_session: Session + ) -> None: + session = sqlite_session with ( patch( "core.rag.index_processor.processor.qa_index_processor.SummaryIndexService.delete_summaries_for_segments" @@ -309,11 +313,15 @@ class TestQAIndexProcessor: vector.delete.assert_called_once() def test_index_adds_documents_and_vectors_for_high_quality( - self, processor: QAIndexProcessor, dataset: Dataset, dataset_document: DatasetDocument + self, + processor: QAIndexProcessor, + dataset: Dataset, + dataset_document: DatasetDocument, + sqlite_session: Session, ) -> None: - session = MagicMock() + session = sqlite_session phase_events: list[str] = [] - session.commit.side_effect = lambda: phase_events.append("commit") + event.listen(session, "after_commit", lambda _session: phase_events.append("commit")) qa_chunks = SimpleNamespace( qa_chunks=[ SimpleNamespace(question="Q1", answer="A1"), @@ -353,9 +361,13 @@ class TestQAIndexProcessor: mock_vector_cls.return_value.create.assert_called_once() def test_index_requires_high_quality( - self, processor: QAIndexProcessor, dataset: Dataset, dataset_document: DatasetDocument + self, + processor: QAIndexProcessor, + dataset: Dataset, + dataset_document: DatasetDocument, + sqlite_session: Session, ) -> None: - session = MagicMock() + session = sqlite_session dataset.indexing_technique = IndexTechniqueType.ECONOMY qa_chunks = SimpleNamespace(qa_chunks=[SimpleNamespace(question="Q1", answer="A1")]) @@ -389,10 +401,10 @@ class TestQAIndexProcessor: assert preview["total_segments"] == 1 assert preview["qa_preview"] == [{"question": "Q1", "answer": "A1"}] - def test_generate_summary_preview_returns_input(self, processor: QAIndexProcessor) -> None: + def test_generate_summary_preview_returns_input(self, processor: QAIndexProcessor, sqlite_session: Session) -> None: preview_items = [PreviewDetail(content="Q1")] assert ( - processor.generate_summary_preview("tenant-1", preview_items, {"enable": False}, session=MagicMock()) + processor.generate_summary_preview("tenant-1", preview_items, {"enable": False}, session=sqlite_session) is preview_items ) diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py b/api/tests/unit_tests/core/rag/indexing/test_index_processor.py index 11137e74380..51d817d8f53 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/test_index_processor.py @@ -1,12 +1,51 @@ -import datetime +import uuid from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, patch +from sqlalchemy import event +from sqlalchemy.orm import Session, sessionmaker + from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.index_processor import IndexProcessor from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem -from models.dataset import Dataset, Document +from models.dataset import Dataset, Document, DocumentSegment +from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus + + +def _persist_dataset_and_document( + session: Session, + *, + indexing_technique: IndexTechniqueType = IndexTechniqueType.HIGH_QUALITY, + summary_index_setting: dict | None = None, +) -> tuple[Dataset, Document]: + tenant_id = str(uuid.uuid4()) + created_by = str(uuid.uuid4()) + dataset = Dataset( + id=str(uuid.uuid4()), + tenant_id=tenant_id, + name="Dataset", + data_source_type=DataSourceType.UPLOAD_FILE, + indexing_technique=indexing_technique, + chunk_structure="text_model", + summary_index_setting=summary_index_setting, + created_by=created_by, + ) + document = Document( + id=str(uuid.uuid4()), + tenant_id=tenant_id, + dataset_id=dataset.id, + position=1, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="Document", + created_from=DocumentCreatedFrom.WEB, + created_by=created_by, + doc_language="English", + ) + session.add_all([dataset, document]) + session.flush() + return dataset, document class TestIndexProcessor: @@ -22,28 +61,10 @@ class TestIndexProcessor: assert preview.qa_preview[0].question == "Q1" assert preview.qa_preview[0].answer == "A1" - def test_index_and_clean_ends_transactions_around_index_io(self) -> None: - document = SimpleNamespace( - id="document-1", - name="Document", - created_at=datetime.datetime(2026, 1, 1), - indexing_latency=None, - indexing_status=None, - completed_at=None, - word_count=0, - need_summary=False, - ) - dataset = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - name="Dataset", - chunk_structure="text_model", - summary_index_setting=None, - ) - session = MagicMock() - session.scalar.side_effect = [dataset, document, 3] + def test_index_and_clean_ends_transactions_around_index_io(self, sqlite_session: Session) -> None: + dataset, document = _persist_dataset_and_document(sqlite_session) phase_events: list[str] = [] - session.commit.side_effect = lambda: phase_events.append("commit") + event.listen(sqlite_session, "after_commit", lambda _session: phase_events.append("commit")) index_processor = MagicMock() index_processor.index.side_effect = lambda *args: phase_events.append("index") @@ -65,7 +86,7 @@ class TestIndexProcessor: original_document_id="", chunks=chunks, batch="batch-1", - session=session, + session=sqlite_session, ) assert phase_events == ["commit", "index", "commit"] @@ -75,30 +96,13 @@ class TestIndexProcessor: chunk_structure=dataset.chunk_structure, chunks=chunks, include_summaries=False, - session=session, + session=sqlite_session, ) - def test_index_and_clean_skips_admission_for_replacement_without_existing_vector_points(self) -> None: - document = SimpleNamespace( - id="document-1", - name="Document", - created_at=datetime.datetime(2026, 1, 1), - indexing_latency=None, - indexing_status=None, - completed_at=None, - word_count=0, - need_summary=False, - ) - dataset = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - name="Dataset", - chunk_structure="text_model", - summary_index_setting=None, - ) - session = MagicMock() - session.scalar.side_effect = [dataset, document, 3] - session.scalars.return_value.all.return_value = [] + def test_index_and_clean_skips_admission_for_replacement_without_existing_vector_points( + self, sqlite_session: Session + ) -> None: + dataset, document = _persist_dataset_and_document(sqlite_session) index_processor = MagicMock() processor = IndexProcessor() chunks = {"general_chunks": ["content"]} @@ -114,39 +118,50 @@ class TestIndexProcessor: original_document_id=document.id, chunks=chunks, batch="batch-1", - session=session, + session=sqlite_session, ) admission_service_class.assert_not_called() - def test_index_and_clean_scopes_replacement_queries_to_dataset_owner(self) -> None: - dataset = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - name="Dataset", - summary_index_setting=None, - chunk_structure="text_model", + def test_index_and_clean_scopes_replacement_queries_to_dataset_owner(self, sqlite_session: Session) -> None: + dataset, document = _persist_dataset_and_document(sqlite_session) + original_document = Document( + id=str(uuid.uuid4()), + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + position=2, + data_source_type=DataSourceType.UPLOAD_FILE, + batch="batch-1", + name="Original document", + created_from=DocumentCreatedFrom.WEB, + created_by=dataset.created_by, ) - document = SimpleNamespace( - id="doc-1", - tenant_id="tenant-1", - dataset_id="dataset-1", - name="Document", - created_at=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id=original_document.id, + position=1, + content="old content", + word_count=3, + tokens=3, + created_by=dataset.created_by, + index_node_id="node-1", + status=SegmentStatus.COMPLETED, ) - segment = SimpleNamespace(index_node_id="node-1") - session = MagicMock() - - def resolve_owner(statement): - entity = statement.column_descriptions[0]["entity"] - if entity is Dataset: - return dataset - if entity is Document: - return document - return 3 - - session.scalar.side_effect = resolve_owner - session.scalars.return_value.all.return_value = [segment] + control_segment = DocumentSegment( + tenant_id=str(uuid.uuid4()), + dataset_id=str(uuid.uuid4()), + document_id=original_document.id, + position=1, + content="other tenant content", + word_count=4, + tokens=4, + created_by=str(uuid.uuid4()), + index_node_id="other-node", + status=SegmentStatus.COMPLETED, + ) + sqlite_session.add_all([original_document, segment, control_segment]) + sqlite_session.flush() processor = IndexProcessor() with ( @@ -155,91 +170,54 @@ class TestIndexProcessor: ): index_backend = index_processor_factory.return_value.init_index_processor.return_value processor.index_and_clean( - dataset_id="dataset-1", - document_id="doc-1", - original_document_id="original-doc", + dataset_id=dataset.id, + document_id=document.id, + original_document_id=original_document.id, chunks={}, batch="batch-1", - session=session, + session=sqlite_session, ) - document_statement = next( - call.args[0] - for call in session.scalar.call_args_list - if call.args[0].column_descriptions[0]["entity"] is Document - ) - segment_statement = session.scalars.call_args.args[0] - delete_statement = session.execute.call_args_list[0].args[0] - word_count_statement = session.scalar.call_args_list[-1].args[0] - segment_update_statement = session.execute.call_args_list[1].args[0] - document_owner = {"doc-1", "dataset-1", "tenant-1"} - original_document_owner = {"original-doc", "dataset-1", "tenant-1"} - - assert document_owner <= set(document_statement.compile().params.values()) - assert original_document_owner <= set(segment_statement.compile().params.values()) - assert original_document_owner <= set(delete_statement.compile().params.values()) - assert document_owner <= set(word_count_statement.compile().params.values()) - assert document_owner <= set(segment_update_statement.compile().params.values()) + assert sqlite_session.get(DocumentSegment, segment.id) is None + assert sqlite_session.get(DocumentSegment, control_segment.id) is control_segment + assert sqlite_session.get(Document, document.id).indexing_status == "completed" index_backend.clean.assert_called_once_with( dataset, ["node-1"], with_keywords=True, delete_child_chunks=True, - session=session, + session=sqlite_session, ) - index_backend.index.assert_called_once_with(dataset, document, {}, session) + index_backend.index.assert_called_once_with(dataset, document, {}, sqlite_session) admission_service_class.assert_not_called() - def test_get_preview_output_scopes_document_to_dataset_owner(self) -> None: - dataset = SimpleNamespace( - id="dataset-1", - tenant_id="tenant-1", - indexing_technique=IndexTechniqueType.ECONOMY, - summary_index_setting=None, - ) - document = SimpleNamespace(doc_language="English") - session = MagicMock() - - def resolve_owner(statement): - entity = statement.column_descriptions[0]["entity"] - if entity is Dataset: - return dataset - if entity is Document: - return document - raise AssertionError(f"Unexpected entity: {entity}") - - session.scalar.side_effect = resolve_owner + def test_get_preview_output_scopes_document_to_dataset_owner(self, sqlite_session: Session) -> None: + dataset, document = _persist_dataset_and_document(sqlite_session, indexing_technique=IndexTechniqueType.ECONOMY) processor = IndexProcessor() expected_preview = MagicMock() with patch.object(processor, "format_preview", return_value=expected_preview): result = processor.get_preview_output( chunks={}, - dataset_id="dataset-1", - document_id="doc-1", + dataset_id=dataset.id, + document_id=document.id, chunk_structure="text_model", summary_index_setting=None, - session=session, + session=sqlite_session, ) - document_statement = next( - call.args[0] - for call in session.scalar.call_args_list - if call.args[0].column_descriptions[0]["entity"] is Document - ) - assert {"doc-1", "dataset-1", "tenant-1"} <= set(document_statement.compile().params.values()) assert result is expected_preview - def test_preview_summary_workers_use_independent_sessions(self) -> None: - caller_session = MagicMock() - phase_events: list[str] = [] - caller_session.commit.side_effect = lambda: phase_events.append("commit") - caller_session.scalar.return_value = SimpleNamespace( - indexing_technique=IndexTechniqueType.HIGH_QUALITY, + def test_preview_summary_workers_use_independent_sessions( + self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] + ) -> None: + dataset, _ = _persist_dataset_and_document( + sqlite_session, summary_index_setting={"enable": True}, - tenant_id="tenant-1", ) - worker_sessions = [MagicMock(), MagicMock()] + phase_events: list[str] = [] + event.listen(sqlite_session, "after_commit", lambda _session: phase_events.append("commit")) + worker_sessions: list[Session] = [] preview = Preview( chunk_structure="text_model", total_segments=2, @@ -247,11 +225,11 @@ class TestIndexProcessor: ) flask_app = SimpleNamespace(app_context=lambda: nullcontext()) processor = IndexProcessor() - worker_contexts = iter(nullcontext(worker_session) for worker_session in worker_sessions) - def create_worker_session(): + def generate_summary(*_args, **kwargs): phase_events.append("worker") - return next(worker_contexts) + worker_sessions.append(kwargs["session"]) + return "summary", None with ( patch.object(processor, "format_preview", return_value=preview), @@ -260,27 +238,25 @@ class TestIndexProcessor: SimpleNamespace(_get_current_object=lambda: flask_app), ), patch( - "core.rag.index_processor.index_processor.session_factory.create_session", - side_effect=create_worker_session, + "core.rag.index_processor.index_processor.session_factory", + SimpleNamespace(create_session=sqlite_session_factory), ), patch( "core.rag.index_processor.index_processor.ParagraphIndexProcessor.generate_summary", - return_value=("summary", None), + side_effect=generate_summary, ) as generate_summary, ): result = processor.get_preview_output( chunks=[], - dataset_id="dataset-1", + dataset_id=dataset.id, document_id="", chunk_structure="text_model", summary_index_setting={"enable": True}, - session=caller_session, + session=sqlite_session, ) assert all(item.summary == "summary" for item in result.preview) assert phase_events == ["commit", "worker", "worker"] call_sessions = [call.kwargs["session"] for call in generate_summary.call_args_list] - assert all(call_session is not caller_session for call_session in call_sessions) - assert all( - any(call_session is worker_session for worker_session in worker_sessions) for call_session in call_sessions - ) + assert all(call_session is not sqlite_session for call_session in call_sessions) + assert call_sessions == worker_sessions