zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-15 20:27:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: overhaul task executor with layered architecture and comprehensive test suite (#15471)

Browse Source 
## Summary

Decomposes the monolithic `task_executor.py` (1945 lines) into a 6-layer
architecture with clear separation of concerns. The refactored code is
functionally equivalent to the original, verified through 400 passing
tests and a production-vs-dry-run comparison framework.

## Architecture

```
entry (task_manager)
  └─ orchestration (task_handler)
       ├─ services (chunk_service, embedding_service, dataflow_service, raptor_service, post_processor)
       │    └─ utilities (chunk_builder, chunk_post_processor, embedding_utils)
       └─ infrastructure (task_context, recording_context, interceptor)
```

Key design decisions:
- **TaskContext** — typed facade over raw task dict, injects rate
limiters + callbacks via composition
- **RecordingContext + Comparator** — enables side-by-side production vs
dry-run execution for safe migration
- **NullRecordingContext** — zero-allocation no-op for production, uses
`__slots__`
- **WriteOperationInterceptor** — FIFO replay of previous runs function
returns for comparison mode

## Migration Strategy

The original `handle_task()` in `task_executor.py` uses a 3-way switch
via `TE_RUN_MODE`:
- `TE_RUN_MODE=0` (default) → runs refactored code
- `TE_RUN_MODE=1` → runs both original + refactored, compares all
intermediate results
- `TE_RUN_MODE=2` → runs original code (fallback)

The comparison mode (`TE_RUN_MODE=1`) records ~40 intermediate values
(chunks, vectors, token counts, func return values) from the production
run and replays them during dry-run, then uses `ContextComparator` to
report mismatches.

## Functional Equivalence Fixes

All divergences between original and refactored code were identified and
fixed:
- Timeout decorators (handle/build_chunks/raptor/embedding)
- NullRecordingContext leak in finally block causing RuntimeError
- MinIO None-binary check with proper FileNotFoundError
- Dataflow dispatch after embedding binding + init_kb
- Memory task missing return after processing
- RAPTOR checkpoint progress reporting
- Tag cache (get_tags_from_cache/set_tags_to_cache) restoration
- dataflow_id correction in _load_dsl
- Language default Chinese, dead code guard removal
- embed_chunks made async with proper thread_pool_exec
- Full GraphRAG default configuration (10 parameters)
- Hardcoded q_768_vec fallback removal in RAPTOR

## Test Changes

- 20 new tests covering table parser manual mode, tag cache, embedding
edge cases, RAPTOR checkpoint, dataflow_id correction, storage binary
None, cancel cleanup, metadata=None boundary
- Unified `make_task_context`/`make_task_dict` factories eliminated 10+
duplicated helpers
- DataflowService tests migrated from internal method mocks to IO
boundary mocks (real orchestration code executes)
- Parametrized duplicate build_chunks post-processor tests
- 7 raptor tests modernized to @pytest.mark.asyncio
- Mock count per test reduced through boundary-level mocking strategy

**Test count: 400 passing, 0 warnings, 0 skips**

## Files Changed

| File | Change |
|------|--------|
| `rag/svr/task_executor.py` | +1 line (NullRecordingContext fix) |
| `rag/svr/task_executor_refactor/task_handler.py` | Orchestration
layer, 8 logic fixes |
| `rag/svr/task_executor_refactor/chunk_service.py` | +timeout +
None-check |
| `rag/svr/task_executor_refactor/embedding_service.py` | sync→async
rewrite |
| `rag/svr/task_executor_refactor/dataflow_service.py` | dataflow_id fix
+ timeout |
| `rag/svr/task_executor_refactor/raptor_service.py` | checkpoint fix +
assert |
| `rag/svr/task_executor_refactor/chunk_post_processor.py` | tag cache
restore |
| `rag/svr/task_executor_refactor/task_context.py` | language default
fix |
| `test/.../conftest.py` | +294 lines shared helpers |
| `test/.../*.py` | 15 test files refactored, 20 new tests |

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Jack
2026-06-03 17:18:31 +08:00
committed by GitHub
parent d736f358ba
commit b363146997
21 changed files with 1317 additions and 1222 deletions
Download Patch File Download Diff File Expand all files Collapse all files

296
test/unit_test/rag/svr/task_executor_refactor/conftest.py
View File

@@ -272,6 +272,7 @@ class MockChatModel:
def __init__(self):
self.llm_name = "mock_chat"
self.max_length = 4096
def __enter__(self):
return self
@@ -469,26 +470,293 @@ def mock_chunk_service_factory():
# =============================================================================
# RaptorService Fixtures
# Unified Mock TaskContext Factory
# =============================================================================
def make_task_context(**overrides):
"""Build a MagicMock TaskContext with sensible defaults for all services.
Every test file that needs a mock ``TaskContext`` should use this factory
with keyword-only overrides instead of defining its own ``_create_mock_context``.
Usage::
ctx = make_task_context(parser_id="table", kb_parser_config={"tag_kb_ids": ["kb_1"]})
"""
defaults = {
"id": "task_1",
"tenant_id": "tenant_1",
"kb_id": "kb_1",
"doc_id": "doc_1",
"name": "test.pdf",
"location": "/path/to/test.pdf",
"language": "en",
"parser_id": "naive",
"parser_config": {},
"kb_parser_config": {},
"llm_id": "llm_1",
"embd_id": "embd_1",
"from_page": 0,
"to_page": -1,
"size": 1000,
"pagerank": 0,
"task_type": "standard",
"dataflow_id": "",
"doc_ids": [],
"file": None,
"memory_id": "",
"source_id": "",
"message_dict": {},
}
ctx = MagicMock()
for k, v in defaults.items():
setattr(ctx, k, v if k not in overrides else overrides.pop(k))
# Callbacks
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.recording_context = MagicMock()
ctx.write_interceptor = None
# Raw task dict — derive from context attributes
ctx.raw_task = MagicMock()
# Limiters — all use AsyncMockLimiter so services that acquire them work
limiter = AsyncMockLimiter()
ctx.chunk_limiter = limiter
ctx.chat_limiter = limiter
ctx.embed_limiter = limiter
ctx.kg_limiter = limiter
ctx.minio_limiter = limiter
# Apply remaining overrides
for k, v in overrides.items():
setattr(ctx, k, v)
return ctx
# =============================================================================
# RaptorService Fixtures (kept for backward compatibility)
# =============================================================================
def create_mock_raptor_context():
"""Create a mock TaskContext suitable for RaptorService tests."""
ctx = MagicMock()
ctx.tenant_id = "tenant_1"
ctx.kb_id = "kb_1"
ctx.write_interceptor = None
ctx.progress_cb = MagicMock()
ctx.raw_task = {"type": ""}
ctx.parser_id = "naive"
ctx.parser_config = {}
ctx.name = "test.pdf"
ctx.pagerank = 0
ctx.id = "task_1"
return ctx
return make_task_context()
@pytest.fixture
def mock_raptor_context():
"""Provide a mock TaskContext for RaptorService tests."""
return create_mock_raptor_context()
return make_task_context()
# =============================================================================
# Embedding Binding Patch Helper
# =============================================================================
class patch_embedding_binding:
"""Context manager that patches embedding model binding at the external boundary.
Patches ``LLMBundle``, ``get_model_config_from_provider_instance``, and
``get_tenant_default_model_by_type`` so that ``TaskHandler._bind_embedding_model``
executes its real logic without making actual API calls.
Usage::
with patch_embedding_binding(vector_size=128):
handler = TaskHandler(ctx)
await handler.handle()
"""
def __init__(self, vector_size: int = 128):
self._vector_size = vector_size
self._patches = []
def __enter__(self):
mock_model = MagicMock()
mock_model.encode = MagicMock(
return_value=(
np.random.rand(1, self._vector_size).astype(np.float32),
10,
)
)
mock_model.max_length = 512
mock_model.llm_name = "mock_embedding"
mock_model.__enter__ = MagicMock(return_value=mock_model)
mock_model.__exit__ = MagicMock(return_value=False)
self._patches = [
patch(
"rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance",
return_value=MagicMock(),
),
patch(
"rag.svr.task_executor_refactor.task_handler.LLMBundle",
return_value=mock_model,
),
patch(
"rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type",
return_value=MagicMock(),
),
]
for p in self._patches:
p.__enter__()
return self
def __exit__(self, *args):
for p in reversed(self._patches):
p.__exit__(*args)
# =============================================================================
# Common mock callbacks
# =============================================================================
async def mock_thread_return_binary(func, *args, **kwargs):
"""Reusable mock for thread_pool_exec — returns fake binary."""
return b"fake pdf binary"
async def mock_thread_return_none(func, *args, **kwargs):
"""Reusable mock for thread_pool_exec — returns None."""
return None
# =============================================================================
# Patch helpers for integration tests
# =============================================================================
def patch_get_storage_binary():
"""Patch TaskHandler._get_storage_binary to return fake binary."""
return patch("rag.svr.task_executor_refactor.task_handler.TaskHandler._get_storage_binary",
new_callable=AsyncMock, return_value=b"fake pdf binary")
def patch_task_handler_settings(mock_settings):
"""Patch the settings module-level import in task_handler."""
return patch("rag.svr.task_executor_refactor.task_handler.settings", mock_settings)
# =============================================================================
# Shared Task Dictionary Factory
# =============================================================================
def make_task_dict(**overrides):
"""Build a task dict with sensible defaults for integration tests.
All ``_create_standard_task_dict`` / ``_create_raptor_task_dict`` / etc.
helpers in integration tests should be replaced with this single factory.
Usage::
task_dict = make_task_dict(task_type="raptor", doc_ids=["doc1"])
"""
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": "doc_test",
"name": "test_document.pdf",
"location": "/path/to/test_document.pdf",
"size": 1024,
"parser_id": "naive",
"parser_config": {"auto_keywords": 0, "auto_questions": 0, "enable_metadata": False},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "standard",
"pagerank": 0,
**overrides,
}
# =============================================================================
# Shared Pipeline Mock Block for Integration Tests
# =============================================================================
class patch_pipeline_mocks:
"""Context manager bundling common integration-test mock blocks.
Patches external boundaries so ``TaskHandler.handle()`` executes without
actual API calls. Use ``mode="raptor"`` or ``mode="graphrag"``.
Usage::
with patch_pipeline_mocks() as m:
m.get_model_config_from_provider_instance.return_value = MagicMock()
handler = TaskHandler(ctx)
await handler.handle()
"""
_MODULES = {
"task_handler": "rag.svr.task_executor_refactor.task_handler",
"chunk_service": "rag.svr.task_executor_refactor.chunk_service",
}
# (module_key, attr_name, use_AsyncMock)
_COMMON = [
("task_handler", "get_model_config_from_provider_instance", False),
("task_handler", "LLMBundle", False),
("task_handler", "get_tenant_default_model_by_type", False),
("task_handler", "search.index_name", False),
("task_handler", "thread_pool_exec", False),
("task_handler", "DocumentService", False),
]
_STANDARD = [
("task_handler", "File2DocumentService", False),
("chunk_service", "thread_pool_exec", False),
("task_handler", "ChunkService", False),
]
_RAPTOR = [
("task_handler", "KnowledgebaseService", False),
("task_handler", "RaptorService", False),
("task_handler", "ChunkService", False),
]
_GRAPH_RAG = [
("task_handler", "KnowledgebaseService", False),
("task_handler", "run_graphrag_for_kb", True),
]
def __init__(self, mode: str = "standard"):
self._mode = mode
self._stack = None
def __enter__(self):
import contextlib
from unittest.mock import patch, MagicMock, AsyncMock
prefixes = list(self._COMMON)
if self._mode == "standard":
prefixes += self._STANDARD
elif self._mode == "raptor":
prefixes += self._RAPTOR
elif self._mode == "graphrag":
prefixes += self._GRAPH_RAG
mocks = MagicMock()
ctx_managers = []
for mod_key, attr, use_async in prefixes:
target = f"{self._MODULES[mod_key]}.{attr}"
if use_async:
cm = patch(target, new_callable=AsyncMock)
else:
cm = patch(target)
mock_handle = cm.__enter__()
setattr(mocks, attr.replace(".", "_"), mock_handle)
ctx_managers.append(cm)
self._stack = contextlib.ExitStack()
self._ctx_managers = ctx_managers
return mocks
def __exit__(self, *args):
for cm in reversed(self._ctx_managers):
cm.__exit__(*args)

103
test/unit_test/rag/svr/task_executor_refactor/test_chunk_builder.py
View File

@@ -18,12 +18,13 @@ Unit tests for ChunkBuilder module.
"""
import pytest
from unittest.mock import MagicMock, patch, AsyncMock
from unittest.mock import MagicMock, patch
from rag.svr.task_executor_refactor.chunk_builder import (
get_parser,
run_chunking,
extract_outline,
)
from test.unit_test.rag.svr.task_executor_refactor.conftest import make_task_context
class TestGetParser:
@@ -49,151 +50,99 @@ class TestGetParser:
class TestRunChunking:
"""Tests for run_chunking function."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.name = "test.pdf"
ctx.location = "/path/to/test.pdf"
ctx.from_page = 0
ctx.to_page = -1
ctx.language = "en"
ctx.kb_id = "kb_1"
ctx.parser_config = {}
ctx.tenant_id = "tenant_1"
ctx.progress_cb = MagicMock()
ctx.raw_task = {}
ctx.chunk_limiter = MagicMock()
ctx.chunk_limiter.__aenter__ = AsyncMock()
ctx.chunk_limiter.__aexit__ = AsyncMock()
return ctx
@pytest.mark.asyncio
async def test_run_chunking_success(self):
"""Test successful chunking."""
ctx = self._create_mock_context()
ctx = make_task_context()
mock_chunker = MagicMock()
mock_chunker.chunk = MagicMock(return_value=[{"content_with_weight": "chunk1"}])
with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread:
# thread_pool_exec returns an awaitable that returns the list
mock_thread.return_value = [{"content_with_weight": "chunk1"}]
result = await run_chunking(mock_chunker, b"binary", ctx)
assert result is not None
assert len(result) == 1
@pytest.mark.asyncio
async def test_run_chunking_with_parser_config(self):
"""Test chunking merges table parser config."""
ctx = self._create_mock_context()
ctx = make_task_context()
ctx.raw_task = {"parser_config": {"chunk_token_num": 128}}
mock_chunker = MagicMock()
mock_chunker.chunk = MagicMock(return_value=[])
with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread:
with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread, \
patch("rag.svr.task_executor_refactor.chunk_builder.merge_table_parser_config_from_kb") as mock_merge:
mock_thread.return_value = []
with patch("rag.svr.task_executor_refactor.chunk_builder.merge_table_parser_config_from_kb") as mock_merge:
mock_merge.return_value = {"chunk_token_num": 128}
await run_chunking(mock_chunker, b"binary", ctx)
mock_merge.assert_called_once_with(ctx.raw_task)
mock_merge.return_value = {"chunk_token_num": 128}
await run_chunking(mock_chunker, b"binary", ctx)
mock_merge.assert_called_once_with(ctx.raw_task)
@pytest.mark.asyncio
async def test_run_chunking_exception(self):
"""Test chunking handles exception."""
ctx = self._create_mock_context()
ctx = make_task_context()
mock_chunker = MagicMock()
mock_chunker.chunk = MagicMock(side_effect=Exception("Test error"))
with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread:
mock_thread.side_effect = Exception("Test error")
with pytest.raises(Exception):
await run_chunking(mock_chunker, b"binary", ctx)
# Verify progress_cb was called with error message
ctx.progress_cb.assert_called()
class TestExtractOutline:
"""Tests for extract_outline function."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.doc_id = "doc_1"
ctx.write_interceptor = None
ctx.progress_cb = MagicMock()
@staticmethod
def _ctx(recording_ctx=None, **overrides):
ctx = make_task_context(**overrides)
ctx.recording_context = recording_ctx or MagicMock()
return ctx
@pytest.mark.asyncio
async def test_extract_outline_with_data(self):
"""Test outline extraction when outline data is present."""
ctx = self._create_mock_context()
ctx = self._ctx()
outline_data = [{"title": "Chapter 1", "page": 1}]
cks = [{"__outline__": outline_data}]
mock_rec_ctx = MagicMock()
ctx.recording_context = mock_rec_ctx
with patch("rag.svr.task_executor_refactor.chunk_builder.DocMetadataService") as mock_meta:
mock_meta.get_document_metadata.return_value = {}
mock_meta.update_document_metadata = MagicMock()
await extract_outline(cks, ctx)
mock_rec_ctx.record.assert_called_with("outline_data", outline_data)
# Outline should be popped from first chunk
ctx.recording_context.record.assert_called_with("outline_data", outline_data)
assert "__outline__" not in cks[0]
mock_meta.update_document_metadata.assert_called_once()
@pytest.mark.asyncio
async def test_extract_outline_without_data(self):
"""Test outline extraction when no outline data."""
ctx = self._create_mock_context()
ctx = self._ctx()
cks = [{"content_with_weight": "test"}]
mock_rec_ctx = MagicMock()
ctx.recording_context = mock_rec_ctx
await extract_outline(cks, ctx)
mock_rec_ctx.record.assert_called_with("outline_data", None)
ctx.recording_context.record.assert_called_with("outline_data", None)
@pytest.mark.asyncio
async def test_extract_outline_empty_chunks(self):
"""Test outline extraction with empty chunks list."""
ctx = self._create_mock_context()
mock_rec_ctx = MagicMock()
ctx.recording_context = mock_rec_ctx
ctx = self._ctx()
await extract_outline([], ctx)
mock_rec_ctx.record.assert_called_with("outline_data", None)
ctx.recording_context.record.assert_called_with("outline_data", None)
@pytest.mark.asyncio
async def test_extract_outline_with_write_interceptor(self):
"""Test outline extraction with write interceptor."""
ctx = self._create_mock_context()
ctx.write_interceptor = MagicMock()
ctx = self._ctx(write_interceptor=MagicMock())
outline_data = [{"title": "Chapter 1", "page": 1}]
cks = [{"__outline__": outline_data}]
mock_rec_ctx = MagicMock()
ctx.recording_context = mock_rec_ctx
await extract_outline(cks, ctx)
ctx.write_interceptor.intercept.assert_called_once_with(
"DocMetadataService.update_document_metadata"
)
@@ -201,19 +150,13 @@ class TestExtractOutline:
@pytest.mark.asyncio
async def test_extract_outline_persistence_exception(self):
"""Test outline extraction handles persistence exception."""
ctx = self._create_mock_context()
ctx = self._ctx()
outline_data = [{"title": "Chapter 1", "page": 1}]
cks = [{"__outline__": outline_data}]
mock_rec_ctx = MagicMock()
ctx.recording_context = mock_rec_ctx
with patch("rag.svr.task_executor_refactor.chunk_builder.DocMetadataService") as mock_meta:
mock_meta.get_document_metadata.return_value = {}
mock_meta.update_document_metadata.side_effect = Exception("DB error")
# Should not raise exception, just log warning
await extract_outline(cks, ctx)
mock_rec_ctx.record.assert_called_with("outline_data", outline_data)
ctx.recording_context.record.assert_called_with("outline_data", outline_data)

524
test/unit_test/rag/svr/task_executor_refactor/test_chunk_post_processor.py
View File

@@ -15,10 +15,16 @@
"""
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, AsyncMock
from unittest.mock import MagicMock, patch
from rag.svr.task_executor_refactor.chunk_post_processor import (
extract_keywords,
generate_questions,
@@ -27,394 +33,314 @@ from rag.svr.task_executor_refactor.chunk_post_processor import (
count_with_key,
build_metadata_config,
)
from test.unit_test.rag.svr.task_executor_refactor.conftest import (
make_task_context,
MockChatModel,
)
class TestExtractKeywords:
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."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.tenant_id = "tenant_1"
ctx.llm_id = "llm_1"
ctx.language = "en"
ctx.parser_config = {"auto_keywords": 5}
ctx.id = "task_1"
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.chat_limiter = MagicMock()
ctx.chat_limiter.__aenter__ = AsyncMock()
ctx.chat_limiter.__aexit__ = AsyncMock()
return ctx
@pytest.mark.asyncio
async def test_extract_keywords_success(self):
"""Test successful keyword extraction."""
ctx = self._create_mock_context()
"""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"},
]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = "keyword1, keyword2"
with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
with patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer") as mock_tokenizer:
mock_tokenizer.tokenize.return_value = "keyword1 keyword2"
await extract_keywords(docs, ctx)
# Verify keywords were set
assert "important_kwd" in docs[0]
assert "important_tks" in docs[0]
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 = self._create_mock_context()
ctx.has_canceled_func = MagicMock(return_value=True)
ctx = make_task_context(parser_config={"auto_keywords": 5},
has_canceled_func=MagicMock(return_value=True))
docs = [{"content_with_weight": "This is test content"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = None # No cache
await extract_keywords(docs, ctx)
# Should return early due to cancellation
assert "important_kwd" not in docs[0]
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 = self._create_mock_context()
ctx = make_task_context(parser_config={"auto_keywords": 5})
docs = []
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
await extract_keywords(docs, ctx)
# Should complete without error
ctx.progress_cb.assert_called()
p1, p2 = self._mock_llm_binding()
with p1, p2:
await extract_keywords(docs, ctx)
ctx.progress_cb.assert_called()
class TestGenerateQuestions:
class TestGenerateQuestions(_BasePostProcessorTest):
"""Tests for generate_questions function."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.tenant_id = "tenant_1"
ctx.llm_id = "llm_1"
ctx.language = "en"
ctx.parser_config = {"auto_questions": 3}
ctx.id = "task_1"
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.chat_limiter = MagicMock()
ctx.chat_limiter.__aenter__ = AsyncMock()
ctx.chat_limiter.__aexit__ = AsyncMock()
return ctx
@pytest.mark.asyncio
async def test_generate_questions_success(self):
"""Test successful question generation."""
ctx = self._create_mock_context()
docs = [
{"content_with_weight": "This is test content one"},
]
"""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"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = "Question 1\nQuestion 2"
with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
with patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer") as mock_tokenizer:
mock_tokenizer.tokenize.return_value = "Question 1 Question 2"
await generate_questions(docs, ctx)
# Verify questions were set
assert "question_kwd" in docs[0]
assert "question_tks" in docs[0]
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 = self._create_mock_context()
ctx.has_canceled_func = MagicMock(return_value=True)
ctx = make_task_context(parser_config={"auto_questions": 3},
has_canceled_func=MagicMock(return_value=True))
docs = [{"content_with_weight": "This is test content"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = None # No cache
await generate_questions(docs, ctx)
# Should return early due to cancellation
assert "question_kwd" not in docs[0]
p1, p2 = self._mock_llm_binding()
with p1, p2:
await generate_questions(docs, ctx)
assert "question_kwd" not in docs[0]
class TestGenerateMetadata:
class TestGenerateMetadata(_BasePostProcessorTest):
"""Tests for generate_metadata function."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.tenant_id = "tenant_1"
ctx.llm_id = "llm_1"
ctx.language = "en"
ctx.parser_config = {
"enable_metadata": True,
"metadata": [{"name": "category", "type": "string"}],
"built_in_metadata": ["author", "date"],
}
ctx.doc_id = "doc_1"
ctx.id = "task_1"
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.write_interceptor = None
ctx.chat_limiter = MagicMock()
ctx.chat_limiter.__aenter__ = AsyncMock()
ctx.chat_limiter.__aexit__ = AsyncMock()
return ctx
@pytest.mark.asyncio
async def test_generate_metadata_success(self):
"""Test successful metadata generation."""
ctx = self._create_mock_context()
docs = [
{"content_with_weight": "This is test content", "metadata_obj": {"category": "test"}},
]
"""Test successful metadata generation — cache miss → LLM prompt runs."""
ctx = make_task_context(
parser_config={
"enable_metadata": True,
"metadata": [{"name": "category", "type": "string"}],
"built_in_metadata": ["author", "date"],
},
)
docs = [{"content_with_weight": "This is test content"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = {"category": "test"}
with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
with patch("rag.svr.task_executor_refactor.chunk_post_processor.update_metadata_to") as mock_update:
mock_update.return_value = {"category": "test"}
with patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService") as mock_meta:
mock_meta.get_document_metadata.return_value = {}
mock_meta.update_document_metadata = MagicMock()
await generate_metadata(docs, ctx)
# Verify metadata_obj was processed
mock_meta.update_document_metadata.assert_called_once()
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.gen_metadata",
return_value={"category": "test"},
)
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
with p1, p2, p3, p4, p5, p6 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 = self._create_mock_context()
ctx.write_interceptor = MagicMock()
docs = [
{"content_with_weight": "This is test content", "metadata_obj": {"category": "test"}},
]
ctx = make_task_context(
parser_config={
"enable_metadata": True,
"metadata": [{"name": "category", "type": "string"}],
"built_in_metadata": ["author", "date"],
},
write_interceptor=MagicMock(),
)
docs = [{"content_with_weight": "This is test content"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = {"category": "test"}
with patch("rag.svr.task_executor_refactor.chunk_post_processor.update_metadata_to") as mock_update:
mock_update.return_value = {"category": "test"}
with patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService") as mock_meta:
mock_meta.get_document_metadata.return_value = {}
mock_meta.update_document_metadata = MagicMock()
await generate_metadata(docs, ctx)
ctx.write_interceptor.intercept.assert_called_once_with(
"DocMetadataService.update_document_metadata"
)
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.gen_metadata",
return_value={"category": "test"},
)
p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService")
with p1, p2, p3, p4, p5, p6:
await generate_metadata(docs, ctx)
ctx.write_interceptor.intercept.assert_called_once_with(
"DocMetadataService.update_document_metadata"
)
class TestApplyTags:
class TestApplyTags(_BasePostProcessorTest):
"""Tests for apply_tags function."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.tenant_id = "tenant_1"
ctx.llm_id = "llm_1"
ctx.language = "en"
ctx.kb_parser_config = {"tag_kb_ids": ["kb_1"], "topn_tags": 3}
ctx.id = "task_1"
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.chat_limiter = MagicMock()
ctx.chat_limiter.__aenter__ = AsyncMock()
ctx.chat_limiter.__aexit__ = AsyncMock()
return ctx
@pytest.mark.asyncio
async def test_apply_tags_success(self):
"""Test successful tag application."""
ctx = self._create_mock_context()
docs = [
{"content_with_weight": "This is test content"},
]
"""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"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") as mock_settings:
mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10, "tag2": 5}
mock_settings.retriever.tag_content.return_value = True
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
mock_cache.return_value = '{"tag1": 1}'
with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
await apply_tags(docs, ctx)
# Verify tags were applied
assert len(docs) == 1
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 = self._create_mock_context()
ctx.has_canceled_func = MagicMock(return_value=True)
docs = [
{"content_with_weight": "This is test content"},
]
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"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") as mock_config:
mock_config.return_value = MagicMock()
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)
with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
mock_llm_instance = MagicMock()
mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
mock_llm.return_value.__exit__ = MagicMock(return_value=False)
@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"}]
with patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") as mock_settings:
mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10}
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()
await apply_tags(docs, ctx)
@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"}]
# Should return early due to cancellation
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):
"""Test counting when all docs have the key."""
docs = [{"tag": 1}, {"tag": 2}, {"tag": 3}]
result = count_with_key(docs, "tag")
assert result == 3
assert count_with_key(docs, "tag") == 3
def test_count_with_key_some_have_key(self):
"""Test counting when some docs have the key."""
docs = [{"tag": 1}, {"other": 2}, {"tag": 3}]
result = count_with_key(docs, "tag")
assert result == 2
assert count_with_key(docs, "tag") == 2
def test_count_with_key_none_have_key(self):
"""Test counting when no docs have the key."""
docs = [{"other": 1}, {"other": 2}]
result = count_with_key(docs, "tag")
assert result == 0
assert count_with_key(docs, "tag") == 0
def test_count_with_key_empty_docs(self):
"""Test counting with empty docs list."""
result = count_with_key([], "tag")
assert result == 0
assert count_with_key([], "tag") == 0
def test_count_with_key_falsy_value(self):
"""Test counting when key exists but has falsy value."""
docs = [{"tag": 0}, {"tag": ""}, {"tag": None}]
result = count_with_key(docs, "tag")
# Falsy values should not be counted (since d.get(key) returns falsy)
assert result == 0
assert count_with_key(docs, "tag") == 0
def test_count_with_key_truthy_value(self):
"""Test counting when key has truthy value."""
docs = [{"tag": 1}, {"tag": "value"}, {"tag": [1, 2]}]
result = count_with_key(docs, "tag")
assert result == 3
assert count_with_key(docs, "tag") == 3
class TestBuildMetadataConfig:
"""Tests for build_metadata_config function."""
def test_dict_without_properties_returns_schema(self):
"""When metadata is a dict without properties, return {type: object, properties: {}}."""
parser_config = {"metadata": {"type": "object"}, "built_in_metadata": []}
result = build_metadata_config(parser_config)
assert result == {"type": "object", "properties": {}}
assert build_metadata_config(parser_config) == {"type": "object", "properties": {}}
def test_dict_with_properties_and_built_in(self):
"""When metadata is a dict with properties AND built_in_metadata, merge them."""
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 result["type"] == "object"
assert "a" in result["properties"]
assert "author" in result["properties"]
def test_dict_with_properties_no_built_in(self):
"""When metadata is a dict with properties and no built_in, return as-is."""
parser_config = {
"metadata": {"type": "object", "properties": {"a": {"type": "string"}}},
"built_in_metadata": [],
@@ -423,38 +349,32 @@ class TestBuildMetadataConfig:
assert result == {"type": "object", "properties": {"a": {"type": "string"}}}
def test_list_with_built_in(self):
"""When metadata is a list and built_in_metadata is present, concatenate."""
parser_config = {
"metadata": [{"key": "category"}],
"built_in_metadata": [{"key": "author"}],
}
result = build_metadata_config(parser_config)
assert result == [{"key": "category"}, {"key": "author"}]
assert build_metadata_config(parser_config) == [{"key": "category"}, {"key": "author"}]
def test_list_without_built_in(self):
"""When metadata is a list and built_in_metadata is empty, return metadata as-is."""
parser_config = {"metadata": [{"key": "category"}], "built_in_metadata": []}
result = build_metadata_config(parser_config)
assert result == [{"key": "category"}]
assert build_metadata_config(parser_config) == [{"key": "category"}]
def test_other_type_with_built_in(self):
"""When metadata is not dict or list (empty list), return built_in_metadata only."""
parser_config = {"metadata": [], "built_in_metadata": [{"key": "author"}]}
result = build_metadata_config(parser_config)
assert result == [{"key": "author"}]
assert build_metadata_config(parser_config) == [{"key": "author"}]
def test_idempotent_same_input(self):
"""Same input produces structurally equal results."""
parser_config = {
"metadata": [{"key": "category"}],
"built_in_metadata": [{"key": "author"}],
}
result1 = build_metadata_config(parser_config)
result2 = build_metadata_config(parser_config)
assert result1 == result2
assert build_metadata_config(parser_config) == build_metadata_config(parser_config)
def test_missing_metadata_key(self):
"""When parser_config has no 'metadata' key, built_in_metadata alone is returned."""
parser_config = {"built_in_metadata": []}
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 == []
assert result == [{"key": "author"}]

199
test/unit_test/rag/svr/task_executor_refactor/test_chunk_service.py
View File

@@ -29,6 +29,7 @@ This test file now focuses on ChunkService-specific functionality:
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:
@@ -44,40 +45,11 @@ class TestChunkServiceInit:
class TestChunkServiceBuildChunks:
"""Tests for build_chunks method."""
def _create_mock_context(self, parser_id="naive", size=1000, parser_config=None, kb_parser_config=None):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.parser_id = parser_id
ctx.name = "test.pdf"
ctx.size = size
ctx.from_page = 0
ctx.to_page = -1
ctx.parser_config = parser_config or {}
ctx.kb_parser_config = kb_parser_config or {}
ctx.language = "en"
ctx.id = "task_1"
ctx.tenant_id = "tenant_1"
ctx.kb_id = "kb_1"
ctx.doc_id = "doc_1"
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.write_interceptor = None
ctx.raw_task = {}
ctx.llm_id = "llm_1"
ctx.pagerank = 0
ctx.location = "/path/to/test.pdf"
ctx.chunk_limiter = MagicMock()
ctx.chunk_limiter.__aenter__ = AsyncMock()
ctx.chunk_limiter.__aexit__ = AsyncMock()
ctx.chat_limiter = MagicMock()
ctx.chat_limiter.__aenter__ = AsyncMock()
ctx.chat_limiter.__aexit__ = AsyncMock()
return ctx
@pytest.mark.asyncio
async def test_build_chunks_file_size_exceeded(self):
"""Test build_chunks returns empty list when file size exceeds limit."""
ctx = self._create_mock_context(size=1000000000) # Very large size
ctx = make_task_context(size=1000000000) # Very large size
service = ChunkService(ctx=ctx)
@@ -95,7 +67,7 @@ class TestChunkServiceBuildChunks:
@pytest.mark.asyncio
async def test_build_chunks_file_size_ok(self):
"""Test build_chunks proceeds when file size is within limit."""
ctx = self._create_mock_context(size=1000)
ctx = make_task_context(size=1000)
service = ChunkService(ctx=ctx)
@@ -123,138 +95,43 @@ class TestChunkServiceBuildChunks:
mock_get_parser.assert_called_once_with("naive")
@pytest.mark.asyncio
async def test_build_chunks_with_auto_keywords(self):
"""Test build_chunks triggers keyword extraction when configured."""
ctx = self._create_mock_context(parser_config={"auto_keywords": 5})
@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)
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
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_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_get_parser.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
mock_run_chunking.return_value = []
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"}]
with patch("rag.svr.task_executor_refactor.chunk_service.extract_keywords", new_callable=AsyncMock) as mock_extract:
await service.build_chunks(b"test binary")
mock_extract.assert_called_once()
@pytest.mark.asyncio
async def test_build_chunks_with_auto_questions(self):
"""Test build_chunks triggers question generation when configured."""
ctx = self._create_mock_context(parser_config={"auto_questions": 3})
service = ChunkService(ctx=ctx)
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
mock_settings.DOC_MAXIMUM_SIZE = 10000000
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_get_parser.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
mock_run_chunking.return_value = []
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"}]
with patch("rag.svr.task_executor_refactor.chunk_service.generate_questions", new_callable=AsyncMock) as mock_gen:
await service.build_chunks(b"test binary")
mock_gen.assert_called_once()
@pytest.mark.asyncio
async def test_build_chunks_with_tag_kb_ids(self):
"""Test build_chunks triggers tag application when tag_kb_ids configured."""
ctx = self._create_mock_context(kb_parser_config={"tag_kb_ids": ["kb_1"]})
service = ChunkService(ctx=ctx)
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
mock_settings.DOC_MAXIMUM_SIZE = 10000000
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_get_parser.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
mock_run_chunking.return_value = []
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"}]
with patch("rag.svr.task_executor_refactor.chunk_service.apply_tags", new_callable=AsyncMock) as mock_apply:
await service.build_chunks(b"test binary")
mock_apply.assert_called_once()
@pytest.mark.asyncio
async def test_build_chunks_with_metadata(self):
"""Test build_chunks triggers metadata generation when configured."""
ctx = self._create_mock_context(
parser_config={
"enable_metadata": True,
"metadata": [{"name": "category", "type": "string"}]
}
)
service = ChunkService(ctx=ctx)
with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
mock_settings.DOC_MAXIMUM_SIZE = 10000000
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_get_parser.return_value = MagicMock()
with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
mock_run_chunking.return_value = []
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"}]
with patch("rag.svr.task_executor_refactor.chunk_service.generate_metadata", new_callable=AsyncMock) as mock_meta:
await service.build_chunks(b"test binary")
mock_meta.assert_called_once()
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."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.doc_id = "doc_1"
ctx.kb_id = "kb_1"
ctx.tenant_id = "tenant_1"
ctx.name = "test.pdf"
ctx.location = "/path/to/test.pdf"
ctx.pagerank = 0
ctx.progress_cb = MagicMock()
return ctx
@pytest.mark.asyncio
async def test_prepare_docs_and_upload_basic(self):
"""Test basic document preparation."""
ctx = self._create_mock_context()
ctx = make_task_context()
service = ChunkService(ctx=ctx)
cks = [{"content_with_weight": "test chunk"}]
@@ -274,7 +151,7 @@ class TestChunkServicePrepareDocsAndUpload:
@pytest.mark.asyncio
async def test_prepare_docs_and_upload_with_pagerank(self):
"""Test document preparation with pagerank."""
ctx = self._create_mock_context()
ctx = make_task_context()
ctx.pagerank = 5
service = ChunkService(ctx=ctx)
@@ -293,23 +170,11 @@ class TestChunkServicePrepareDocsAndUpload:
class TestChunkServiceInsertChunks:
"""Tests for insert_chunks method."""
def _create_mock_context(self):
"""Helper to create a mock TaskContext."""
ctx = MagicMock()
ctx.id = "task_1"
ctx.tenant_id = "tenant_1"
ctx.kb_id = "kb_1"
ctx.doc_id = "doc_1"
ctx.parser_id = "naive"
ctx.progress_cb = MagicMock()
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.write_interceptor = None
return ctx
@pytest.mark.asyncio
async def test_insert_chunks_success(self):
"""Test successful chunk insertion."""
ctx = self._create_mock_context()
ctx = make_task_context()
service = ChunkService(ctx=ctx)
chunks = [
@@ -338,7 +203,7 @@ class TestChunkServiceInsertChunks:
@pytest.mark.asyncio
async def test_insert_chunks_canceled(self):
"""Test chunk insertion when task is canceled."""
ctx = self._create_mock_context()
ctx = make_task_context()
ctx.has_canceled_func = MagicMock(return_value=True)
service = ChunkService(ctx=ctx)
@@ -363,7 +228,7 @@ class TestChunkServiceInsertChunks:
@pytest.mark.asyncio
async def test_insert_chunks_doc_store_error(self):
"""Test chunk insertion when doc store returns error."""
ctx = self._create_mock_context()
ctx = make_task_context()
service = ChunkService(ctx=ctx)
chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]

15
test/unit_test/rag/svr/task_executor_refactor/test_comparator.py
View File

@@ -138,25 +138,18 @@ class TestComparisonReport:
assert "No comparison details" in md
class TestContextComparatorInit:
"""Tests for ContextComparator initialization."""
class TestContextComparatorCompareValue:
"""Tests for ContextComparator.compare_value method and initialization."""
def test_init_default_tolerance(self):
"""Test initialization with default tolerance."""
comparator = ContextComparator()
assert comparator.float_tolerance == 1e-6
assert ContextComparator().float_tolerance == 1e-6
def test_init_custom_tolerance(self):
"""Test initialization with custom tolerance."""
comparator = ContextComparator(float_tolerance=0.01)
assert comparator.float_tolerance == 0.01
class TestContextComparatorCompareValue:
"""Tests for ContextComparator.compare_value method."""
assert ContextComparator(float_tolerance=0.01).float_tolerance == 0.01
def setup_method(self):
"""Set up test fixtures."""
self.comparator = ContextComparator()
def test_compare_none_values(self):

140
test/unit_test/rag/svr/task_executor_refactor/test_dataflow_service.py
View File

@@ -66,10 +66,11 @@ class TestDataflowServiceRunDataflow:
await service.run_dataflow()
@pytest.mark.asyncio
@pytest.mark.parametrize("output_key", ["chunks", "json"])
@patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
@patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
async def test_run_dataflow_with_chunks_output(self, mock_canvas, mock_pipeline_class, task_context):
"""Test run_dataflow processes 'chunks' output type end-to-end."""
async def test_run_dataflow_with_output_type(self, mock_canvas, mock_pipeline_class, task_context, output_key):
"""Test run_dataflow processes output end-to-end (chunks / json)."""
task_context._task["task_type"] = "dataflow"
task_context._task["dataflow_id"] = "dataflow_test"
task_context._task["tenant_id"] = "tenant_test"
@@ -79,59 +80,22 @@ class TestDataflowServiceRunDataflow:
task_context._write_interceptor = None
mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
chunks = {
"chunks": [
{"text": "Hello world", "content_with_weight": "Hello world"},
],
"embedding_token_consumption": 5,
}
data = {output_key: [{"text": "content", "content_with_weight": "content"}]}
data["embedding_token_consumption"] = 5
mock_pipeline = MagicMock()
mock_pipeline.run = AsyncMock(return_value=chunks)
mock_pipeline.run = AsyncMock(return_value=data)
mock_pipeline_class.return_value = mock_pipeline
# Patch internal heavy dependencies so run_dataflow completes
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["chunks"], 5)):
with patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True):
with patch.object(DataflowService, '_update_document_metadata'):
with patch.object(DataflowService, '_record_pipeline_log'):
with patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
service = DataflowService(ctx=task_context)
await service.run_dataflow()
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock,
return_value=(data[output_key], 5)), \
patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True), \
patch.object(DataflowService, '_update_document_metadata'), \
patch.object(DataflowService, '_record_pipeline_log'), \
patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
# Verify chunks were inserted
DataflowService._insert_chunks.assert_called_once()
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
@patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
async def test_run_dataflow_with_json_output(self, mock_canvas, mock_pipeline_class, task_context):
"""Test run_dataflow processes 'json' output type."""
task_context._task["task_type"] = "dataflow"
task_context._task["dataflow_id"] = "dataflow_test"
task_context._task["tenant_id"] = "tenant_test"
task_context._task["kb_id"] = "kb_test"
task_context._task["doc_id"] = "doc_test"
task_context._task["name"] = "test.pdf"
task_context._write_interceptor = None
mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
chunks = {
"json": [
{"text": "JSON content"},
],
"embedding_token_consumption": 2,
}
mock_pipeline = MagicMock()
mock_pipeline.run = AsyncMock(return_value=chunks)
mock_pipeline_class.return_value = mock_pipeline
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["json"], 2)):
with patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True):
with patch.object(DataflowService, '_update_document_metadata'):
with patch.object(DataflowService, '_record_pipeline_log'):
with patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
service = DataflowService(ctx=task_context)
await service.run_dataflow()
service = DataflowService(ctx=task_context)
await service.run_dataflow()
DataflowService._insert_chunks.assert_called_once()
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
@@ -154,13 +118,11 @@ class TestDataflowServiceRunDataflow:
mock_pipeline.run = AsyncMock(return_value=chunks)
mock_pipeline_class.return_value = mock_pipeline
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(None, 0)):
with patch.object(DataflowService, '_record_pipeline_log'):
service = DataflowService(ctx=task_context)
await service.run_dataflow()
# Should not insert chunks when embedding fails
service._record_pipeline_log.assert_called()
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(None, 0)), \
patch.object(DataflowService, '_record_pipeline_log'):
service = DataflowService(ctx=task_context)
await service.run_dataflow()
service._record_pipeline_log.assert_called()
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
@@ -190,16 +152,16 @@ class TestDataflowServiceRunDataflow:
billing_hook.on_pipeline_success = AsyncMock()
billing_hook.on_pipeline_error = AsyncMock()
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["chunks"], 1)):
with patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True):
with patch.object(DataflowService, '_update_document_metadata'):
with patch.object(DataflowService, '_record_pipeline_log'):
with patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
service = DataflowService(ctx=task_context, billing_hook=billing_hook)
await service.run_dataflow()
with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["chunks"], 1)), \
patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True), \
patch.object(DataflowService, '_update_document_metadata'), \
patch.object(DataflowService, '_record_pipeline_log'), \
patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
billing_hook.on_pipeline_success.assert_called_once()
billing_hook.on_pipeline_error.assert_not_called()
service = DataflowService(ctx=task_context, billing_hook=billing_hook)
await service.run_dataflow()
billing_hook.on_pipeline_success.assert_called_once()
billing_hook.on_pipeline_error.assert_not_called()
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
@@ -378,4 +340,46 @@ class TestDataflowServiceInit:
hook = MagicMock()
service = DataflowService(ctx=ctx, billing_hook=hook)
assert service._task_context is ctx
assert service._billing_hook is hook
assert service._billing_hook is hook
class TestDataflowServiceLoadDsl:
"""Tests for _load_dsl with dataflow_id correction."""
@pytest.mark.asyncio
async def test_load_dsl_for_dataflow_task_type_returns_unchanged_id(self):
"""When task_type == 'dataflow', dataflow_id is returned unchanged."""
ctx = MagicMock()
ctx.task_type = "dataflow"
dataflow_id = "original_dataflow_id"
with patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService") as mock_canvas:
mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
service = DataflowService(ctx=ctx)
dsl, corrected_id = await service._load_dsl(dataflow_id)
assert dsl == '{"id": "test"}'
assert corrected_id == "original_dataflow_id"
mock_canvas.get_by_id.assert_called_once_with(dataflow_id)
@pytest.mark.asyncio
async def test_load_dsl_for_pipeline_log_task_type_returns_corrected_id(self):
"""When task_type != 'dataflow', dataflow_id comes from pipeline_log.pipeline_id."""
ctx = MagicMock()
ctx.task_type = "raptor"
dataflow_id = "pipeline_log_id"
with patch("rag.svr.task_executor_refactor.dataflow_service.PipelineOperationLogService") as mock_log:
mock_log_instance = MagicMock()
mock_log_instance.dsl = '{"id": "test_pipeline"}'
mock_log_instance.pipeline_id = "corrected_pipeline_id"
mock_log.get_by_id.return_value = (True, mock_log_instance)
service = DataflowService(ctx=ctx)
dsl, corrected_id = await service._load_dsl(dataflow_id)
assert dsl == '{"id": "test_pipeline"}'
assert corrected_id == "corrected_pipeline_id"
mock_log.get_by_id.assert_called_once_with(dataflow_id)

179
test/unit_test/rag/svr/task_executor_refactor/test_embedding_service.py
View File

@@ -17,13 +17,13 @@
Unit tests for EmbeddingService module.
All tests validate behavior through the public API (embed_chunks) rather than
reaching into private orchestration methods like _encode_single, _encode_batch,
or _run_encode. Those internal boundaries may be reshaped during a refactor
without changing the external behavior; the suite should not break in that case.
reaching into private orchestration methods. The new async implementation uses
thread_pool_exec for model.encode calls; tests mock that boundary.
"""
import numpy as np
from unittest.mock import MagicMock, patch
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from rag.svr.task_executor_refactor.embedding_service import EmbeddingService
@@ -56,17 +56,18 @@ class TestEmbeddingServiceInit:
class TestEmbeddingServiceEmbedChunks:
"""Tests for the public embed_chunks method.
Internal helpers _encode_single, _encode_batch, and _run_encode are
exercised through this public entry point so the suite stays resilient to
method-boundary reshuffles.
The async implementation uses thread_pool_exec for model.encode calls.
Tests mock thread_pool_exec at the module level to control returned vectors.
"""
@patch.object(EmbeddingService, '_run_encode')
def test_embed_chunks_basic(self, mock_run_encode):
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_basic(self, mock_thread_pool):
"""Test basic chunk embedding."""
mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
@@ -74,23 +75,20 @@ class TestEmbeddingServiceEmbedChunks:
docs = [
{"docnm_kwd": "Title1", "content_with_weight": "Content1"},
]
tk_count, vector_size = service.embed_chunks(docs, model)
tk_count, vector_size = await service.embed_chunks(docs, model)
assert tk_count > 0
assert vector_size == 2
assert "q_2_vec" in docs[0]
@patch.object(EmbeddingService, '_run_encode')
def test_embed_chunks_uses_embedding_utils(self, mock_run_encode):
"""Test that embed_chunks uses EmbeddingUtils internally.
The internal path runs _encode_batch -> EmbeddingUtils.truncate_texts
-> _run_encode. We verify via the public embed_chunks that the chain
is wired correctly without asserting on individual private method calls.
"""
mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_uses_embedding_utils(self, mock_thread_pool):
"""Test that embed_chunks uses thread_pool_exec for encoding."""
mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
@@ -98,16 +96,19 @@ class TestEmbeddingServiceEmbedChunks:
docs = [
{"docnm_kwd": "Title1", "content_with_weight": "Content1"},
]
service.embed_chunks(docs, model)
await service.embed_chunks(docs, model)
mock_run_encode.assert_called()
# thread_pool_exec should be called at least once for encoding
mock_thread_pool.assert_called()
@patch.object(EmbeddingService, '_run_encode')
def test_embed_chunks_with_title_content_combination(self, mock_run_encode):
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_with_title_content_combination(self, mock_thread_pool):
"""Test that title and content vectors are combined."""
mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
@@ -115,21 +116,23 @@ class TestEmbeddingServiceEmbedChunks:
docs = [
{"docnm_kwd": "Title1", "content_with_weight": "Content1"},
]
_, vector_size = service.embed_chunks(docs, model, parser_config={"filename_embd_weight": 0.5})
_, vector_size = await service.embed_chunks(docs, model, parser_config={"filename_embd_weight": 0.5})
assert vector_size == 2
@patch.object(EmbeddingService, '_run_encode')
def test_embed_chunks_handles_long_text(self, mock_run_encode):
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_handles_long_text(self, mock_thread_pool):
"""Test that long texts are handled by embedding pipeline.
Even with content exceeding model.max_length, embed_chunks produces
valid vectors, meaning truncation (via EmbeddingUtils) is wired
correctly in the encode path.
"""
mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
@@ -137,9 +140,125 @@ class TestEmbeddingServiceEmbedChunks:
docs = [
{"docnm_kwd": "Title1", "content_with_weight": "a" * 200},
]
tk_count, vector_size = service.embed_chunks(docs, model)
tk_count, vector_size = await service.embed_chunks(docs, model)
# Public contract: embed_chunks returns valid token counts and vectors
assert tk_count > 0
assert vector_size == 2
assert "q_2_vec" in docs[0]
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_empty_docs(self, mock_thread_pool):
"""Test embedding with empty docs list returns zero results."""
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx)
model = MagicMock()
model.max_length = 100
tk_count, vector_size = await service.embed_chunks([], model)
assert tk_count == 0
assert vector_size == 0
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_no_title(self, mock_thread_pool):
"""Test embedding when chunks have no title — content vectors used directly."""
mock_thread_pool.return_value = (np.array([[3.0, 4.0]]), 5)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
docs = [{"content_with_weight": "Content only, no title"}]
tk_count, vector_size = await service.embed_chunks(docs, model)
# With no title, only content is encoded (1 call); vector_size comes from content vec
assert vector_size == 2
assert "q_2_vec" in docs[0]
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_title_weight_zero(self, mock_thread_pool):
"""Test embedding with filename_embd_weight=0.0 — no title contribution."""
mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 5)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
docs = [{"docnm_kwd": "Title1", "content_with_weight": "Content1"}]
_, vector_size = await service.embed_chunks(docs, model,
parser_config={"filename_embd_weight": 0.0})
assert vector_size == 2
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_title_weight_one(self, mock_thread_pool):
"""Test embedding with filename_embd_weight=1.0 — full title contribution."""
mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 5)
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
docs = [{"docnm_kwd": "Title1", "content_with_weight": "Content1"}]
_, vector_size = await service.embed_chunks(docs, model,
parser_config={"filename_embd_weight": 1.0})
assert vector_size == 2
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_encode_failure_propagates(self, mock_thread_pool):
"""Test that model.encode exceptions are propagated to caller."""
mock_thread_pool.side_effect = RuntimeError("embedding service unavailable")
ctx = MagicMock()
ctx.progress_cb = None
ctx.embed_limiter = AsyncMockLimiter()
service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
model = MagicMock()
model.max_length = 100
docs = [{"docnm_kwd": "Title1", "content_with_weight": "Content1"}]
with pytest.raises(RuntimeError, match="embedding service unavailable"):
await service.embed_chunks(docs, model)
@pytest.mark.asyncio
@patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock)
async def test_embed_chunks_multiple_batches(self, mock_thread_pool):
"""Test embedding with more chunks than batch size — multiple encode calls."""
# Each call returns vectors matching input count
def side_effect(func, texts, *args, **kw):
n = len(texts) if isinstance(texts, list) else 1
return np.random.rand(n, 2).astype(np.float32), 10 * n
mock_thread_pool.side_effect = side_effect
ctx = MagicMock()
ctx.progress_cb = MagicMock()
ctx.embed_limiter = AsyncMockLimiter()
# batch_size=2, 5 chunks with titles → 1 title + ceil(5/2)=3 content = 4 calls
service = EmbeddingService(ctx=ctx, embedding_batch_size=2)
model = MagicMock()
model.max_length = 100
docs = [{"docnm_kwd": f"Title{i}", "content_with_weight": f"Content{i}"} for i in range(5)]
_, vector_size = await service.embed_chunks(docs, model)
assert mock_thread_pool.call_count == 4
assert vector_size > 0
# Reuse from conftest
from test.unit_test.rag.svr.task_executor_refactor.conftest import AsyncMockLimiter

52
test/unit_test/rag/svr/task_executor_refactor/test_post_processor.py
View File

@@ -60,6 +60,58 @@ class TestPostProcessorProcessTableParserMetadata:
mock_merge.assert_called_once()
@pytest.mark.asyncio
async def test_processes_table_parser_manual_mode(self):
"""Test that table parser in manual mode aggregates and persists metadata."""
ctx = MagicMock()
ctx.parser_id = "table"
ctx.raw_task = {"parser_config": {}}
ctx.write_interceptor = None
chunks = [{"col_key": "val"}]
with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge, \
patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_manual_doc_metadata") as mock_agg, \
patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip, \
patch("rag.svr.task_executor_refactor.post_processor.update_metadata_to") as mock_update, \
patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta:
mock_merge.return_value = {"table_column_mode": "manual"}
mock_agg.return_value = {"col_key": ["val1", "val2"]}
mock_strip.return_value = set()
mock_meta.get_document_metadata.return_value = {}
mock_update.return_value = {"col_key": ["val1", "val2"]}
service = PostProcessor(ctx=ctx)
await service.process_table_parser_metadata("doc_1", chunks)
mock_agg.assert_called_once_with(chunks, ctx.raw_task)
mock_meta.update_document_metadata.assert_called_once()
@pytest.mark.asyncio
async def test_processes_table_parser_with_write_interceptor(self):
"""Test table parser with write interceptor bypasses DB."""
ctx = MagicMock()
ctx.parser_id = "table"
ctx.raw_task = {}
ctx.write_interceptor = MagicMock()
with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge, \
patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_manual_doc_metadata") as mock_agg, \
patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip, \
patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta:
mock_merge.return_value = {"table_column_mode": "manual"}
mock_agg.return_value = {"key": ["v"]}
mock_strip.return_value = set()
mock_meta.get_document_metadata.return_value = {}
service = PostProcessor(ctx=ctx)
await service.process_table_parser_metadata("doc_1", [])
ctx.write_interceptor.intercept.assert_called_once_with(
"DocMetadataService.update_document_metadata"
)
class TestPostProcessorInsertTocChunk:
"""Tests for insert_toc_chunk method."""

163
test/unit_test/rag/svr/task_executor_refactor/test_raptor_service.py
View File

@@ -33,6 +33,7 @@ import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from rag.svr.task_executor_refactor.raptor_service import RaptorService
from test.unit_test.rag.svr.task_executor_refactor.conftest import make_task_context
# =============================================================================
@@ -216,7 +217,8 @@ class TestRaptorServiceRunRaptorForKb:
# ---- Basic dispatch (file-level scope) ----
def test_run_raptor_for_kb_file_scope_delegates_to_file_level(
@pytest.mark.asyncio
async def test_run_raptor_for_kb_file_scope_delegates_to_file_level(
self, mock_raptor_context, sample_chunks, raptor_config_file_scope
):
"""When scope='file', _run_file_level_raptor is called."""
@@ -233,20 +235,7 @@ class TestRaptorServiceRunRaptorForKb:
patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock) as mock_dataset:
mock_file.return_value = (sample_chunks, 42)
AsyncMock(return_value=(sample_chunks, 42, []))
with patch.object(RaptorService, "run_raptor_for_kb", new=AsyncMock(wraps=svc.run_raptor_for_kb)):
pass # let's just call directly
# Direct call since we need to invoke the async method properly
import asyncio
loop = asyncio.new_event_loop()
try:
chunks, tk_count, cleanup = loop.run_until_complete(
svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
)
finally:
loop.close()
chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
mock_file.assert_called_once()
mock_dataset.assert_not_called()
@@ -255,7 +244,8 @@ class TestRaptorServiceRunRaptorForKb:
# ---- Basic dispatch (dataset-level scope) ----
def test_run_raptor_for_kb_dataset_scope_delegates_to_dataset_level(
@pytest.mark.asyncio
async def test_run_raptor_for_kb_dataset_scope_delegates_to_dataset_level(
self, mock_raptor_context, sample_chunks, raptor_config_dataset_scope
):
"""When scope='dataset', _run_dataset_level_raptor is called."""
@@ -271,15 +261,7 @@ class TestRaptorServiceRunRaptorForKb:
patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock) as mock_dataset:
mock_dataset.return_value = (sample_chunks, 99)
import asyncio
loop = asyncio.new_event_loop()
try:
chunks, tk_count, cleanup = loop.run_until_complete(
svc.run_raptor_for_kb(raptor_config_dataset_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
)
finally:
loop.close()
chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_dataset_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
mock_dataset.assert_called_once()
mock_file.assert_not_called()
@@ -288,7 +270,8 @@ class TestRaptorServiceRunRaptorForKb:
# ---- Empty / no documents ----
def test_run_raptor_for_kb_empty_doc_ids(self, mock_raptor_context, raptor_config_file_scope):
@pytest.mark.asyncio
async def test_run_raptor_for_kb_empty_doc_ids(self, mock_raptor_context, raptor_config_file_scope):
"""Empty doc_ids returns empty results."""
svc = RaptorService(mock_raptor_context)
chat_mdl = MagicMock()
@@ -299,15 +282,7 @@ class TestRaptorServiceRunRaptorForKb:
patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock):
mock_file.return_value = ([], 0)
import asyncio
loop = asyncio.new_event_loop()
try:
chunks, tk_count, cleanup = loop.run_until_complete(
svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, [])
)
finally:
loop.close()
chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, [])
assert chunks == []
assert tk_count == 0
@@ -315,15 +290,11 @@ class TestRaptorServiceRunRaptorForKb:
# ---- Cleanup scheduling through the public API ----
def test_run_raptor_for_kb_returns_cleanup_list(
@pytest.mark.asyncio
async def test_run_raptor_for_kb_returns_cleanup_list(
self, mock_raptor_context, raptor_config_file_scope
):
"""Cleanup list from internal method is propagated to caller.
_run_file_level_raptor receives cleanup_raptor_chunks by reference (as
a positional arg) and may mutate it. This test verifies the public
method propagates whatever ends up in that list.
"""
"""Cleanup list from internal method is propagated to caller."""
svc = RaptorService(mock_raptor_context)
doc_ids = ["doc_1"]
chat_mdl = MagicMock()
@@ -336,28 +307,19 @@ class TestRaptorServiceRunRaptorForKb:
}), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file:
async def mock_run_file(*args, **kwargs):
# _run_file_level_raptor takes 12 positional args;
# cleanup_raptor_chunks is args[11] (0-indexed, last positional).
cleanup_list = args[11]
cleanup_list.append(("doc_1", "tree_builder_a"))
return [{"id": "c1"}], 10
mock_file.side_effect = mock_run_file
import asyncio
loop = asyncio.new_event_loop()
try:
chunks, tk_count, cleanup = loop.run_until_complete(
svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
)
finally:
loop.close()
chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
assert cleanup == expected_cleanup
# ---- Dispatch with missing raptor config key ----
def test_run_raptor_for_kb_defaults_to_file_scope_when_no_raptor_key(
@pytest.mark.asyncio
async def test_run_raptor_for_kb_defaults_to_file_scope_when_no_raptor_key(
self, mock_raptor_context
):
"""When kb_parser_config has no 'raptor' key, defaults to file scope."""
@@ -373,22 +335,15 @@ class TestRaptorServiceRunRaptorForKb:
patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock) as mock_dataset:
mock_file.return_value = ([], 0)
import asyncio
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(
svc.run_raptor_for_kb(config, chat_mdl, embd_mdl, 128, doc_ids)
)
finally:
loop.close()
await svc.run_raptor_for_kb(config, chat_mdl, embd_mdl, 128, doc_ids)
mock_file.assert_called_once()
mock_dataset.assert_not_called()
# ---- Vector dimension name construction ----
def test_run_raptor_for_kb_passes_vector_size_to_file_level(
@pytest.mark.asyncio
async def test_run_raptor_for_kb_passes_vector_size_to_file_level(
self, mock_raptor_context, sample_chunks, raptor_config_file_scope
):
"""Vector size is used to construct vctr_nm and passed to internal method."""
@@ -403,15 +358,7 @@ class TestRaptorServiceRunRaptorForKb:
}), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file:
mock_file.return_value = (sample_chunks, 10)
import asyncio
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(
svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
)
finally:
loop.close()
await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
# Verify _run_file_level_raptor received vctr_nm with the correct vector size
# Positional args: 0=raptor_config, 1=tree_builder, 2=clustering_method,
@@ -421,7 +368,8 @@ class TestRaptorServiceRunRaptorForKb:
# ---- Document info collection through public API ----
def test_run_raptor_for_kb_collects_doc_info(
@pytest.mark.asyncio
async def test_run_raptor_for_kb_collects_doc_info(
self, mock_raptor_context, raptor_config_file_scope
):
"""Document info is collected before dispatching to internal methods."""
@@ -436,17 +384,66 @@ class TestRaptorServiceRunRaptorForKb:
patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file:
mock_file.return_value = ([], 0)
import asyncio
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(
svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
)
finally:
loop.close()
await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
mock_collect.assert_called_once_with(doc_ids)
# Verify doc_info_by_id was passed as positional arg[7] to _run_file_level_raptor
positional_args = mock_file.call_args[0]
assert positional_args[7] == expected_info
class TestRaptorServiceFileLevelRaptorCheckpoint:
"""Tests for _run_file_level_raptor checkpoint behavior.
Verifies the fix that moves progress_cb and continue to the outer if
block so progress is reported even when existing_methods == {tree_builder}.
"""
@pytest.mark.asyncio
async def test_file_level_raptor_existing_methods_exact_match_updates_progress(self):
"""When existing_methods == {tree_builder}, progress_cb is still called."""
ctx = make_task_context()
svc = RaptorService(ctx)
doc_ids = ["doc_1"]
doc_info_by_id = {
"doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}}
}
raptor_config = {
"scope": "file",
"max_cluster": 64, "prompt": "test prompt",
"max_token": 256, "threshold": 0.1, "random_seed": 0,
"clustering_method": "gmm", "tree_builder": "raptor",
"ext": {},
}
with patch.object(svc, "_get_raptor_chunk_methods", new_callable=AsyncMock) as mock_methods, \
patch.object(svc, "_should_skip_raptor", return_value=False):
mock_methods.return_value = {"raptor"}
result = await svc._run_file_level_raptor(
raptor_config=raptor_config, tree_builder="raptor",
clustering_method="gmm", chat_mdl=MagicMock(),
embd_mdl=MagicMock(), vctr_nm="q_128_vec",
doc_ids=doc_ids, doc_info_by_id=doc_info_by_id,
max_errors=3, res=[], tk_count=0,
cleanup_raptor_chunks=[],
)
msg_calls = [
call.kwargs.get("msg", "")
for call in ctx.progress_cb.call_args_list
if call.kwargs.get("msg") is not None
]
assert any("already has" in m for m in msg_calls), \
f"Expected 'already has' progress message, got: {msg_calls}"
prog_calls = [
call.kwargs.get("prog")
for call in ctx.progress_cb.call_args_list
if call.kwargs.get("prog") is not None
]
assert len(prog_calls) > 0, \
"Expected progress_cb to be called with prog update"
assert result[0] == []

4
test/unit_test/rag/svr/task_executor_refactor/test_task_context.py
View File

@@ -216,10 +216,10 @@ class TestTaskContextLanguageAndModelProperties:
"""Tests for language and model properties."""
def test_language_default(self):
"""Test language property defaults to 'en'."""
"""Test language property defaults to 'Chinese'."""
task = {"id": "task_1", "tenant_id": "tenant_1"}
ctx = _make_ctx(task=task)
assert ctx.language == "en"
assert ctx.language == "Chinese"
def test_language(self):
"""Test language property."""

282
test/unit_test/rag/svr/task_executor_refactor/test_task_handler.py
View File

@@ -16,21 +16,27 @@
"""
Unit tests for TaskHandler module.
All orchestration tests validate behavior through the public handle()/handle_task()
entry points. Internal helpers (_run_standard_chunking, _run_dataflow, _run_raptor,
_run_graphrag, _bind_embedding_model, _get_storage_binary, etc.) are exercised
implicitly; no test reaches directly into those private orchestration methods.
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, _get_vector_size) are tested directly since they
are side-effect-free data transformations.
Stable pure helpers (_build_toc) are tested directly.
"""
import pytest
import numpy as np
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."""
@@ -68,63 +74,86 @@ class TestTaskHandlerHandleTask:
ctx.recording_context = MagicMock()
handler = TaskHandler(ctx=ctx)
handler.handle = AsyncMock(side_effect=Exception("test error"))
# Should raise the exception
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.
Internal orchestration methods (_run_standard_chunking, _run_dataflow,
_run_raptor, _run_graphrag, _bind_embedding_model) are exercised through
handle() so the suite stays resilient when those private methods change.
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 dispatches memory tasks correctly."""
ctx = MagicMock()
ctx.task_type = "memory"
ctx.id = "task_1"
ctx.raw_task = {"memory_id": "mem_1"}
ctx.write_interceptor = None
ctx.has_canceled_func = MagicMock(return_value=False)
"""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:
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._bind_embedding_model = AsyncMock()
handler._get_vector_size = MagicMock(return_value=1024)
handler._init_kb = MagicMock()
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."""
ctx = MagicMock()
ctx.task_type = "dataflow"
ctx.id = "task_1"
ctx.doc_id = "doc_1"
ctx.has_canceled_func = MagicMock(return_value=False)
"""Test handle dispatches dataflow tasks (after embedding binding + init_kb)."""
ctx = make_task_context(task_type="dataflow", doc_id="doc_1")
handler = TaskHandler(ctx=ctx)
handler._run_dataflow = AsyncMock()
await handler.handle()
handler._run_dataflow.assert_called_once()
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 = MagicMock()
ctx.task_type = "standard"
ctx.id = "task_1"
ctx.has_canceled_func = MagicMock(return_value=True)
ctx.progress_cb = MagicMock()
ctx = make_task_context(has_canceled_func=MagicMock(return_value=True))
handler = TaskHandler(ctx=ctx)
await handler.handle()
@@ -132,102 +161,87 @@ class TestTaskHandlerHandle:
@pytest.mark.asyncio
async def test_handle_standard_chunking(self):
"""Test handle dispatches standard chunking end-to-end."""
ctx = MagicMock()
ctx.task_type = "standard"
ctx.id = "task_1"
ctx.tenant_id = "tenant_1"
ctx.kb_id = "kb_1"
ctx.doc_id = "doc_1"
ctx.embd_id = "embd_1"
ctx.language = "en"
ctx.parser_config = {}
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.progress_cb = MagicMock()
ctx.recording_context = MagicMock()
ctx.name = "test.pdf"
ctx.from_page = 0
ctx.to_page = -1
"""Test handle routes to standard chunking.
handler = TaskHandler(ctx=ctx)
handler._bind_embedding_model = AsyncMock(return_value=MagicMock())
handler._get_vector_size = MagicMock(return_value=128)
handler._init_kb = MagicMock()
handler._run_standard_chunking = AsyncMock()
``_bind_embedding_model`` and ``_init_kb`` run their real code;
only the external boundary (LLM API, settings) is mocked.
"""
ctx = make_task_context()
await handler.handle()
handler._run_standard_chunking.assert_called_once()
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 dispatches raptor tasks."""
ctx = MagicMock()
ctx.task_type = "raptor"
ctx.id = "task_1"
ctx.tenant_id = "tenant_1"
ctx.kb_id = "kb_1"
ctx.embd_id = "embd_1"
ctx.language = "en"
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.progress_cb = MagicMock()
ctx.recording_context = MagicMock()
"""Test handle routes to RAPTOR with real embedding binding."""
ctx = make_task_context(task_type="raptor")
handler = TaskHandler(ctx=ctx)
handler._bind_embedding_model = AsyncMock(return_value=MagicMock())
handler._get_vector_size = MagicMock(return_value=128)
handler._init_kb = MagicMock()
handler._run_raptor = AsyncMock()
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"):
await handler.handle()
handler._run_raptor.assert_called_once()
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 dispatches graphrag tasks."""
ctx = MagicMock()
ctx.task_type = "graphrag"
ctx.id = "task_1"
ctx.tenant_id = "tenant_1"
ctx.kb_id = "kb_1"
ctx.embd_id = "embd_1"
ctx.language = "en"
ctx.has_canceled_func = MagicMock(return_value=False)
ctx.progress_cb = MagicMock()
ctx.recording_context = MagicMock()
"""Test handle routes to GraphRAG with real embedding binding."""
ctx = make_task_context(task_type="graphrag")
handler = TaskHandler(ctx=ctx)
handler._bind_embedding_model = AsyncMock(return_value=MagicMock())
handler._get_vector_size = MagicMock(return_value=128)
handler._init_kb = MagicMock()
handler._run_graphrag = AsyncMock()
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"):
await handler.handle()
handler._run_graphrag.assert_called_once()
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."""
ctx = MagicMock()
ctx.task_type = "standard"
ctx.id = "task_1"
ctx.has_canceled_func = MagicMock(return_value=False)
"""Test handle returns early when embedding model binding fails.
handler = TaskHandler(ctx=ctx)
handler._bind_embedding_model = AsyncMock(return_value=None)
``LLMBundle`` is patched to raise, so ``_bind_embedding_model``
itself raises — no need to mock the private method.
"""
ctx = make_task_context()
await handler.handle()
# Should not call _run_standard_chunking when model is None
assert not hasattr(handler, '_run_standard_chunking_called')
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")
class TestTaskHandlerGetVectorSize:
"""Tests for _get_vector_size — stable pure helper."""
handler = TaskHandler(ctx=ctx)
with pytest.raises(RuntimeError, match="embedding service unavailable"):
await handler.handle()
def test_get_vector_size(self):
mock_model = MagicMock()
mock_model.encode.return_value = (np.array([[1.0, 2.0, 3.0]]), 10)
result = TaskHandler._get_vector_size(mock_model)
assert result == 3
@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:
@@ -243,18 +257,19 @@ class TestTaskHandlerBuildToc:
docs = [{"id": "chunk_1", "content_with_weight": "text", "page_num_int": [1], "top_int": [0]}]
def mock_asyncio_run(coro):
# Close the coroutine to prevent "never awaited" warnings
coro.close()
return []
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg:
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()
with patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle:
mock_msg = MagicMock()
mock_bundle.return_value.__enter__.return_value = mock_msg
with patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run):
result = TaskHandler._build_toc(ctx, docs, MagicMock())
assert result is None
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."""
@@ -267,21 +282,22 @@ class TestTaskHandlerBuildToc:
toc_result = [{"chunk_id": "0", "title": "Section 1"}]
def mock_asyncio_run(coro):
# Close the coroutine to prevent "never awaited" warnings
coro.close()
return toc_result
with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg:
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()
with patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle:
mock_msg = MagicMock()
mock_bundle.return_value.__enter__.return_value = mock_msg
with patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run):
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
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:
@@ -297,4 +313,4 @@ class TestTaskHandlerInit:
def test_init_default_hook_none(self):
ctx = MagicMock()
handler = TaskHandler(ctx=ctx)
assert handler._billing_hook is None
assert handler._billing_hook is None

371
test/unit_test/rag/svr/task_executor_refactor/test_task_handler_integration.py
View File

@@ -19,7 +19,6 @@ Integration tests for TaskHandler orchestration.
import asyncio
import gc
import uuid
from typing import Any, Dict
from unittest.mock import MagicMock, AsyncMock, patch
@@ -37,6 +36,11 @@ from test.unit_test.rag.svr.task_executor_refactor.conftest import (
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,
)
@@ -82,48 +86,13 @@ def create_task_context(
return ctx
# Common patcher for _get_storage_binary since it imports settings internally
def patch_get_storage_binary():
return patch.object(TaskHandler, '_get_storage_binary', new_callable=AsyncMock, return_value=b"fake pdf binary")
def patch_task_handler_settings(mock_settings):
"""Patch the settings module-level import in task_handler."""
return patch("rag.svr.task_executor_refactor.task_handler.settings", mock_settings)
class TestStandardChunkingPipelineIntegration:
"""P0: Integration tests for the complete standard chunking pipeline."""
def _create_standard_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": "doc_test",
"name": "test_document.pdf",
"location": "/path/to/test_document.pdf",
"size": 1024,
"parser_id": "naive",
"parser_config": {
"auto_keywords": 0,
"auto_questions": 0,
"enable_metadata": False,
},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "standard",
"pagerank": 0,
}
@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 = self._create_standard_task_dict()
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()
@@ -152,11 +121,8 @@ class TestStandardChunkingPipelineIntegration:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -168,7 +134,7 @@ class TestStandardChunkingPipelineIntegration:
@pytest.mark.asyncio
async def test_full_chunking_pipeline_records_insertion_result(self):
"""Verify that insertion_result is recorded as 'success'."""
task_dict = self._create_standard_task_dict()
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()
@@ -197,11 +163,8 @@ class TestStandardChunkingPipelineIntegration:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -213,7 +176,7 @@ class TestStandardChunkingPipelineIntegration:
@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 = self._create_standard_task_dict()
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()
@@ -245,11 +208,8 @@ class TestStandardChunkingPipelineIntegration:
mock_chunk_service_cls.return_value = mock_chunk_service
mock_run_toc.return_value = [] # TOC returns empty when not enabled
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -262,7 +222,7 @@ class TestStandardChunkingPipelineIntegration:
@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 = self._create_standard_task_dict()
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()
@@ -291,11 +251,8 @@ class TestStandardChunkingPipelineIntegration:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -311,7 +268,7 @@ class TestStandardChunkingPipelineIntegration:
@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 = self._create_standard_task_dict()
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()
@@ -340,11 +297,8 @@ class TestStandardChunkingPipelineIntegration:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -357,31 +311,10 @@ class TestStandardChunkingPipelineIntegration:
class TestTaskCancellationCleanupIntegration:
"""P0: Integration tests for task cancellation cleanup flow."""
def _create_standard_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": "doc_test",
"name": "test_document.pdf",
"location": "/path/to/test_document.pdf",
"size": 1024,
"parser_id": "naive",
"parser_config": {},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "standard",
"pagerank": 0,
}
@pytest.mark.asyncio
async def test_canceled_task_calls_docstore_delete(self):
"""Verify that docStoreConn.delete is called when task is canceled."""
task_dict = self._create_standard_task_dict()
task_dict = make_task_dict()
ctx = create_task_context(task_dict, is_canceled=True)
mock_settings = create_mock_settings()
@@ -411,7 +344,7 @@ class TestTaskCancellationCleanupIntegration:
@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 = self._create_standard_task_dict()
task_dict = make_task_dict()
ctx = create_task_context(task_dict, is_canceled=True)
mock_settings = create_mock_settings()
@@ -450,7 +383,7 @@ class TestTaskCancellationCleanupIntegration:
@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 = self._create_standard_task_dict()
task_dict = make_task_dict()
ctx = create_task_context(task_dict, is_canceled=True)
mock_settings = create_mock_settings()
@@ -481,30 +414,10 @@ class TestTaskCancellationCleanupIntegration:
class TestRaptorPipelineIntegration:
"""P1: Integration tests for the RAPTOR pipeline."""
def _create_raptor_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": GRAPH_RAPTOR_FAKE_DOC_ID,
"doc_ids": ["doc1", "doc2"],
"name": "raptor_task",
"parser_id": "naive",
"parser_config": {"raptor": {"use_raptor": False}},
"kb_parser_config": {"raptor": {"use_raptor": False}},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "raptor",
"pagerank": 0,
}
@pytest.mark.asyncio
async def test_raptor_pipeline_records_task_status(self):
"""Verify that RAPTOR pipeline records task_status."""
task_dict = self._create_raptor_task_dict()
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()
@@ -534,10 +447,7 @@ class TestRaptorPipelineIntegration:
mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
mock_doc_service.increment_chunk_num = MagicMock()
async def mock_thread_impl(func, *args, **kwargs):
return None
mock_thread_exec.side_effect = mock_thread_impl
mock_thread_exec.side_effect = mock_thread_return_none
handler = TaskHandler(ctx=ctx)
await handler.handle()
@@ -549,7 +459,7 @@ class TestRaptorPipelineIntegration:
@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 = self._create_raptor_task_dict()
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()
@@ -579,10 +489,7 @@ class TestRaptorPipelineIntegration:
mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
mock_doc_service.increment_chunk_num = MagicMock()
async def mock_thread_impl(func, *args, **kwargs):
return None
mock_thread_exec.side_effect = mock_thread_impl
mock_thread_exec.side_effect = mock_thread_return_none
handler = TaskHandler(ctx=ctx)
await handler.handle()
@@ -598,31 +505,10 @@ class TestRaptorPipelineIntegration:
class TestEmbeddingModelBindingFailureIntegration:
"""P1: Integration tests for embedding model binding failure."""
def _create_standard_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": "doc_test",
"name": "test_document.pdf",
"location": "/path/to/test_document.pdf",
"size": 1024,
"parser_id": "naive",
"parser_config": {},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "standard",
"pagerank": 0,
}
@pytest.mark.asyncio
async def test_embedding_binding_failure_raises_exception(self):
"""Verify that embedding model binding failure raises an exception."""
task_dict = self._create_standard_task_dict()
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, \
@@ -639,7 +525,7 @@ class TestEmbeddingModelBindingFailureIntegration:
@pytest.mark.asyncio
async def test_embedding_binding_failure_calls_progress_callback(self):
"""Verify that embedding model binding failure calls progress_callback."""
task_dict = self._create_standard_task_dict()
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, \
@@ -659,49 +545,26 @@ class TestEmbeddingModelBindingFailureIntegration:
class TestDataflowPipelineIntegration:
"""P2: Integration tests for the dataflow pipeline."""
def _create_dataflow_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": CANVAS_DEBUG_DOC_ID,
"name": "dataflow_debug",
"parser_id": "naive",
"parser_config": {},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "dataflow",
"pagerank": 0,
}
@pytest.mark.asyncio
async def test_dataflow_pipeline_calls_dataflow_service(self):
"""Verify that dataflow pipeline calls DataflowService.run_dataflow()."""
task_dict = self._create_dataflow_task_dict()
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.DataflowService") as mock_dataflow_service:
mock_instance = MagicMock()
mock_instance.run_dataflow = AsyncMock(return_value=None)
mock_dataflow_service.return_value = mock_instance
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:
handler = TaskHandler(ctx=ctx)
await handler.handle()
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_dataflow_service.assert_called_once()
mock_instance.run_dataflow.assert_called_once()
@pytest.mark.asyncio
async def test_dataflow_debug_mode_calls_dataflow_service(self):
"""Verify that dataflow debug mode also calls DataflowService."""
task_dict = self._create_dataflow_task_dict()
ctx = create_task_context(task_dict)
with patch("rag.svr.task_executor_refactor.task_handler.DataflowService") as mock_dataflow_service:
mock_instance = MagicMock()
mock_instance.run_dataflow = AsyncMock(return_value=None)
mock_dataflow_service.return_value = mock_instance
@@ -716,37 +579,11 @@ class TestDataflowPipelineIntegration:
class TestTocAsyncFlowIntegration:
"""P2: Integration tests for TOC async flow."""
def _create_toc_enabled_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": "doc_test",
"name": "test_document.pdf",
"location": "/path/to/test_document.pdf",
"size": 1024,
"parser_id": "naive",
"parser_config": {
"auto_keywords": 0,
"auto_questions": 0,
"enable_metadata": False,
"toc_extraction": True,
},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "standard",
"pagerank": 0,
}
@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 = self._create_toc_enabled_task_dict()
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()
@@ -779,11 +616,8 @@ class TestTocAsyncFlowIntegration:
mock_run_toc.return_value = [{"title": "Test TOC", "level": 1}]
mock_post_doc_service.increment_chunk_num = MagicMock()
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -811,7 +645,7 @@ class TestTocAsyncFlowIntegration:
by pytest-asyncio.
"""
task_dict = self._create_toc_enabled_task_dict()
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)
@@ -842,11 +676,8 @@ class TestTocAsyncFlowIntegration:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -873,35 +704,10 @@ class TestTocAsyncFlowIntegration:
class TestRecordingContextDataFlowAssertions:
"""P2: Integration tests for RecordingContext data flow assertions."""
def _create_standard_task_dict(self) -> Dict[str, Any]:
return {
"id": f"task_{uuid.uuid4().hex[:8]}",
"tenant_id": "tenant_test",
"kb_id": "kb_test",
"doc_id": "doc_test",
"name": "test_document.pdf",
"location": "/path/to/test_document.pdf",
"size": 1024,
"parser_id": "naive",
"parser_config": {
"auto_keywords": 0,
"auto_questions": 0,
"enable_metadata": False,
},
"kb_parser_config": {},
"language": "en",
"llm_id": "llm_test",
"embd_id": "embd_test",
"from_page": 0,
"to_page": -1,
"task_type": "standard",
"pagerank": 0,
}
@pytest.mark.asyncio
async def test_recording_context_captures_file_size_check(self):
"""Verify that RecordingContext captures file_size_exceeded result."""
task_dict = self._create_standard_task_dict()
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()
@@ -930,11 +736,8 @@ class TestRecordingContextDataFlowAssertions:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -947,7 +750,7 @@ class TestRecordingContextDataFlowAssertions:
@pytest.mark.asyncio
async def test_recording_context_captures_parser_id(self):
"""Verify that RecordingContext captures parser_id from task context."""
task_dict = self._create_standard_task_dict()
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()
@@ -976,11 +779,8 @@ class TestRecordingContextDataFlowAssertions:
mock_doc_service.update_document_metadata = MagicMock()
mock_chunk_service_cls.return_value = mock_chunk_service
async def mock_thread_impl(func, *args, **kwargs):
return b"fake pdf binary"
mock_thread_exec.side_effect = mock_thread_impl
mock_chunk_thread_exec.side_effect = mock_thread_impl
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()
@@ -991,3 +791,66 @@ class TestRecordingContextDataFlowAssertions:
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"