From 482ec5e7129e834616fe3bbe64537305eaf90782 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 4 Aug 2026 19:53:55 +0900 Subject: [PATCH] test: use sqlite3 session in test_qa_index_processor (#38765) --- .../processor/qa_index_processor.py | 2 +- .../processor/test_qa_index_processor.py | 88 +++++++++++++------ 2 files changed, 64 insertions(+), 26 deletions(-) diff --git a/api/core/rag/index_processor/processor/qa_index_processor.py b/api/core/rag/index_processor/processor/qa_index_processor.py index 4b55d0d75a9..e1971da27fa 100644 --- a/api/core/rag/index_processor/processor/qa_index_processor.py +++ b/api/core/rag/index_processor/processor/qa_index_processor.py @@ -1,4 +1,4 @@ -"""Paragraph index processor.""" +"""Question-and-answer index processor.""" import logging import re 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 ec390fcc578..3d7cd0a98f3 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 @@ -5,12 +5,17 @@ from unittest.mock import MagicMock, Mock, patch import pandas as pd import pytest +from sqlalchemy.orm import Session from werkzeug.datastructures import FileStorage from core.entities.knowledge_entities import PreviewDetail from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.processor.qa_index_processor import QAIndexProcessor 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: @@ -32,20 +37,19 @@ class TestQAIndexProcessor: return QAIndexProcessor() @pytest.fixture - def dataset(self) -> Mock: - dataset = Mock() - dataset.id = "dataset-1" - dataset.tenant_id = "tenant-1" - dataset.indexing_technique = IndexTechniqueType.HIGH_QUALITY - dataset.is_multimodal = True - return dataset + def dataset(self) -> Dataset: + return Dataset( + id="dataset-1", + tenant_id="tenant-1", + name="QA Dataset", + created_by="user-1", + indexing_technique=IndexTechniqueType.HIGH_QUALITY, + is_multimodal=True, + ) @pytest.fixture - def dataset_document(self) -> Mock: - document = Mock() - document.id = "doc-1" - document.created_by = "user-1" - return document + def dataset_document(self) -> DatasetDocument: + return DatasetDocument(id="doc-1", created_by="user-1") @pytest.fixture def process_rule(self) -> dict: @@ -212,7 +216,7 @@ 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: Mock) -> None: + def test_load_creates_vectors_for_high_quality_dataset(self, processor: QAIndexProcessor, dataset: Dataset) -> None: session = MagicMock() docs = [Document(page_content="Q1", metadata={"answer": "A1"})] multimodal_docs = [AttachmentDocument(page_content="image", metadata={})] @@ -225,7 +229,7 @@ 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: Mock) -> None: + def test_load_skips_vector_for_non_high_quality(self, processor: QAIndexProcessor, dataset: Dataset) -> None: session = MagicMock() dataset.indexing_technique = IndexTechniqueType.ECONOMY docs = [Document(page_content="Q1", metadata={"answer": "A1"})] @@ -235,14 +239,48 @@ 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, dataset: Mock + self, + processor: QAIndexProcessor, + dataset: Dataset, + sqlite_session: Session, ) -> None: - mock_segment = SimpleNamespace(id="seg-1") - scalars_result = Mock() - scalars_result.all.return_value = [mock_segment] - mock_session = MagicMock() - mock_session.scalars.return_value = scalars_result + matching_segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc-1", + position=1, + content="Q1", + word_count=1, + tokens=1, + created_by="user-1", + index_node_id="node-1", + ) + other_node_segment = DocumentSegment( + tenant_id=dataset.tenant_id, + dataset_id=dataset.id, + document_id="doc-1", + position=2, + content="Q2", + word_count=1, + tokens=1, + created_by="user-1", + index_node_id="node-2", + ) + other_dataset_segment = DocumentSegment( + tenant_id="tenant-2", + dataset_id="dataset-2", + document_id="doc-2", + position=1, + content="Q3", + word_count=1, + tokens=1, + created_by="user-2", + index_node_id="node-1", + ) + sqlite_session.add_all([matching_segment, other_node_segment, other_dataset_segment]) + sqlite_session.commit() with ( patch( @@ -251,12 +289,12 @@ class TestQAIndexProcessor: patch("core.rag.index_processor.processor.qa_index_processor.Vector") as mock_vector_cls, ): vector = mock_vector_cls.return_value - processor.clean(dataset, ["node-1"], delete_summaries=True, session=mock_session) + processor.clean(dataset, ["node-1"], delete_summaries=True, session=sqlite_session) - mock_summary.assert_called_once_with(dataset, ["seg-1"], session=mock_session) + 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: Mock) -> None: + def test_clean_handles_dataset_wide_cleanup(self, processor: QAIndexProcessor, dataset: Dataset) -> None: session = MagicMock() with ( patch( @@ -271,7 +309,7 @@ class TestQAIndexProcessor: vector.delete.assert_called_once() def test_index_adds_documents_and_vectors_for_high_quality( - self, processor: QAIndexProcessor, dataset: Mock, dataset_document: Mock + self, processor: QAIndexProcessor, dataset: Dataset, dataset_document: DatasetDocument ) -> None: session = MagicMock() phase_events: list[str] = [] @@ -315,7 +353,7 @@ class TestQAIndexProcessor: mock_vector_cls.return_value.create.assert_called_once() def test_index_requires_high_quality( - self, processor: QAIndexProcessor, dataset: Mock, dataset_document: Mock + self, processor: QAIndexProcessor, dataset: Dataset, dataset_document: DatasetDocument ) -> None: session = MagicMock() dataset.indexing_technique = IndexTechniqueType.ECONOMY