mirror of
https://github.com/langgenius/dify.git
synced 2026-09-29 17:07:38 +08:00
test: migrate document indexing task sessions to SQLite (#40092)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Asuka Minato
co-authored by
autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
parent
7039217c99
commit
6530f2bcdb
@@ -1,520 +1,276 @@
|
||||
"""
|
||||
Unit tests for document_indexing_update_task summary generation.
|
||||
"""SQLite-backed tests for document update indexing and summary generation."""
|
||||
|
||||
After updating a document via the API, the summary index should be
|
||||
regenerated under the same conditions as during initial creation:
|
||||
- indexing_technique is HIGH_QUALITY
|
||||
- summary_index_setting has enable=True
|
||||
- document.indexing_status is COMPLETED
|
||||
- document.doc_form is not QA_INDEX
|
||||
- document.need_summary is True
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
import tasks.document_indexing_update_task as task_module
|
||||
from core.indexing_runner import DocumentIsPausedError
|
||||
from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
|
||||
from models.dataset import Dataset, Document, DocumentSegment
|
||||
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
|
||||
from tasks.document_indexing_update_task import document_indexing_update_task
|
||||
|
||||
|
||||
class _SessionContext:
|
||||
"""Minimal context manager that yields a mock session."""
|
||||
|
||||
def __init__(self, session: MagicMock) -> None:
|
||||
self._session = session
|
||||
|
||||
def __enter__(self) -> MagicMock:
|
||||
return self._session
|
||||
|
||||
def __exit__(self, exc_type, exc, tb) -> None: # type: ignore[override]
|
||||
return None
|
||||
@pytest.fixture
|
||||
def task_harness(
|
||||
sqlite_session: Session,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> tuple[MagicMock, MagicMock]:
|
||||
"""Bind task-owned sessions to SQLite and keep only external boundaries mocked."""
|
||||
engine = sqlite_session.get_bind()
|
||||
monkeypatch.setattr(
|
||||
task_module.session_factory,
|
||||
"create_session",
|
||||
lambda: Session(engine, expire_on_commit=False),
|
||||
)
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
monkeypatch.setattr(task_module, "IndexingRunner", MagicMock(return_value=runner))
|
||||
monkeypatch.setattr(
|
||||
task_module,
|
||||
"IndexProcessorFactory",
|
||||
MagicMock(return_value=MagicMock(init_index_processor=MagicMock(return_value=processor))),
|
||||
)
|
||||
return runner, processor
|
||||
|
||||
|
||||
def _make_dataset_and_documents(
|
||||
def _persist_rows(
|
||||
session: Session,
|
||||
*,
|
||||
dataset_id: str = "ds-1",
|
||||
document_id: str = "doc-1",
|
||||
indexing_technique: str = "high_quality",
|
||||
indexing_technique: IndexTechniqueType = IndexTechniqueType.HIGH_QUALITY,
|
||||
summary_index_setting: dict | None = None,
|
||||
doc_form: str = "text_model",
|
||||
doc_form: IndexStructureType = IndexStructureType.PARAGRAPH_INDEX,
|
||||
need_summary: bool = True,
|
||||
):
|
||||
"""Create mock dataset and document objects.
|
||||
|
||||
Returns (dataset, doc_for_session1, doc_for_session3).
|
||||
|
||||
session1 doc: before IndexingRunner runs (status irrelevant for summary).
|
||||
session3 doc: re-queried after IndexingRunner completes — normally COMPLETED.
|
||||
"""
|
||||
dataset = SimpleNamespace(
|
||||
with_segment: bool = False,
|
||||
) -> tuple[Dataset, Document]:
|
||||
tenant_id = str(uuid.uuid4())
|
||||
dataset_id = str(uuid.uuid4())
|
||||
document_id = str(uuid.uuid4())
|
||||
created_by = str(uuid.uuid4())
|
||||
dataset = Dataset(
|
||||
id=dataset_id,
|
||||
tenant_id=tenant_id,
|
||||
name="Update dataset",
|
||||
data_source_type=DataSourceType.UPLOAD_FILE,
|
||||
created_by=created_by,
|
||||
indexing_technique=indexing_technique,
|
||||
summary_index_setting=summary_index_setting,
|
||||
)
|
||||
doc_s1 = SimpleNamespace(
|
||||
document = Document(
|
||||
id=document_id,
|
||||
tenant_id=tenant_id,
|
||||
dataset_id=dataset_id,
|
||||
indexing_status="waiting",
|
||||
position=1,
|
||||
data_source_type=DataSourceType.UPLOAD_FILE,
|
||||
batch="batch-1",
|
||||
name="document.txt",
|
||||
created_from=DocumentCreatedFrom.WEB,
|
||||
created_by=created_by,
|
||||
indexing_status=IndexingStatus.WAITING,
|
||||
doc_form=doc_form,
|
||||
need_summary=need_summary,
|
||||
)
|
||||
# After IndexingRunner.run the document status is COMPLETED in the DB
|
||||
doc_s3 = SimpleNamespace(
|
||||
id=document_id,
|
||||
dataset_id=dataset_id,
|
||||
indexing_status="completed",
|
||||
doc_form=doc_form,
|
||||
need_summary=need_summary,
|
||||
rows: list[object] = [dataset, document]
|
||||
if with_segment:
|
||||
rows.append(
|
||||
DocumentSegment(
|
||||
tenant_id=tenant_id,
|
||||
dataset_id=dataset_id,
|
||||
document_id=document_id,
|
||||
position=1,
|
||||
content="segment",
|
||||
word_count=1,
|
||||
tokens=1,
|
||||
created_by=created_by,
|
||||
index_node_id="node-1",
|
||||
)
|
||||
)
|
||||
session.add_all(rows)
|
||||
session.commit()
|
||||
return dataset, document
|
||||
|
||||
|
||||
def _complete_indexing(documents: list[Document], _session: Session) -> None:
|
||||
for document in documents:
|
||||
document.indexing_status = IndexingStatus.COMPLETED
|
||||
|
||||
|
||||
def test_queues_summary_when_all_persisted_conditions_match(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, _processor = task_harness
|
||||
dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True})
|
||||
runner.run.side_effect = _complete_indexing
|
||||
delay = MagicMock()
|
||||
monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay)
|
||||
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
delay.assert_called_once_with(dataset.id, document.id, None)
|
||||
sqlite_session.expire_all()
|
||||
assert sqlite_session.get(Document, document.id).indexing_status == IndexingStatus.COMPLETED # type: ignore[union-attr]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("dataset_changes", "document_changes"),
|
||||
[
|
||||
({"indexing_technique": IndexTechniqueType.ECONOMY}, {}),
|
||||
({"summary_index_setting": None}, {}),
|
||||
({"summary_index_setting": {"enable": False}}, {}),
|
||||
({"summary_index_setting": {"enable": True}}, {"need_summary": False}),
|
||||
(
|
||||
{"summary_index_setting": {"enable": True}},
|
||||
{"doc_form": IndexStructureType.QA_INDEX},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_skips_summary_when_persisted_eligibility_does_not_match(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
dataset_changes: dict,
|
||||
document_changes: dict,
|
||||
) -> None:
|
||||
runner, _processor = task_harness
|
||||
dataset, document = _persist_rows(
|
||||
sqlite_session,
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
return dataset, doc_s1, doc_s3
|
||||
for key, value in dataset_changes.items():
|
||||
setattr(dataset, key, value)
|
||||
for key, value in document_changes.items():
|
||||
setattr(document, key, value)
|
||||
sqlite_session.commit()
|
||||
runner.run.side_effect = _complete_indexing
|
||||
delay = MagicMock()
|
||||
monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay)
|
||||
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
delay.assert_not_called()
|
||||
|
||||
|
||||
def _patch_all(monkeypatch: pytest.MonkeyPatch, *, sessions, runner, processor):
|
||||
"""Wire up all mocks for document_indexing_update_task."""
|
||||
@pytest.mark.parametrize(
|
||||
"error_factory",
|
||||
[
|
||||
lambda _document_id: RuntimeError("indexing failed"),
|
||||
lambda document_id: DocumentIsPausedError(f"{document_id} is paused"),
|
||||
],
|
||||
)
|
||||
def test_skips_summary_when_indexing_fails_or_is_paused(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
error_factory: Callable[[str], Exception],
|
||||
) -> None:
|
||||
runner, _processor = task_harness
|
||||
dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True})
|
||||
runner.run.side_effect = error_factory(document.id)
|
||||
delay = MagicMock()
|
||||
monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay)
|
||||
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
delay.assert_not_called()
|
||||
|
||||
|
||||
def test_returns_without_opening_external_boundaries_when_document_is_missing(
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
) -> None:
|
||||
runner, processor = task_harness
|
||||
|
||||
document_indexing_update_task(str(uuid.uuid4()), str(uuid.uuid4()))
|
||||
|
||||
runner.run.assert_not_called()
|
||||
processor.clean.assert_not_called()
|
||||
|
||||
|
||||
def test_skips_summary_when_dataset_is_removed_after_indexing(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, _processor = task_harness
|
||||
dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True})
|
||||
|
||||
def complete_and_remove(documents: list[Document], session: Session) -> None:
|
||||
_complete_indexing(documents, session)
|
||||
persisted_dataset = session.get(Dataset, dataset.id)
|
||||
assert persisted_dataset is not None
|
||||
session.delete(persisted_dataset)
|
||||
|
||||
runner.run.side_effect = complete_and_remove
|
||||
delay = MagicMock()
|
||||
monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay)
|
||||
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
delay.assert_not_called()
|
||||
|
||||
|
||||
def test_skips_summary_when_runner_leaves_document_incomplete(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, _processor = task_harness
|
||||
dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True})
|
||||
runner.run.return_value = None
|
||||
delay = MagicMock()
|
||||
monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay)
|
||||
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
delay.assert_not_called()
|
||||
|
||||
|
||||
def test_queue_failure_is_swallowed_after_successful_indexing(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, _processor = task_harness
|
||||
dataset, document = _persist_rows(sqlite_session, summary_index_setting={"enable": True})
|
||||
runner.run.side_effect = _complete_indexing
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.session_factory.create_session",
|
||||
MagicMock(side_effect=sessions),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.IndexProcessorFactory",
|
||||
MagicMock(return_value=MagicMock(init_index_processor=MagicMock(return_value=processor))),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.IndexingRunner",
|
||||
MagicMock(return_value=runner),
|
||||
task_module.generate_summary_index_task,
|
||||
"delay",
|
||||
MagicMock(side_effect=RuntimeError("queue unavailable")),
|
||||
)
|
||||
|
||||
|
||||
def _session_with_begin():
|
||||
"""Create a mock session with a begin() context manager."""
|
||||
s = MagicMock()
|
||||
s.begin.return_value = nullcontext()
|
||||
return s
|
||||
|
||||
|
||||
class TestUpdateTaskSummaryGeneration:
|
||||
"""Tests for summary index generation in the document update task.
|
||||
|
||||
The update task creates sessions in this order:
|
||||
1. session1: fetch document + dataset + segments (uses begin())
|
||||
2. session2: delete segments — only if segments exist (uses begin())
|
||||
3. session3: summary check — only if indexing succeeded (no begin())
|
||||
|
||||
With empty segments (default), only sessions 1 and 3 are created.
|
||||
"""
|
||||
|
||||
def test_should_queue_summary_when_conditions_met(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary task is queued when all conditions are met."""
|
||||
dataset, doc_s1, doc_s3 = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s3, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_called_once_with("ds-1", "doc-1", None)
|
||||
|
||||
def test_should_not_queue_when_not_high_quality(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when indexing_technique is not high_quality."""
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
indexing_technique="economy",
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s1, dataset] # dataset.indexing_technique == "economy"
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_summary_setting_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when summary_index_setting has enable=False."""
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": False},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s1, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_summary_setting_none(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when summary_index_setting is None."""
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
summary_index_setting=None,
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s1, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_need_summary_false(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when document.need_summary is False."""
|
||||
dataset, doc_s1, doc_s3 = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
need_summary=False,
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s3, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_qa_index_form(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when doc_form is QA_INDEX."""
|
||||
dataset, doc_s1, doc_s3 = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
doc_form="qa_model",
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s3, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_indexing_fails(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when IndexingRunner.run raises."""
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
runner = MagicMock()
|
||||
runner.run.side_effect = Exception("indexing failed")
|
||||
processor = MagicMock()
|
||||
|
||||
# Only session1 needed — task returns early after indexing failure
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_document_is_paused(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when IndexingRunner raises DocumentIsPausedError."""
|
||||
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
runner = MagicMock()
|
||||
runner.run.side_effect = DocumentIsPausedError("doc-1 is paused")
|
||||
processor = MagicMock()
|
||||
|
||||
# Only session1 needed — task returns early after paused error
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_dataset_not_found_after_indexing(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when the dataset disappears after indexing."""
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
# Session 3: dataset is None
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s1, None]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_not_queue_when_document_not_completed_after_indexing(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Summary is skipped when document indexing_status is not COMPLETED after indexing."""
|
||||
dataset, doc_s1, _ = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
# Document still in error status after indexing
|
||||
doc_s3_error = SimpleNamespace(
|
||||
id="doc-1",
|
||||
dataset_id="ds-1",
|
||||
indexing_status="error",
|
||||
doc_form="text_model",
|
||||
need_summary=True,
|
||||
)
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s3_error, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_not_called()
|
||||
|
||||
def test_should_swallow_summary_queue_error(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Task should not raise when generate_summary_index_task.delay raises."""
|
||||
dataset, doc_s1, doc_s3 = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s3, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[_SessionContext(session1), _SessionContext(session3)],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock(side_effect=Exception("queue full"))
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_called_once_with("ds-1", "doc-1", None)
|
||||
|
||||
def test_should_queue_summary_with_segments_and_session2(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""When segments exist, session2 is also created for deletion.
|
||||
Verify summary generation still works correctly."""
|
||||
dataset, doc_s1, doc_s3 = _make_dataset_and_documents(
|
||||
summary_index_setting={"enable": True},
|
||||
)
|
||||
|
||||
session1 = _session_with_begin()
|
||||
session1.scalar.side_effect = [doc_s1, dataset]
|
||||
seg = SimpleNamespace(index_node_id="node-1")
|
||||
session1.scalars.return_value = MagicMock(all=MagicMock(return_value=[seg]))
|
||||
|
||||
session3 = MagicMock()
|
||||
session3.scalar.side_effect = [doc_s3, dataset]
|
||||
|
||||
runner = MagicMock()
|
||||
processor = MagicMock()
|
||||
|
||||
_patch_all(
|
||||
monkeypatch,
|
||||
sessions=[
|
||||
_SessionContext(session1),
|
||||
_SessionContext(session3),
|
||||
],
|
||||
runner=runner,
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
delay_mock = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"tasks.document_indexing_update_task.generate_summary_index_task.delay",
|
||||
delay_mock,
|
||||
)
|
||||
|
||||
document_indexing_update_task("ds-1", "doc-1")
|
||||
|
||||
delay_mock.assert_called_once_with("ds-1", "doc-1", None)
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
sqlite_session.expire_all()
|
||||
assert sqlite_session.get(Document, document.id).indexing_status == IndexingStatus.COMPLETED # type: ignore[union-attr]
|
||||
|
||||
|
||||
def test_cleans_and_deletes_persisted_segments_with_real_session(
|
||||
sqlite_session: Session,
|
||||
task_harness: tuple[MagicMock, MagicMock],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, processor = task_harness
|
||||
dataset, document = _persist_rows(
|
||||
sqlite_session,
|
||||
summary_index_setting={"enable": True},
|
||||
with_segment=True,
|
||||
)
|
||||
runner.run.side_effect = _complete_indexing
|
||||
delay = MagicMock()
|
||||
monkeypatch.setattr(task_module.generate_summary_index_task, "delay", delay)
|
||||
|
||||
document_indexing_update_task(dataset.id, document.id)
|
||||
|
||||
processor.clean.assert_called_once()
|
||||
assert isinstance(processor.clean.call_args.kwargs["session"], Session)
|
||||
assert sqlite_session.scalars(select(DocumentSegment).where(DocumentSegment.document_id == document.id)).all() == []
|
||||
delay.assert_called_once_with(dataset.id, document.id, None)
|
||||
|
||||
Reference in New Issue
Block a user