mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-06 19:38:36 +08:00
Feat: Add knowledge compilation workflows (#16515)
## Summary - Add knowledge compilation template APIs, services, and builtin template seed data - Add advanced knowledge compile structure/artifact/RAPTOR workflow support - Update parsing, dataset/document APIs, and supporting services for compilation workflows
This commit is contained in:
Kevin Hu
@@ -120,7 +120,16 @@ class FileTestModel(BaseTestModel):
|
||||
db_table = "file"
|
||||
|
||||
|
||||
_TABLES = [FileCommitTestModel, FileCommitItemTestModel, FileTestModel]
|
||||
class UserTestModel(BaseTestModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
nickname = CharField(max_length=100, null=False, index=True)
|
||||
email = CharField(max_length=255, null=False)
|
||||
|
||||
class Meta:
|
||||
db_table = "user"
|
||||
|
||||
|
||||
_TABLES = [FileCommitTestModel, FileCommitItemTestModel, FileTestModel, UserTestModel]
|
||||
sqlite_db.create_tables(_TABLES)
|
||||
|
||||
|
||||
@@ -241,6 +250,7 @@ def _load_module(monkeypatch):
|
||||
db_models_mod.FileCommit = FileCommitTestModel
|
||||
db_models_mod.FileCommitItem = FileCommitItemTestModel
|
||||
db_models_mod.File = FileTestModel
|
||||
db_models_mod.User = UserTestModel
|
||||
db_models_mod.DataBaseModel = BaseTestModel
|
||||
monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod)
|
||||
|
||||
|
||||
@@ -71,7 +71,7 @@
|
||||
# "rag.utils.base64_image",
|
||||
|
||||
# "rag.prompts.generator",
|
||||
# "rag.raptor",
|
||||
# "rag.advanced_rag.knowlege_compile.raptor",
|
||||
# "rag.app",
|
||||
# "rag.graphrag.utils",
|
||||
# ]
|
||||
|
||||
@@ -1,438 +0,0 @@
|
||||
#
|
||||
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Unit tests for ChunkPostProcessor module.
|
||||
|
||||
Mock strategy: the LLM prompt functions (``keyword_extraction``, ``question_proposal``,
|
||||
``gen_metadata``, ``content_tagging``) are mocked since they make actual LLM API
|
||||
calls. ``get_llm_cache`` / ``set_llm_cache`` run as real code, so cache
|
||||
population and retrieval are exercised. ``rag_tokenizer`` is mocked because
|
||||
it requires NLTK data in the test environment.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
from rag.svr.task_executor_refactor.chunk_post_processor import (
|
||||
extract_keywords,
|
||||
generate_questions,
|
||||
generate_metadata,
|
||||
apply_tags,
|
||||
count_with_key,
|
||||
build_metadata_config,
|
||||
)
|
||||
from test.unit_test.rag.svr.task_executor_refactor.conftest import (
|
||||
make_task_context,
|
||||
MockChatModel,
|
||||
)
|
||||
|
||||
|
||||
class _BasePostProcessorTest:
|
||||
"""Shared helpers for post-processor test classes."""
|
||||
|
||||
@staticmethod
|
||||
def _mock_llm_binding(chat_model_cls=MockChatModel):
|
||||
"""Patch model config lookup + LLMBundle to return a MockChatModel."""
|
||||
p1 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance",
|
||||
return_value=MagicMock())
|
||||
p2 = patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle",
|
||||
return_value=chat_model_cls())
|
||||
return p1, p2
|
||||
|
||||
@staticmethod
|
||||
def _patch_prompt_func(func_path: str, return_value):
|
||||
"""Patch a prompt-level LLM function (the actual API call)."""
|
||||
return patch(func_path, return_value=return_value)
|
||||
|
||||
|
||||
class TestExtractKeywords(_BasePostProcessorTest):
|
||||
"""Tests for extract_keywords function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_success(self):
|
||||
"""Test successful keyword extraction — cache miss → LLM prompt runs."""
|
||||
ctx = make_task_context(parser_config={"auto_keywords": 5})
|
||||
docs = [
|
||||
{"content_with_weight": "This is test content one"},
|
||||
{"content_with_weight": "This is test content two"},
|
||||
]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None) # cache miss
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") # Redis stub
|
||||
p5 = self._patch_prompt_func(
|
||||
"rag.svr.task_executor_refactor.chunk_post_processor.keyword_extraction",
|
||||
return_value="keyword1, keyword2",
|
||||
)
|
||||
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer")
|
||||
with p1, p2, p3, p4, p5, p6 as mock_tok:
|
||||
mock_tok.tokenize.return_value = "keyword1 keyword2"
|
||||
await extract_keywords(docs, ctx)
|
||||
assert "important_kwd" in docs[0]
|
||||
assert "important_tks" in docs[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_canceled(self):
|
||||
"""Test keyword extraction when task is canceled."""
|
||||
ctx = make_task_context(parser_config={"auto_keywords": 5},
|
||||
has_canceled_func=MagicMock(return_value=True))
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
with p1, p2:
|
||||
await extract_keywords(docs, ctx)
|
||||
assert "important_kwd" not in docs[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_empty_docs(self):
|
||||
"""Test keyword extraction with empty docs list."""
|
||||
ctx = make_task_context(parser_config={"auto_keywords": 5})
|
||||
docs = []
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
with p1, p2:
|
||||
await extract_keywords(docs, ctx)
|
||||
ctx.progress_cb.assert_called()
|
||||
|
||||
|
||||
class TestGenerateQuestions(_BasePostProcessorTest):
|
||||
"""Tests for generate_questions function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_questions_success(self):
|
||||
"""Test successful question generation — cache miss → LLM prompt runs."""
|
||||
ctx = make_task_context(parser_config={"auto_questions": 3})
|
||||
docs = [{"content_with_weight": "This is test content one"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None)
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p5 = self._patch_prompt_func(
|
||||
"rag.svr.task_executor_refactor.chunk_post_processor.question_proposal",
|
||||
return_value="Question 1\nQuestion 2",
|
||||
)
|
||||
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer")
|
||||
with p1, p2, p3, p4, p5, p6 as mock_tok:
|
||||
mock_tok.tokenize.return_value = "Question 1 Question 2"
|
||||
await generate_questions(docs, ctx)
|
||||
assert "question_kwd" in docs[0]
|
||||
assert "question_tks" in docs[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_questions_canceled(self):
|
||||
"""Test question generation when task is canceled."""
|
||||
ctx = make_task_context(parser_config={"auto_questions": 3},
|
||||
has_canceled_func=MagicMock(return_value=True))
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
with p1, p2:
|
||||
await generate_questions(docs, ctx)
|
||||
assert "question_kwd" not in docs[0]
|
||||
|
||||
|
||||
class TestGenerateMetadata(_BasePostProcessorTest):
|
||||
"""Tests for generate_metadata function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_metadata_success(self):
|
||||
"""Test successful metadata generation — cache miss → LLM prompt runs."""
|
||||
ctx = make_task_context(
|
||||
parser_config={
|
||||
"enable_metadata": True,
|
||||
"metadata": [{"key": "category", "type": "string"}],
|
||||
"built_in_metadata": [{"key": "update_time", "type": "time"}],
|
||||
},
|
||||
)
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None)
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
|
||||
with p1, p2, p3, p4, p5 as mock_meta:
|
||||
mock_meta.get_document_metadata.return_value = {}
|
||||
mock_meta.update_document_metadata = MagicMock()
|
||||
await generate_metadata(docs, ctx)
|
||||
mock_meta.update_document_metadata.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_metadata_with_write_interceptor(self):
|
||||
"""Test metadata generation with write interceptor."""
|
||||
ctx = make_task_context(
|
||||
parser_config={
|
||||
"enable_metadata": True,
|
||||
"metadata": [{"key": "category", "type": "string"}],
|
||||
"built_in_metadata": [{"key": "update_time", "type": "time"}],
|
||||
},
|
||||
write_interceptor=MagicMock(),
|
||||
)
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None)
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
|
||||
with p1, p2, p3, p4, p5:
|
||||
await generate_metadata(docs, ctx)
|
||||
ctx.write_interceptor.intercept.assert_called_once_with(
|
||||
"DocMetadataService.update_document_metadata"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_metadata_empty_config_does_not_crash(self):
|
||||
"""Empty parser_config — no metadata configured — should not crash."""
|
||||
ctx = make_task_context(parser_config={})
|
||||
docs = [{"content_with_weight": "test"}]
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
|
||||
with p1, p2, p3:
|
||||
await generate_metadata(docs, ctx) # no exception = pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_metadata_enum_none_accepted(self):
|
||||
"""enum: None in metadata — treated as absent, should not crash."""
|
||||
ctx = make_task_context(
|
||||
parser_config={
|
||||
"enable_metadata": True,
|
||||
"metadata": [{"key": "format", "type": "string", "enum": None}],
|
||||
},
|
||||
)
|
||||
docs = [{"content_with_weight": "test"}]
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None)
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
|
||||
with p1, p2, p3, p4, p5:
|
||||
await generate_metadata(docs, ctx) # no exception = pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_metadata_description_none_accepted(self):
|
||||
"""description: None in metadata — should not crash."""
|
||||
ctx = make_task_context(
|
||||
parser_config={
|
||||
"enable_metadata": True,
|
||||
"metadata": [{"key": "test", "type": "string", "description": None}],
|
||||
},
|
||||
)
|
||||
docs = [{"content_with_weight": "test"}]
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None)
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
|
||||
with p1, p2, p3, p4, p5:
|
||||
await generate_metadata(docs, ctx) # no exception = pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_metadata_built_in_with_enum_none(self):
|
||||
"""built_in_metadata with enum: None — should not crash."""
|
||||
ctx = make_task_context(
|
||||
parser_config={
|
||||
"enable_metadata": True,
|
||||
"built_in_metadata": [
|
||||
{"key": "update_time", "type": "time", "description": None, "enum": None},
|
||||
],
|
||||
},
|
||||
)
|
||||
docs = [{"content_with_weight": "test"}]
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value=None)
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
|
||||
with p1, p2, p3, p4, p5:
|
||||
await generate_metadata(docs, ctx) # no exception = pass
|
||||
|
||||
|
||||
class TestApplyTags(_BasePostProcessorTest):
|
||||
"""Tests for apply_tags function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_tags_success(self):
|
||||
"""Test successful tag application with tag cache miss."""
|
||||
ctx = make_task_context(
|
||||
kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3},
|
||||
)
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings")
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value='{"tag1": 1}') # cache hit → skip LLM
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache",
|
||||
return_value=None)
|
||||
p7 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache")
|
||||
with p1, p2, p3 as mock_settings, p4, p5, p6 as mock_get_tags, p7 as mock_set_tags:
|
||||
mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10, "tag2": 5}
|
||||
mock_settings.retriever.tag_content.return_value = True
|
||||
await apply_tags(docs, ctx)
|
||||
assert len(docs) == 1
|
||||
mock_get_tags.assert_called_once()
|
||||
mock_set_tags.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_tags_canceled(self):
|
||||
"""Test tag application when task is canceled."""
|
||||
ctx = make_task_context(
|
||||
kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3},
|
||||
has_canceled_func=MagicMock(return_value=True),
|
||||
)
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings")
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache",
|
||||
return_value=None)
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache")
|
||||
with p1, p2, p3 as mock_settings, p4, p5:
|
||||
mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10}
|
||||
await apply_tags(docs, ctx)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_tags_tag_cache_miss(self):
|
||||
"""Test apply_tags when get_tags_from_cache returns None (cache miss)."""
|
||||
ctx = make_task_context(
|
||||
kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3},
|
||||
)
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings")
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value='{"tag1": 1}') # cache hit → skip LLM
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache",
|
||||
return_value=None)
|
||||
p7 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache")
|
||||
with p1, p2, p3 as mock_settings, p4, p5, p6 as mock_get_tags, p7 as mock_set_tags:
|
||||
mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10, "tag2": 5}
|
||||
mock_settings.retriever.tag_content.return_value = True
|
||||
await apply_tags(docs, ctx)
|
||||
mock_get_tags.assert_called_once_with(["kb_1"])
|
||||
mock_set_tags.assert_called_once()
|
||||
mock_settings.retriever.all_tags_in_portion.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_tags_tag_cache_hit(self):
|
||||
"""Test apply_tags when get_tags_from_cache returns valid data (cache hit)."""
|
||||
ctx = make_task_context(
|
||||
kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3},
|
||||
)
|
||||
docs = [{"content_with_weight": "This is test content"}]
|
||||
|
||||
p1, p2 = self._mock_llm_binding()
|
||||
p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings")
|
||||
p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache",
|
||||
return_value='{"tag1": 1}') # cache hit → skip LLM
|
||||
p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache")
|
||||
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache",
|
||||
return_value='{"cached_tag": 10}')
|
||||
p7 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache")
|
||||
with p1, p2, p3 as mock_settings, p4, p5, p6 as mock_get_tags, p7 as mock_set_tags:
|
||||
mock_settings.retriever.tag_content.return_value = True
|
||||
await apply_tags(docs, ctx)
|
||||
mock_get_tags.assert_called_once_with(["kb_1"])
|
||||
mock_set_tags.assert_not_called()
|
||||
mock_settings.retriever.all_tags_in_portion.assert_not_called()
|
||||
|
||||
|
||||
class TestCountWithKey:
|
||||
"""Tests for count_with_key function."""
|
||||
|
||||
def test_count_with_key_all_have_key(self):
|
||||
docs = [{"tag": 1}, {"tag": 2}, {"tag": 3}]
|
||||
assert count_with_key(docs, "tag") == 3
|
||||
|
||||
def test_count_with_key_some_have_key(self):
|
||||
docs = [{"tag": 1}, {"other": 2}, {"tag": 3}]
|
||||
assert count_with_key(docs, "tag") == 2
|
||||
|
||||
def test_count_with_key_none_have_key(self):
|
||||
docs = [{"other": 1}, {"other": 2}]
|
||||
assert count_with_key(docs, "tag") == 0
|
||||
|
||||
def test_count_with_key_empty_docs(self):
|
||||
assert count_with_key([], "tag") == 0
|
||||
|
||||
def test_count_with_key_falsy_value(self):
|
||||
docs = [{"tag": 0}, {"tag": ""}, {"tag": None}]
|
||||
assert count_with_key(docs, "tag") == 0
|
||||
|
||||
def test_count_with_key_truthy_value(self):
|
||||
docs = [{"tag": 1}, {"tag": "value"}, {"tag": [1, 2]}]
|
||||
assert count_with_key(docs, "tag") == 3
|
||||
|
||||
|
||||
class TestBuildMetadataConfig:
|
||||
"""Tests for build_metadata_config function."""
|
||||
|
||||
def test_dict_without_properties_returns_schema(self):
|
||||
parser_config = {"metadata": {"type": "object"}, "built_in_metadata": []}
|
||||
assert build_metadata_config(parser_config) == {"type": "object", "properties": {}}
|
||||
|
||||
def test_dict_with_properties_and_built_in(self):
|
||||
parser_config = {
|
||||
"metadata": {"type": "object", "properties": {"a": {"type": "string"}}},
|
||||
"built_in_metadata": [{"key": "author", "description": "Author name", "enum": ["alice", "bob"]}],
|
||||
}
|
||||
result = build_metadata_config(parser_config)
|
||||
assert "a" in result["properties"]
|
||||
assert "author" in result["properties"]
|
||||
|
||||
def test_dict_with_properties_no_built_in(self):
|
||||
parser_config = {
|
||||
"metadata": {"type": "object", "properties": {"a": {"type": "string"}}},
|
||||
"built_in_metadata": [],
|
||||
}
|
||||
result = build_metadata_config(parser_config)
|
||||
assert result == {"type": "object", "properties": {"a": {"type": "string"}}}
|
||||
|
||||
def test_list_with_built_in(self):
|
||||
parser_config = {
|
||||
"metadata": [{"key": "category"}],
|
||||
"built_in_metadata": [{"key": "author"}],
|
||||
}
|
||||
assert build_metadata_config(parser_config) == [{"key": "category"}, {"key": "author"}]
|
||||
|
||||
def test_list_without_built_in(self):
|
||||
parser_config = {"metadata": [{"key": "category"}], "built_in_metadata": []}
|
||||
assert build_metadata_config(parser_config) == [{"key": "category"}]
|
||||
|
||||
def test_other_type_with_built_in(self):
|
||||
parser_config = {"metadata": [], "built_in_metadata": [{"key": "author"}]}
|
||||
assert build_metadata_config(parser_config) == [{"key": "author"}]
|
||||
|
||||
def test_idempotent_same_input(self):
|
||||
parser_config = {
|
||||
"metadata": [{"key": "category"}],
|
||||
"built_in_metadata": [{"key": "author"}],
|
||||
}
|
||||
assert build_metadata_config(parser_config) == build_metadata_config(parser_config)
|
||||
|
||||
def test_missing_metadata_key(self):
|
||||
assert build_metadata_config({"built_in_metadata": []}) == []
|
||||
|
||||
def test_metadata_is_none(self):
|
||||
"""When metadata is None, built_in_metadata alone is returned."""
|
||||
parser_config = {"metadata": None, "built_in_metadata": [{"key": "author"}]}
|
||||
result = build_metadata_config(parser_config)
|
||||
assert result == [{"key": "author"}]
|
||||
@@ -1,318 +0,0 @@
|
||||
#
|
||||
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Unit tests for ChunkService module.
|
||||
|
||||
Note: After refactoring, some functionality has been moved to:
|
||||
- chunk_builder.py: Parser factory, run_chunking, extract_outline
|
||||
- chunk_post_processor.py: Keyword extraction, question generation, metadata, tagging
|
||||
|
||||
This test file now focuses on ChunkService-specific functionality:
|
||||
- build_chunks orchestration
|
||||
- _prepare_docs_and_upload
|
||||
- insert_chunks and related methods
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
from rag.svr.task_executor_refactor.chunk_service import ChunkService
|
||||
from test.unit_test.rag.svr.task_executor_refactor.conftest import make_task_context
|
||||
|
||||
|
||||
class TestChunkServiceInit:
|
||||
"""Tests for ChunkService initialization."""
|
||||
|
||||
def test_init_stores_task_context(self):
|
||||
"""Test that task context is stored."""
|
||||
ctx = MagicMock()
|
||||
service = ChunkService(ctx=ctx)
|
||||
assert service._task_context is ctx
|
||||
|
||||
|
||||
class TestChunkServiceBuildChunks:
|
||||
"""Tests for build_chunks method."""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_chunks_file_size_exceeded(self):
|
||||
"""Test build_chunks returns empty list when file size exceeds limit."""
|
||||
ctx = make_task_context(size=1000000000) # Very large size
|
||||
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.DOC_MAXIMUM_SIZE = 1000 # Small limit
|
||||
|
||||
mock_rec_ctx = MagicMock()
|
||||
ctx.recording_context = mock_rec_ctx
|
||||
|
||||
result = await service.build_chunks(b"test binary")
|
||||
|
||||
assert result == []
|
||||
mock_rec_ctx.record.assert_any_call("file_size_exceeded", True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_chunks_file_size_ok(self):
|
||||
"""Test build_chunks proceeds when file size is within limit."""
|
||||
ctx = make_task_context(size=1000)
|
||||
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.DOC_MAXIMUM_SIZE = 10000000 # Large limit
|
||||
|
||||
mock_rec_ctx = MagicMock()
|
||||
ctx.recording_context = mock_rec_ctx
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
|
||||
mock_parser = MagicMock()
|
||||
mock_get_parser.return_value = mock_parser
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
|
||||
mock_run_chunking.return_value = [{"content_with_weight": "test"}]
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
|
||||
with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
|
||||
mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
|
||||
|
||||
await service.build_chunks(b"test binary")
|
||||
|
||||
mock_rec_ctx.record.assert_any_call("file_size_exceeded", False)
|
||||
mock_rec_ctx.record.assert_any_call("parser_id", "naive")
|
||||
mock_get_parser.assert_called_once_with("naive")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("task_kwargs,func_path,func_name", [
|
||||
({"parser_config": {"auto_keywords": 5}}, "extract_keywords", "extract_keywords"),
|
||||
({"parser_config": {"auto_questions": 3}}, "generate_questions", "generate_questions"),
|
||||
({"kb_parser_config": {"tag_kb_ids": ["kb_1"]}}, "apply_tags", "apply_tags"),
|
||||
({"parser_config": {"enable_metadata": True, "metadata": [{"name": "category", "type": "string"}]}},
|
||||
"generate_metadata", "generate_metadata"),
|
||||
])
|
||||
async def test_build_chunks_with_post_processing(self, task_kwargs, func_path, func_name):
|
||||
"""Test build_chunks triggers post-processing when configured."""
|
||||
ctx = make_task_context(**task_kwargs)
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
mock_rec_ctx = MagicMock()
|
||||
ctx.recording_context = mock_rec_ctx
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock), \
|
||||
patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare, \
|
||||
patch(f"rag.svr.task_executor_refactor.chunk_service.{func_path}", new_callable=AsyncMock) as mock_fn:
|
||||
mock_settings.DOC_MAXIMUM_SIZE = 10000000
|
||||
mock_get_parser.return_value = MagicMock()
|
||||
mock_run_chunking.return_value = []
|
||||
mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
|
||||
await service.build_chunks(b"test binary")
|
||||
mock_fn.assert_called_once()
|
||||
|
||||
|
||||
class TestChunkServicePrepareDocsAndUpload:
|
||||
"""Tests for _prepare_docs_and_upload method."""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prepare_docs_and_upload_basic(self):
|
||||
"""Test basic document preparation."""
|
||||
ctx = make_task_context()
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
cks = [{"content_with_weight": "test chunk"}]
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.STORAGE_IMPL = MagicMock()
|
||||
mock_settings.STORAGE_IMPL.put = MagicMock()
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.image2id", new_callable=AsyncMock):
|
||||
|
||||
docs = await service._prepare_docs_and_upload(cks)
|
||||
|
||||
assert len(docs) == 1
|
||||
assert docs[0]["doc_id"] == "doc_1"
|
||||
assert docs[0]["kb_id"] == "kb_1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prepare_docs_and_upload_with_pagerank(self):
|
||||
"""Test document preparation with pagerank."""
|
||||
ctx = make_task_context()
|
||||
ctx.pagerank = 5
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
cks = [{"content_with_weight": "test chunk"}]
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.STORAGE_IMPL = MagicMock()
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.image2id", new_callable=AsyncMock):
|
||||
|
||||
docs = await service._prepare_docs_and_upload(cks)
|
||||
|
||||
assert docs[0].get("pagerank_fea") == 5
|
||||
|
||||
|
||||
class TestChunkServiceInsertChunks:
|
||||
"""Tests for insert_chunks method."""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_chunks_success(self):
|
||||
"""Test successful chunk insertion."""
|
||||
ctx = make_task_context()
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
chunks = [
|
||||
{"id": "chunk_1", "content_with_weight": "test1"},
|
||||
{"id": "chunk_2", "content_with_weight": "test2"},
|
||||
]
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.DOC_BULK_SIZE = 100
|
||||
mock_settings.docStoreConn = MagicMock()
|
||||
mock_settings.docStoreConn.insert = MagicMock(return_value=None)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
|
||||
mock_index.return_value = "test_index"
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
|
||||
mock_thread.return_value = None
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.TaskService") as mock_task:
|
||||
mock_task.update_chunk_ids = MagicMock()
|
||||
|
||||
result = await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_chunks_canceled(self):
|
||||
"""Test chunk insertion when task is canceled."""
|
||||
ctx = make_task_context()
|
||||
ctx.has_canceled_func = MagicMock(return_value=True)
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.DOC_BULK_SIZE = 100
|
||||
mock_settings.docStoreConn = MagicMock()
|
||||
mock_settings.docStoreConn.insert = MagicMock(return_value=None)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
|
||||
mock_index.return_value = "test_index"
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
|
||||
mock_thread.return_value = None
|
||||
|
||||
result = await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
|
||||
|
||||
assert result is False
|
||||
ctx.progress_cb.assert_called_with(-1, msg="Task has been canceled.")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_chunks_doc_store_error(self):
|
||||
"""Test chunk insertion when doc store returns error."""
|
||||
ctx = make_task_context()
|
||||
service = ChunkService(ctx=ctx)
|
||||
|
||||
chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
|
||||
mock_settings.DOC_BULK_SIZE = 100
|
||||
mock_settings.docStoreConn = MagicMock()
|
||||
mock_settings.docStoreConn.insert = MagicMock(return_value="Error message")
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
|
||||
mock_index.return_value = "test_index"
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
|
||||
mock_thread.return_value = "Error"
|
||||
|
||||
with pytest.raises(Exception, match="Insert chunk error"):
|
||||
await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
|
||||
|
||||
|
||||
class TestChunkServiceCreateMotherChunks:
|
||||
"""Tests for _create_mother_chunks class method."""
|
||||
|
||||
def test_create_mother_chunks_with_mom_field(self):
|
||||
"""Test creating mother chunks from mom field."""
|
||||
chunks = [
|
||||
{"id": "chunk_1", "mom": "Summary text 1", "content_with_weight": "test1"},
|
||||
]
|
||||
|
||||
mothers = ChunkService._create_mother_chunks(chunks)
|
||||
|
||||
assert len(mothers) == 1
|
||||
assert mothers[0]["content_with_weight"] == "Summary text 1"
|
||||
assert mothers[0]["available_int"] == 0
|
||||
|
||||
def test_create_mother_chunks_with_mom_with_weight_field(self):
|
||||
"""Test creating mother chunks from mom_with_weight field."""
|
||||
chunks = [
|
||||
{"id": "chunk_1", "mom_with_weight": "Summary text 2", "content_with_weight": "test1"},
|
||||
]
|
||||
|
||||
mothers = ChunkService._create_mother_chunks(chunks)
|
||||
|
||||
assert len(mothers) == 1
|
||||
assert mothers[0]["content_with_weight"] == "Summary text 2"
|
||||
|
||||
def test_create_mother_chunks_no_mom_field(self):
|
||||
"""Test creating mother chunks when no mom field present."""
|
||||
chunks = [
|
||||
{"id": "chunk_1", "content_with_weight": "test1"},
|
||||
]
|
||||
|
||||
mothers = ChunkService._create_mother_chunks(chunks)
|
||||
|
||||
assert len(mothers) == 0
|
||||
|
||||
def test_create_mother_chunks_empty_mom(self):
|
||||
"""Test creating mother chunks with empty mom field."""
|
||||
chunks = [
|
||||
{"id": "chunk_1", "mom": "", "content_with_weight": "test1"},
|
||||
]
|
||||
|
||||
mothers = ChunkService._create_mother_chunks(chunks)
|
||||
|
||||
assert len(mothers) == 0
|
||||
|
||||
def test_create_mother_chunks_deduplicates_ids(self):
|
||||
"""Test that mother chunks deduplicate by ID."""
|
||||
chunks = [
|
||||
{"id": "chunk_1", "mom": "Same summary", "content_with_weight": "test1"},
|
||||
{"id": "chunk_2", "mom": "Same summary", "content_with_weight": "test2"},
|
||||
]
|
||||
|
||||
mothers = ChunkService._create_mother_chunks(chunks)
|
||||
|
||||
assert len(mothers) == 1
|
||||
|
||||
def test_create_mother_chunks_filters_fields(self):
|
||||
"""Test that mother chunks only keep allowed fields."""
|
||||
chunks = [
|
||||
{"id": "chunk_1", "mom": "Summary", "extra_field": "should be removed", "content_with_weight": "test1"},
|
||||
]
|
||||
|
||||
mothers = ChunkService._create_mother_chunks(chunks)
|
||||
|
||||
assert "extra_field" not in mothers[0]
|
||||
assert "id" in mothers[0]
|
||||
assert "content_with_weight" in mothers[0]
|
||||
@@ -1,316 +0,0 @@
|
||||
#
|
||||
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Unit tests for TaskHandler module.
|
||||
|
||||
Mock strategy: external boundaries (LLMBundle, model config services, settings)
|
||||
are mocked so that ``handle()`` and ``_bind_embedding_model`` execute their
|
||||
real logic. Heavy orchestration methods (``_run_standard_chunking``,
|
||||
``_run_raptor``, ``_run_graphrag``) are mocked since they are tested
|
||||
exhaustively in the integration test suite.
|
||||
|
||||
Stable pure helpers (_build_toc) are tested directly.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
from rag.svr.task_executor_refactor.task_handler import TaskHandler
|
||||
|
||||
# Reuse shared helpers from conftest
|
||||
from test.unit_test.rag.svr.task_executor_refactor.conftest import (
|
||||
patch_embedding_binding,
|
||||
create_mock_settings,
|
||||
make_task_context,
|
||||
)
|
||||
|
||||
|
||||
class TestTaskHandlerHandleTask:
|
||||
"""Tests for the public handle_task() entry point."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_task_calls_handle(self):
|
||||
"""Test handle_task delegates to handle()."""
|
||||
ctx = MagicMock()
|
||||
ctx.id = "task_1"
|
||||
ctx.tenant_id = "tenant_1"
|
||||
ctx.kb_id = "kb_1"
|
||||
ctx.doc_id = "doc_1"
|
||||
ctx.has_canceled_func = MagicMock(return_value=False)
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler.handle = AsyncMock()
|
||||
await handler.handle_task()
|
||||
handler.handle.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_task_cleanup_on_cancel(self):
|
||||
"""Test handle_task cleans up docStore when canceled."""
|
||||
from common import settings
|
||||
mock_doc_store = MagicMock()
|
||||
mock_doc_store.index_exist = MagicMock(return_value=True)
|
||||
mock_doc_store.delete = MagicMock(return_value=None)
|
||||
orig = settings.docStoreConn
|
||||
settings.docStoreConn = mock_doc_store
|
||||
try:
|
||||
ctx = MagicMock()
|
||||
ctx.id = "task_1"
|
||||
ctx.tenant_id = "tenant_1"
|
||||
ctx.kb_id = "kb_1"
|
||||
ctx.doc_id = "doc_1"
|
||||
ctx.has_canceled_func = MagicMock(return_value=True)
|
||||
ctx.recording_context = MagicMock()
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler.handle = AsyncMock(side_effect=Exception("test error"))
|
||||
with pytest.raises(Exception, match="test error"):
|
||||
await handler.handle_task()
|
||||
mock_doc_store.delete.assert_called()
|
||||
finally:
|
||||
settings.docStoreConn = orig
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_task_cleanup_skips_when_index_missing(self):
|
||||
"""Cancel cleanup should not call delete when the index doesn't exist."""
|
||||
from common import settings
|
||||
mock_doc_store = MagicMock()
|
||||
mock_doc_store.index_exist = MagicMock(return_value=False)
|
||||
mock_doc_store.delete = MagicMock()
|
||||
orig = settings.docStoreConn
|
||||
settings.docStoreConn = mock_doc_store
|
||||
try:
|
||||
ctx = MagicMock()
|
||||
ctx.id = "task_1"
|
||||
ctx.tenant_id = "tenant_1"
|
||||
ctx.kb_id = "kb_1"
|
||||
ctx.doc_id = "doc_1"
|
||||
ctx.has_canceled_func = MagicMock(return_value=True)
|
||||
ctx.recording_context = MagicMock()
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler.handle = AsyncMock(side_effect=Exception("test error"))
|
||||
with pytest.raises(Exception, match="test error"):
|
||||
await handler.handle_task()
|
||||
mock_doc_store.delete.assert_not_called()
|
||||
finally:
|
||||
settings.docStoreConn = orig
|
||||
|
||||
|
||||
class TestTaskHandlerHandle:
|
||||
"""Tests for the public handle() method.
|
||||
|
||||
External boundaries (LLMBundle, model config services, settings) are mocked
|
||||
so that ``_bind_embedding_model`` and ``_init_kb`` execute their real logic
|
||||
through ``handle()``. Only the heavy orchestration methods
|
||||
(``_run_standard_chunking``, ``_run_raptor``, ``_run_graphrag``) are mocked.
|
||||
"""
|
||||
|
||||
# ── Context factory: make_task_context from conftest — see import above
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_memory_task(self):
|
||||
"""Test handle returns after dispatching memory task — no further processing."""
|
||||
ctx = make_task_context(task_type="memory")
|
||||
ctx.raw_task = {"memory_id": "mem_1", "id": "task_1"}
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.handle_save_to_memory_task",
|
||||
new_callable=AsyncMock) as mock_handle:
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler._run_standard_chunking = AsyncMock()
|
||||
handler._run_dataflow = AsyncMock()
|
||||
await handler.handle()
|
||||
|
||||
mock_handle.assert_called_once_with(ctx.raw_task)
|
||||
# After memory task, should return immediately — no further routing
|
||||
handler._run_standard_chunking.assert_not_called()
|
||||
handler._run_dataflow.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_dataflow_task(self):
|
||||
"""Test handle dispatches dataflow tasks (after embedding binding + init_kb)."""
|
||||
ctx = make_task_context(task_type="dataflow", doc_id="doc_1")
|
||||
|
||||
with patch_embedding_binding(), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler._run_dataflow = AsyncMock()
|
||||
await handler.handle()
|
||||
handler._run_dataflow.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_canceled_task(self):
|
||||
"""Test handle returns early when task is canceled."""
|
||||
ctx = make_task_context(has_canceled_func=MagicMock(return_value=True))
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
ctx.progress_cb.assert_called_once_with(-1, msg="Task has been canceled.")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_standard_chunking(self):
|
||||
"""Test handle routes to standard chunking.
|
||||
|
||||
``_bind_embedding_model`` and ``_init_kb`` run their real code;
|
||||
only the external boundary (LLM API, settings) is mocked.
|
||||
"""
|
||||
ctx = make_task_context()
|
||||
|
||||
with patch_embedding_binding(), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler._run_standard_chunking = AsyncMock()
|
||||
await handler.handle()
|
||||
handler._run_standard_chunking.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_raptor_task(self):
|
||||
"""Test handle routes to RAPTOR with real embedding binding."""
|
||||
ctx = make_task_context(task_type="raptor")
|
||||
|
||||
with patch_embedding_binding(), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler._run_raptor = AsyncMock()
|
||||
await handler.handle()
|
||||
handler._run_raptor.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_graphrag_task(self):
|
||||
"""Test handle routes to GraphRAG with real embedding binding."""
|
||||
ctx = make_task_context(task_type="graphrag")
|
||||
|
||||
with patch_embedding_binding(), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
handler._run_graphrag = AsyncMock()
|
||||
await handler.handle()
|
||||
handler._run_graphrag.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_embedding_model_failure(self):
|
||||
"""Test handle returns early when embedding model binding fails.
|
||||
|
||||
``LLMBundle`` is patched to raise, so ``_bind_embedding_model``
|
||||
itself raises — no need to mock the private method.
|
||||
"""
|
||||
ctx = make_task_context()
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle:
|
||||
|
||||
mock_cfg.return_value = MagicMock()
|
||||
mock_default.return_value = MagicMock()
|
||||
mock_bundle.side_effect = RuntimeError("embedding service unavailable")
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
with pytest.raises(RuntimeError, match="embedding service unavailable"):
|
||||
await handler.handle()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_storage_binary_none_raises_file_not_found(self):
|
||||
"""Verify that None binary from storage raises FileNotFoundError."""
|
||||
ctx = make_task_context()
|
||||
|
||||
with patch_embedding_binding(), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService.get_storage_address",
|
||||
return_value=("bucket_test", "name_test")), \
|
||||
patch.object(TaskHandler, "_get_storage_binary", new_callable=AsyncMock, return_value=None):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
# Do NOT mock _run_standard_chunking — we want real code path for the check
|
||||
with pytest.raises(FileNotFoundError, match="Can not find file <test.pdf> from minio"):
|
||||
await handler.handle()
|
||||
|
||||
|
||||
class TestTaskHandlerBuildToc:
|
||||
"""Tests for _build_toc — stable pure helper (requires LLM mocking)."""
|
||||
|
||||
def test_build_toc_with_empty_docs(self):
|
||||
"""Test _build_toc returns None when run_toc_from_text returns empty."""
|
||||
ctx = MagicMock()
|
||||
ctx.tenant_id = "tenant_1"
|
||||
ctx.llm_id = "llm_1"
|
||||
ctx.language = "en"
|
||||
|
||||
docs = [{"id": "chunk_1", "content_with_weight": "text", "page_num_int": [1], "top_int": [0]}]
|
||||
|
||||
def mock_asyncio_run(coro):
|
||||
coro.close()
|
||||
return []
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run):
|
||||
|
||||
mock_cfg.return_value = MagicMock()
|
||||
mock_msg = MagicMock()
|
||||
mock_bundle.return_value.__enter__.return_value = mock_msg
|
||||
|
||||
result = TaskHandler._build_toc(ctx, docs, MagicMock())
|
||||
assert result is None
|
||||
|
||||
def test_build_toc_with_results(self):
|
||||
"""Test _build_toc builds TOC chunk when results exist."""
|
||||
ctx = MagicMock()
|
||||
ctx.tenant_id = "tenant_1"
|
||||
ctx.llm_id = "llm_1"
|
||||
ctx.language = "en"
|
||||
|
||||
docs = [{"id": "chunk_0", "content_with_weight": "text", "doc_id": "doc_1", "page_num_int": [1], "top_int": [0]}]
|
||||
toc_result = [{"chunk_id": "0", "title": "Section 1"}]
|
||||
|
||||
def mock_asyncio_run(coro):
|
||||
coro.close()
|
||||
return toc_result
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run):
|
||||
|
||||
mock_cfg.return_value = MagicMock()
|
||||
mock_msg = MagicMock()
|
||||
mock_bundle.return_value.__enter__.return_value = mock_msg
|
||||
|
||||
result = TaskHandler._build_toc(ctx, docs, MagicMock())
|
||||
assert result is not None
|
||||
assert "toc_kwd" in result
|
||||
assert result["toc_kwd"] == "toc"
|
||||
assert result["available_int"] == 0
|
||||
|
||||
|
||||
class TestTaskHandlerInit:
|
||||
"""Tests for TaskHandler initialization."""
|
||||
|
||||
def test_init_stores_context_and_hook(self):
|
||||
ctx = MagicMock()
|
||||
hook = MagicMock()
|
||||
handler = TaskHandler(ctx=ctx, billing_hook=hook)
|
||||
assert handler._task_context is ctx
|
||||
assert handler._billing_hook is hook
|
||||
|
||||
def test_init_default_hook_none(self):
|
||||
ctx = MagicMock()
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
assert handler._billing_hook is None
|
||||
@@ -1,856 +0,0 @@
|
||||
#
|
||||
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Integration tests for TaskHandler orchestration.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
from typing import Any, Dict
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from rag.svr.task_executor_refactor.task_handler import TaskHandler
|
||||
from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
|
||||
from rag.svr.task_executor_refactor.recording_context import BaseRecordingContext, RecordingContext
|
||||
from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID
|
||||
|
||||
# Import shared helpers from conftest
|
||||
from test.unit_test.rag.svr.task_executor_refactor.conftest import (
|
||||
AsyncMockLimiter,
|
||||
create_mock_embedding_model,
|
||||
create_default_chunks,
|
||||
create_mock_settings,
|
||||
create_mock_chunk_service,
|
||||
make_task_dict,
|
||||
patch_get_storage_binary,
|
||||
patch_task_handler_settings,
|
||||
mock_thread_return_binary,
|
||||
mock_thread_return_none,
|
||||
)
|
||||
|
||||
|
||||
def create_task_context(
|
||||
task_dict: Dict[str, Any],
|
||||
is_canceled: bool = False,
|
||||
recording_context: BaseRecordingContext | None = None,
|
||||
) -> TaskContext:
|
||||
"""Create a real TaskContext with mocked limiters and callbacks.
|
||||
|
||||
Args:
|
||||
task_dict: Task dictionary with all task attributes.
|
||||
is_canceled: If True, has_canceled_func returns True.
|
||||
recording_context: RecordingContext to inject. If None, a new one
|
||||
is created automatically so that recording_context access works.
|
||||
|
||||
Returns:
|
||||
TaskContext with all required dependencies injected.
|
||||
"""
|
||||
if recording_context is None:
|
||||
recording_context = RecordingContext()
|
||||
limiter = AsyncMockLimiter()
|
||||
progress_callback = MagicMock()
|
||||
ctx = TaskContext(
|
||||
task=task_dict,
|
||||
limiters=TaskLimiters(
|
||||
chat=limiter,
|
||||
minio=limiter,
|
||||
chunk=limiter,
|
||||
embed=limiter,
|
||||
kg=limiter,
|
||||
),
|
||||
callbacks=TaskCallbacks(
|
||||
progress=progress_callback,
|
||||
has_canceled=MagicMock(return_value=is_canceled),
|
||||
),
|
||||
recording_context=recording_context,
|
||||
)
|
||||
# Add progress_callback property for task_handler compatibility
|
||||
ctx.progress_callback = progress_callback
|
||||
# Add set_progress_cb method for task_handler compatibility
|
||||
ctx.set_progress_cb = lambda cb: setattr(ctx.callbacks, 'progress_cb', cb)
|
||||
return ctx
|
||||
|
||||
|
||||
class TestStandardChunkingPipelineIntegration:
|
||||
"""P0: Integration tests for the complete standard chunking pipeline."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_chunking_pipeline_records_task_status(self):
|
||||
"""Verify that the complete pipeline records task_status as 'completed'."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
task_status = recording_ctx.get("task_status")
|
||||
assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_chunking_pipeline_records_insertion_result(self):
|
||||
"""Verify that insertion_result is recorded as 'success'."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
insertion_result = recording_ctx.get("insertion_result")
|
||||
assert insertion_result == "success", f"Expected insertion_result='success', got {insertion_result}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_chunking_pipeline_records_chunk_ids(self):
|
||||
"""Verify that chunk_ids_count is recorded after build_chunks."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunks = create_default_chunks(count=3)
|
||||
mock_chunk_service = create_mock_chunk_service(chunks=mock_chunks)
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
mock_run_toc.return_value = [] # TOC returns empty when not enabled
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
chunk_ids_count = recording_ctx.get("chunk_ids_count")
|
||||
assert chunk_ids_count is not None, "chunk_ids_count should be recorded"
|
||||
assert chunk_ids_count == 3, f"Expected chunk_ids_count=3, got {chunk_ids_count}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_chunking_pipeline_records_token_count(self):
|
||||
"""Verify that token_count and vector_size are recorded after embedding."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
token_count = recording_ctx.get("token_count")
|
||||
vector_size = recording_ctx.get("vector_size")
|
||||
|
||||
assert token_count is not None, "token_count should be recorded"
|
||||
assert vector_size is not None, "vector_size should be recorded"
|
||||
assert vector_size == 128, f"Expected vector_size=128, got {vector_size}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_chunking_pipeline_progress_callback_invoked(self):
|
||||
"""Verify that progress_callback is invoked multiple times during pipeline."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
ctx.progress_callback.assert_called()
|
||||
call_count = ctx.progress_callback.call_count
|
||||
assert call_count > 0, "progress_callback should have been invoked at least once"
|
||||
|
||||
|
||||
class TestTaskCancellationCleanupIntegration:
|
||||
"""P0: Integration tests for task cancellation cleanup flow."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_canceled_task_calls_docstore_delete(self):
|
||||
"""Verify that docStoreConn.delete is called when task is canceled."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict, is_canceled=True)
|
||||
mock_settings = create_mock_settings()
|
||||
|
||||
call_log = []
|
||||
|
||||
def mock_thread_impl(func, *args, **kwargs):
|
||||
# Get the actual method name from the mock
|
||||
func_repr = repr(func)
|
||||
call_log.append(func_repr)
|
||||
if 'index_exist' in func_repr:
|
||||
return True
|
||||
if 'delete' in func_repr:
|
||||
return {"result": "deleted"}
|
||||
return {"result": "deleted"}
|
||||
|
||||
with patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_index"), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec", side_effect=mock_thread_impl):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle_task()
|
||||
|
||||
# Verify delete was called by checking the call log
|
||||
delete_calls = [c for c in call_log if 'delete' in c]
|
||||
assert len(delete_calls) >= 1, f"Expected at least one delete call, got: {call_log}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_canceled_task_progress_callback_with_negative_one(self):
|
||||
"""Verify that progress_callback is called with -1 when task is canceled."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict, is_canceled=True)
|
||||
mock_settings = create_mock_settings()
|
||||
|
||||
def mock_thread_impl(func, *args, **kwargs):
|
||||
func_repr = repr(func)
|
||||
if 'index_exist' in func_repr:
|
||||
return True
|
||||
if 'delete' in func_repr:
|
||||
return {"result": "deleted"}
|
||||
return {"result": "deleted"}
|
||||
|
||||
with patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_index"), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec", side_effect=mock_thread_impl):
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle_task()
|
||||
|
||||
ctx.progress_callback.assert_called()
|
||||
call_args_list = ctx.progress_callback.call_args_list
|
||||
# Check for -1 in any position of the call arguments
|
||||
has_negative_progress = False
|
||||
for call in call_args_list:
|
||||
# Check positional args
|
||||
for arg in call[0]:
|
||||
if arg == -1:
|
||||
has_negative_progress = True
|
||||
break
|
||||
# Check keyword args
|
||||
if call[1].get("prog") == -1:
|
||||
has_negative_progress = True
|
||||
if has_negative_progress:
|
||||
break
|
||||
assert has_negative_progress, f"progress_callback should have been called with -1 progress. Calls: {call_args_list}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_canceled_task_does_not_proceed_to_chunking(self):
|
||||
"""Verify that canceled task does not proceed to embedding model binding."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict, is_canceled=True)
|
||||
mock_settings = create_mock_settings()
|
||||
|
||||
with patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
|
||||
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_settings.docStoreConn.index_exist.return_value = True
|
||||
mock_settings.docStoreConn.delete.return_value = {"result": "deleted"}
|
||||
|
||||
async def mock_thread_impl(func, *args, **kwargs):
|
||||
return {"result": "deleted"}
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_impl
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle_task()
|
||||
|
||||
mock_bundle.assert_not_called()
|
||||
|
||||
|
||||
class TestRaptorPipelineIntegration:
|
||||
"""P1: Integration tests for the RAPTOR pipeline."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raptor_pipeline_records_task_status(self):
|
||||
"""Verify that RAPTOR pipeline records task_status."""
|
||||
task_dict = make_task_dict(doc_id=GRAPH_RAPTOR_FAKE_DOC_ID, doc_ids=["doc1", "doc2"], task_type="raptor", parser_config={"raptor": {"use_raptor": False}}, kb_parser_config={"raptor": {"use_raptor": False}})
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_kb = MagicMock()
|
||||
mock_kb.id = "kb_test"
|
||||
mock_kb.parser_config = {"raptor": {"use_raptor": False}}
|
||||
|
||||
with patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.RaptorService") as mock_raptor_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_kb_service.get_by_id.return_value = (True, mock_kb)
|
||||
mock_kb_service.update_by_id.return_value = True
|
||||
mock_raptor_service.return_value.run_raptor_for_kb = AsyncMock(return_value=([], 0, []))
|
||||
mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_none
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
task_status = recording_ctx.get("task_status")
|
||||
assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raptor_pipeline_enables_raptor_if_not_configured(self):
|
||||
"""Verify that RAPTOR is enabled if not already configured."""
|
||||
task_dict = make_task_dict(doc_id=GRAPH_RAPTOR_FAKE_DOC_ID, doc_ids=["doc1", "doc2"], task_type="raptor", parser_config={"raptor": {"use_raptor": False}}, kb_parser_config={"raptor": {"use_raptor": False}})
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_kb = MagicMock()
|
||||
mock_kb.id = "kb_test"
|
||||
mock_kb.parser_config = {"raptor": {"use_raptor": False}}
|
||||
|
||||
with patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.RaptorService") as mock_raptor_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_kb_service.get_by_id.return_value = (True, mock_kb)
|
||||
mock_kb_service.update_by_id.return_value = True
|
||||
mock_raptor_service.return_value.run_raptor_for_kb = AsyncMock(return_value=([], 0, []))
|
||||
mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_none
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
# Check that the kb parser_config was updated
|
||||
mock_kb_service.update_by_id.assert_called_once()
|
||||
call_args = mock_kb_service.update_by_id.call_args
|
||||
update_dict = call_args[0][1]
|
||||
assert update_dict.get("parser_config", {}).get("raptor", {}).get("use_raptor") is True, \
|
||||
"RAPTOR should be enabled in parser_config after running"
|
||||
|
||||
|
||||
class TestEmbeddingModelBindingFailureIntegration:
|
||||
"""P1: Integration tests for embedding model binding failure."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_binding_failure_raises_exception(self):
|
||||
"""Verify that embedding model binding failure raises an exception."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
|
||||
|
||||
mock_get_config.side_effect = Exception("Model not found")
|
||||
mock_get_default.side_effect = Exception("Model not found")
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
|
||||
with pytest.raises(Exception, match="Model not found"):
|
||||
await handler.handle()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_binding_failure_calls_progress_callback(self):
|
||||
"""Verify that embedding model binding failure calls progress_callback."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
|
||||
|
||||
mock_get_config.side_effect = Exception("Model not found")
|
||||
mock_get_default.side_effect = Exception("Model not found")
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await handler.handle()
|
||||
|
||||
ctx.progress_callback.assert_called()
|
||||
|
||||
|
||||
class TestDataflowPipelineIntegration:
|
||||
"""P2: Integration tests for the dataflow pipeline."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dataflow_pipeline_calls_dataflow_service(self):
|
||||
"""Verify that dataflow pipeline calls DataflowService.run_dataflow()."""
|
||||
task_dict = make_task_dict(doc_id=CANVAS_DEBUG_DOC_ID, task_type="dataflow")
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
|
||||
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.settings") as mock_settings, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DataflowService") as mock_dataflow_service:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_settings.docStoreConn = MagicMock()
|
||||
mock_settings.docStoreConn.create_idx = MagicMock()
|
||||
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.run_dataflow = AsyncMock(return_value=None)
|
||||
mock_dataflow_service.return_value = mock_instance
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
mock_dataflow_service.assert_called_once()
|
||||
mock_instance.run_dataflow.assert_called_once()
|
||||
|
||||
|
||||
class TestTocAsyncFlowIntegration:
|
||||
"""P2: Integration tests for TOC async flow."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_toc_async_flow_creates_toc_thread(self):
|
||||
"""Verify that TOC async flow creates a TOC thread when enabled."""
|
||||
|
||||
task_dict = make_task_dict(parser_config={"auto_keywords": 0, "auto_questions": 0, "enable_metadata": False, "toc_extraction": True})
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls, \
|
||||
patch("rag.svr.task_executor_refactor.post_processor.DocumentService") as mock_post_doc_service:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
mock_run_toc.return_value = [{"title": "Test TOC", "level": 1}]
|
||||
mock_post_doc_service.increment_chunk_num = MagicMock()
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
mock_run_toc.assert_called()
|
||||
|
||||
# Explicit cleanup to prevent resource leaks
|
||||
del mock_embedding, mock_settings, mock_chunk_service
|
||||
del mock_get_config, mock_get_default, mock_bundle, mock_file_service
|
||||
del mock_index_name, mock_doc_service, mock_chunk_service_cls, mock_run_toc, mock_post_doc_service
|
||||
del mock_thread_exec, mock_chunk_thread_exec
|
||||
# Allow pending callbacks to execute
|
||||
await asyncio.sleep(0)
|
||||
gc.collect()
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="function")
|
||||
@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
|
||||
async def test_toc_async_flow_does_not_create_thread_when_disabled(self):
|
||||
"""Verify that TOC async flow does not create a thread when disabled.
|
||||
|
||||
Note: This test has a known issue with resource leaks (unclosed sockets and
|
||||
event loops) when run as part of the full test suite. The warning filter
|
||||
above suppresses these warnings temporarily. The root cause is related to
|
||||
asyncio.to_thread creating new event loops that are not properly cleaned up
|
||||
by pytest-asyncio.
|
||||
"""
|
||||
|
||||
task_dict = make_task_dict(parser_config={"auto_keywords": 0, "auto_questions": 0, "enable_metadata": False, "toc_extraction": True})
|
||||
task_dict["parser_config"]["toc_extraction"] = False
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
mock_run_toc.assert_not_called()
|
||||
|
||||
# Explicit cleanup to prevent resource leaks
|
||||
del mock_embedding, mock_settings, mock_chunk_service
|
||||
del mock_get_config, mock_get_default, mock_bundle, mock_file_service
|
||||
del mock_index_name, mock_doc_service, mock_chunk_service_cls, mock_run_toc
|
||||
del mock_thread_exec, mock_chunk_thread_exec
|
||||
# Allow pending callbacks to execute and close event loop
|
||||
await asyncio.sleep(0)
|
||||
# Cancel all pending tasks
|
||||
current_task = asyncio.current_task()
|
||||
pending = [t for t in asyncio.all_tasks() if t is not current_task and not t.done()]
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if pending:
|
||||
await asyncio.gather(*pending, return_exceptions=True)
|
||||
gc.collect()
|
||||
|
||||
|
||||
class TestRecordingContextDataFlowAssertions:
|
||||
"""P2: Integration tests for RecordingContext data flow assertions."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recording_context_captures_file_size_check(self):
|
||||
"""Verify that RecordingContext captures file_size_exceeded result."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
file_size_exceeded = recording_ctx.get("file_size_exceeded")
|
||||
assert file_size_exceeded is None or file_size_exceeded is False, \
|
||||
f"Expected file_size_exceeded to be False/None for small file, got {file_size_exceeded}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recording_context_captures_parser_id(self):
|
||||
"""Verify that RecordingContext captures parser_id from task context."""
|
||||
task_dict = make_task_dict()
|
||||
ctx = create_task_context(task_dict)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_chunk_service = create_mock_chunk_service()
|
||||
|
||||
with patch_get_storage_binary(), \
|
||||
patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_doc_service.increment_chunk_num = MagicMock()
|
||||
mock_doc_service.get_document_metadata.return_value = {}
|
||||
mock_doc_service.update_document_metadata = MagicMock()
|
||||
mock_chunk_service_cls.return_value = mock_chunk_service
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_binary
|
||||
mock_chunk_thread_exec.side_effect = mock_thread_return_binary
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
recording_ctx = ctx.recording_context
|
||||
# parser_id is available in the task context, verify task completion
|
||||
task_status = recording_ctx.get("task_status")
|
||||
assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
|
||||
# Verify the parser_id is accessible from the task context
|
||||
assert ctx.parser_id == "naive", f"Expected parser_id='naive', got {ctx.parser_id}"
|
||||
|
||||
|
||||
class TestGraphragPipelineIntegration:
|
||||
"""P2: Integration tests for GraphRAG pipeline default configuration."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graphrag_pipeline_configures_full_defaults(self):
|
||||
"""Verify that GraphRAG configures all default parameters when not already set."""
|
||||
task_dict = make_task_dict(doc_ids=["doc1", "doc2"], task_type="graphrag")
|
||||
rec_ctx = RecordingContext()
|
||||
ctx = create_task_context(task_dict, recording_context=rec_ctx)
|
||||
mock_embedding = create_mock_embedding_model(vector_size=128)
|
||||
mock_settings = create_mock_settings()
|
||||
mock_kb = MagicMock()
|
||||
mock_kb.id = "kb_test"
|
||||
mock_kb.parser_config = {}
|
||||
|
||||
with patch_task_handler_settings(mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.run_graphrag_for_kb") as mock_run_graphrag, \
|
||||
patch("rag.svr.task_executor_refactor.task_handler.DocumentService"):
|
||||
|
||||
mock_get_config.return_value = MagicMock()
|
||||
mock_get_default.return_value = MagicMock()
|
||||
mock_bundle.return_value = mock_embedding
|
||||
mock_index_name.return_value = "test_index"
|
||||
mock_kb_service.get_by_id.return_value = (True, mock_kb)
|
||||
mock_kb_service.update_by_id.return_value = True
|
||||
mock_run_graphrag.return_value = {"status": "completed"}
|
||||
|
||||
mock_thread_exec.side_effect = mock_thread_return_none
|
||||
|
||||
handler = TaskHandler(ctx=ctx)
|
||||
await handler.handle()
|
||||
|
||||
# Verify update_by_id was called with full default config
|
||||
mock_kb_service.update_by_id.assert_called_once()
|
||||
call_args = mock_kb_service.update_by_id.call_args
|
||||
config = call_args[0][1]["parser_config"]["graphrag"]
|
||||
assert config["use_graphrag"] is True
|
||||
assert "organization" in config["entity_types"]
|
||||
assert "person" in config["entity_types"]
|
||||
assert "geo" in config["entity_types"]
|
||||
assert "event" in config["entity_types"]
|
||||
assert "category" in config["entity_types"]
|
||||
assert config["method"] == "light"
|
||||
assert "batch_chunk_token_size" in config
|
||||
assert "retry_attempts" in config
|
||||
assert "retry_backoff_seconds" in config
|
||||
assert "retry_backoff_max_seconds" in config
|
||||
assert "build_subgraph_timeout_per_chunk_seconds" in config
|
||||
assert "build_subgraph_min_timeout_seconds" in config
|
||||
assert "merge_timeout_seconds" in config
|
||||
assert "resolution_timeout_seconds" in config
|
||||
assert "community_timeout_seconds" in config
|
||||
assert "lock_acquire_timeout_seconds" in config, \
|
||||
"All GraphRAG default config parameters should be present"
|
||||
@@ -15,8 +15,10 @@
|
||||
#
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -113,10 +115,28 @@ def raptor_module(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "common.token_utils", token_utils_module)
|
||||
monkeypatch.setitem(sys.modules, "rag.graphrag.utils", graphrag_utils_module)
|
||||
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_module)
|
||||
monkeypatch.delitem(sys.modules, "rag.raptor", raising=False)
|
||||
module = importlib.import_module("rag.raptor")
|
||||
# Create stub parent packages and load raptor directly via spec_from_file_location
|
||||
# to bypass rag/advanced_rag/__init__.py (which triggers ES connection etc.).
|
||||
_test_dir = os.path.dirname(__file__)
|
||||
_rag_adv_kc_dir = os.path.normpath(os.path.join(_test_dir, "../../../rag/advanced_rag/knowlege_compile"))
|
||||
_rag_adv = types.ModuleType("rag.advanced_rag")
|
||||
_rag_adv.__path__ = [os.path.normpath(os.path.join(_test_dir, "../../../rag/advanced_rag"))]
|
||||
_rag_adv.__package__ = "rag.advanced_rag"
|
||||
monkeypatch.setitem(sys.modules, "rag.advanced_rag", _rag_adv)
|
||||
_rag_adv_kc = types.ModuleType("rag.advanced_rag.knowlege_compile")
|
||||
_rag_adv_kc.__path__ = [_rag_adv_kc_dir]
|
||||
_rag_adv_kc.__package__ = "rag.advanced_rag.knowlege_compile"
|
||||
monkeypatch.setitem(sys.modules, "rag.advanced_rag.knowlege_compile", _rag_adv_kc)
|
||||
monkeypatch.delitem(sys.modules, "rag.advanced_rag.knowlege_compile.raptor", raising=False)
|
||||
_raptor_spec = importlib.util.spec_from_file_location(
|
||||
"rag.advanced_rag.knowlege_compile.raptor",
|
||||
os.path.join(_rag_adv_kc_dir, "raptor.py"),
|
||||
)
|
||||
module = importlib.util.module_from_spec(_raptor_spec)
|
||||
sys.modules["rag.advanced_rag.knowlege_compile.raptor"] = module
|
||||
_raptor_spec.loader.exec_module(module)
|
||||
yield module
|
||||
monkeypatch.delitem(sys.modules, "rag.raptor", raising=False)
|
||||
monkeypatch.delitem(sys.modules, "rag.advanced_rag.knowlege_compile.raptor", raising=False)
|
||||
|
||||
|
||||
class FakeChatModel:
|
||||
@@ -372,7 +392,7 @@ async def test_psi_tree_builder_materializes_rebalanced_summary_layers_without_u
|
||||
def fail_umap(*args, **kwargs):
|
||||
raise AssertionError("Psi tree builder must use original embeddings, not UMAP")
|
||||
|
||||
monkeypatch.setattr(raptor_module.umap, "UMAP", fail_umap)
|
||||
monkeypatch.setattr("umap.UMAP", fail_umap)
|
||||
raptor = _make_raptor(raptor_module, max_cluster=2)
|
||||
|
||||
chunks, layers = await raptor(_chunks(), random_state=0)
|
||||
|
||||
Reference in New Issue
Block a user