Refactor: Task Executor (#15154)

### What problem does this PR solve?

1. Break huge function into smaller pieces
2. Add unit test for the smaller pieces function
3. Layer-ed design
a. infra layer - task_context.py, recording_context.py,
write_operation_interceptor.py, ...
    b. service layer - *_service.py
    c. business layer - task_handler.py
4. Default behavior: use "refactor-ed version" - can switch to original
version by change env variable

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement

---------

Co-authored-by: Liu An <asiro@qq.com>
Co-authored-by: Zhichang Yu <yuzhichang@gmail.com>

test/unit_test/common/test_settings_queue.py

@@ -0,0 +1,208 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+"""Test cases for get_svr_queue_name and get_svr_queue_names functions in common.settings."""
+
+from common.settings import get_svr_queue_name, get_svr_queue_names
+
+
+class TestGetSvrQueueName:
+    """Test cases for get_svr_queue_name function."""
+
+    def test_default_suffix(self):
+        """Test that default suffix is 'common'."""
+
+        result = get_svr_queue_name(0)
+        assert result == "te.0.common"
+
+    def test_priority_zero(self):
+        """Test queue name with priority 0 (low)."""
+
+        result = get_svr_queue_name(0)
+        assert result == "te.0.common"
+
+    def test_priority_one(self):
+        """Test queue name with priority 1 (high)."""
+
+        result = get_svr_queue_name(1)
+        assert result == "te.1.common"
+
+    def test_explicit_suffix_common(self):
+        """Test with explicit 'common' suffix."""
+
+        result = get_svr_queue_name(0, "common")
+        assert result == "te.0.common"
+
+    def test_suffix_parameter_ignored(self):
+        """Test that suffix parameter is currently ignored (hardcoded to 'common').
+        
+        Note: The function signature accepts a suffix parameter but currently
+        hardcodes 'common' in the return value. This test documents this behavior.
+        """
+
+        # Even with different suffix values, result should be the same
+        result_default = get_svr_queue_name(0, "common")
+        result_resume = get_svr_queue_name(0, "resume")
+        result_graphrag = get_svr_queue_name(0, "graphrag")
+
+        # All should return the same value since suffix is hardcoded
+        assert result_default == result_resume == result_graphrag == "te.0.common"
+
+    def test_format_structure(self):
+        """Test that queue name follows expected format: {SVR_QUEUE_NAME}.{priority}.common."""
+
+        for priority in [0, 1]:
+            result = get_svr_queue_name(priority)
+            parts = result.split(".")
+            assert len(parts) == 3
+            assert parts[0] == "te"  # SVR_QUEUE_NAME
+            assert parts[1] == str(priority)
+            assert parts[2] == "common"
+
+    def test_different_priorities_produce_different_results(self):
+        """Test that different priorities produce different queue names."""
+
+        result_0 = get_svr_queue_name(0)
+        result_1 = get_svr_queue_name(1)
+
+        assert result_0 != result_1
+        assert result_0 == "te.0.common"
+        assert result_1 == "te.1.common"
+
+    def test_with_various_priority_values(self):
+        """Test with various priority values beyond 0 and 1."""
+
+        # Test with other priority values to ensure format is correct
+        for priority in [2, 5, 10, 100]:
+            result = get_svr_queue_name(priority)
+            expected = f"te.{priority}.common"
+            assert result == expected
+
+    def test_returns_string_type(self):
+        """Test that function returns a string."""
+
+        result = get_svr_queue_name(0)
+        assert isinstance(result, str)
+
+    def test_no_whitespace_issues(self):
+        """Test that queue name has no unexpected whitespace."""
+
+        for priority in [0, 1]:
+            result = get_svr_queue_name(priority)
+            assert " " not in result
+            assert "\t" not in result
+            assert "\n" not in result
+
+
+class TestGetSvrQueueNames:
+    """Test cases for get_svr_queue_names function."""
+
+    def test_returns_list(self):
+        """Test that function returns a list."""
+
+        result = get_svr_queue_names("common")
+        assert isinstance(result, list)
+
+    def test_returns_two_queues(self):
+        """Test that function returns exactly two queue names."""
+
+        result = get_svr_queue_names("common")
+        assert len(result) == 2
+
+    def test_sorted_high_to_low(self):
+        """Test that queue names are sorted from high priority to low priority."""
+
+        result = get_svr_queue_names("common")
+        assert result[0] == "te.1.common"  # High priority first
+        assert result[1] == "te.0.common"  # Low priority second
+
+    def test_expected_values(self):
+        """Test that returned values match expected queue names."""
+
+        result = get_svr_queue_names("common")
+        expected = ["te.1.common", "te.0.common"]
+        assert result == expected
+
+    def test_suffix_parameter_passed_through(self):
+        """Test that suffix parameter is passed to get_svr_queue_name.
+        
+        Note: Since get_svr_queue_name currently hardcodes 'common' as the suffix,
+        different suffix values will still produce the same result.
+        """
+
+        # All suffixes should produce same result due to hardcoded suffix in get_svr_queue_name
+        result_common = get_svr_queue_names("common")
+        result_resume = get_svr_queue_names("resume")
+        result_graphrag = get_svr_queue_names("graphrag")
+
+        expected = ["te.1.common", "te.0.common"]
+        assert result_common == expected
+        assert result_resume == expected  # suffix is currently ignored
+        assert result_graphrag == expected  # suffix is currently ignored
+
+    def test_all_elements_are_strings(self):
+        """Test that all elements in the returned list are strings."""
+
+        result = get_svr_queue_names("common")
+        for item in result:
+            assert isinstance(item, str)
+
+    def test_consistent_results(self):
+        """Test that multiple calls return consistent results."""
+
+        result1 = get_svr_queue_names("common")
+        result2 = get_svr_queue_names("common")
+        result3 = get_svr_queue_names("common")
+
+        assert result1 == result2 == result3
+
+    def test_with_empty_suffix(self):
+        """Test with empty string suffix."""
+
+        result = get_svr_queue_names("")
+        # Should still work since suffix is ignored
+        assert result == ["te.1.common", "te.0.common"]
+
+
+class TestGetSvrQueueNameWithMockedConstant:
+    """Test cases with mocked SVR_QUEUE_NAME constant."""
+
+    def test_with_custom_queue_name(self):
+        """Test with a custom SVR_QUEUE_NAME constant."""
+        # Need to patch where the constant is imported in settings module
+        import common.settings as settings_mod
+
+        original_value = settings_mod.SVR_QUEUE_NAME
+        try:
+            settings_mod.SVR_QUEUE_NAME = "custom_queue"
+            result = settings_mod.get_svr_queue_name(0)
+            assert result == "custom_queue.0.common"
+
+            result = settings_mod.get_svr_queue_name(1)
+            assert result == "custom_queue.1.common"
+        finally:
+            settings_mod.SVR_QUEUE_NAME = original_value
+
+    def test_with_custom_queue_names(self):
+        """Test get_svr_queue_names with a custom SVR_QUEUE_NAME constant."""
+        import common.settings as settings_mod
+
+        original_value = settings_mod.SVR_QUEUE_NAME
+        try:
+            settings_mod.SVR_QUEUE_NAME = "custom_queue"
+            result = settings_mod.get_svr_queue_names("common")
+            assert result == ["custom_queue.1.common", "custom_queue.0.common"]
+        finally:
+            settings_mod.SVR_QUEUE_NAME = original_value

test/unit_test/rag/svr/task_executor_refactor/conftest.py

@@ -0,0 +1,494 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Shared pytest fixtures for task_executor_refactor integration tests.
+
+This module provides reusable fixtures for integration tests that verify
+the complete orchestration flow of TaskHandler and its collaborating services.
+
+Design principles:
+- Mock external system boundaries (LLM, ES, MinIO, MySQL)
+- Use real TaskContext, TaskHandler, and service instances
+- Verify RecordingContext for data flow assertions
+"""
+# =============================================================================
+# TensorFlow/UMAP Import Workaround
+# =============================================================================
+# Mock umap.parametric_umap before any other imports to prevent TensorFlow
+# dependency errors during test collection. This allows tests to run without
+# requiring TensorFlow to be installed.
+import sys
+from unittest.mock import MagicMock
+
+# Create a mock module for parametric_umap to satisfy umap's import check
+_mock_parametric_umap = MagicMock()
+sys.modules.setdefault("umap.parametric_umap", _mock_parametric_umap)
+sys.modules.setdefault("umap", MagicMock())
+
+import asyncio
+import uuid
+from typing import Any, Dict, List
+from unittest.mock import MagicMock, AsyncMock, patch
+
+import numpy as np
+import pytest
+
+from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+from rag.svr.task_executor_refactor.recording_context import (
+    RecordingContext,
+    set_recording_context,
+)
+
+
+# =============================================================================
+# Async Limiter Fixtures
+# =============================================================================
+
+class AsyncMockLimiter:
+    """Mock asyncio semaphore that does not actually limit."""
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        pass
+
+
+@pytest.fixture
+def mock_limiter():
+    """Provide a no-op async limiter."""
+    return asyncio.Semaphore(5)
+
+
+# =============================================================================
+# Task Dictionary Fixtures
+# =============================================================================
+
+@pytest.fixture
+def standard_task_dict() -> Dict[str, Any]:
+    """Provide a minimal but complete task dict for standard chunking."""
+    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.fixture
+def dataflow_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for dataflow tasks."""
+    task = standard_task_dict()
+    task["task_type"] = "dataflow"
+    task["dataflow_id"] = "dataflow_test"
+    return task
+
+
+@pytest.fixture
+def raptor_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for RAPTOR tasks."""
+    task = standard_task_dict()
+    task["task_type"] = "raptor"
+    task["doc_ids"] = ["doc_1", "doc_2"]
+    return task
+
+
+@pytest.fixture
+def graphrag_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for GraphRAG tasks."""
+    task = standard_task_dict()
+    task["task_type"] = "graphrag"
+    task["doc_ids"] = ["doc_1"]
+    return task
+
+
+@pytest.fixture
+def memory_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for memory tasks."""
+    return {
+        "id": f"task_{uuid.uuid4().hex[:8]}",
+        "task_type": "memory",
+        "memory_id": "mem_test",
+        "source_id": "src_test",
+        "message_dict": {"role": "user", "content": "test"},
+    }
+
+
+# =============================================================================
+# TaskContext Fixtures
+# =============================================================================
+
+@pytest.fixture
+def task_context(standard_task_dict, mock_limiter, recording_context):
+    """Provide a real TaskContext instance with mocked limiters."""
+    ctx = TaskContext(
+        task=standard_task_dict,
+        limiters=TaskLimiters(
+            chat=mock_limiter,
+            minio=mock_limiter,
+            chunk=mock_limiter,
+            embed=mock_limiter,
+            kg=mock_limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=MagicMock(),
+            has_canceled=MagicMock(return_value=False),
+        ),
+        recording_context=recording_context,
+    )
+    return ctx
+
+
+@pytest.fixture
+def canceled_task_context(standard_task_dict, mock_limiter, recording_context):
+    """Provide a TaskContext where the task is already canceled."""
+    ctx = TaskContext(
+        task=standard_task_dict,
+        limiters=TaskLimiters(
+            chat=mock_limiter,
+            minio=mock_limiter,
+            chunk=mock_limiter,
+            embed=mock_limiter,
+            kg=mock_limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=MagicMock(),
+            has_canceled=MagicMock(return_value=True),
+        ),
+        recording_context=recording_context,
+    )
+    return ctx
+
+
+# =============================================================================
+# RecordingContext Fixtures
+# =============================================================================
+
+@pytest.fixture(autouse=True)
+def recording_context():
+    """Provide a fresh RecordingContext for each test.
+
+    This fixture is autouse=True to ensure every test has a clean
+    recording context for assertions.
+    """
+    ctx = RecordingContext()
+    set_recording_context(ctx)
+    yield ctx
+    # Cleanup: reset the global context after test
+    set_recording_context(RecordingContext())
+
+
+@pytest.fixture(autouse=True)
+def cleanup_resources(request):
+    """Global resource cleanup fixture.
+
+    Runs after each test to clean up:
+    - Unclosed event loops
+    - Unclosed sockets (via garbage collection)
+    - Unawaited coroutines
+    - MagicMock objects that may hold unclosed resources
+
+    This prevents ResourceWarning and RuntimeWarning from failing
+    tests when filterwarnings is set to "error".
+
+    Optimization: Uses minimal gc cycles and generation-2 collection
+    for faster teardown.
+    """
+    yield
+    import warnings
+
+    # Suppress warnings during cleanup to avoid recursive warning issues
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+
+        # Close any unclosed event loops
+        try:
+            policy = asyncio.get_event_loop_policy()
+            loop = policy.get_event_loop()
+            if not loop.is_closed():
+                loop.close()
+        except RuntimeError:
+            # No event loop exists, which is fine
+            pass
+
+
+# =============================================================================
+# External System Mocks (Boundary Mocks)
+# =============================================================================
+
+class MockEmbeddingModel:
+    """Mock embedding model that returns deterministic vectors."""
+
+    def __init__(self, vector_size: int = 128):
+        self.vector_size = vector_size
+        self.max_length = 512
+        self.llm_name = "mock_embedding"
+
+    def encode(self, texts: List[str]):
+        """Return random vectors for the given texts."""
+        vectors = np.random.rand(len(texts), self.vector_size).astype(np.float32)
+        token_count = sum(len(t.split()) for t in texts)
+        return vectors, token_count
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+class MockChatModel:
+    """Mock chat model that returns canned responses."""
+
+    def __init__(self):
+        self.llm_name = "mock_chat"
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+@pytest.fixture
+def mock_embedding_model():
+    """Provide a mock embedding model."""
+    return MockEmbeddingModel(vector_size=128)
+
+
+@pytest.fixture
+def mock_chat_model():
+    """Provide a mock chat model."""
+    return MockChatModel()
+
+
+# =============================================================================
+# Patching Helpers
+# =============================================================================
+
+def create_patch_embedding_model(vectors=None, vector_size=128):
+    """Create a patcher for the embedding model binding.
+
+    This patches the entire _bind_embedding_model flow to return a mock model.
+    """
+    if vectors is None:
+        vectors = np.random.rand(1, vector_size).astype(np.float32)
+
+    mock_model = MagicMock()
+    mock_model.encode.return_value = (vectors, 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)
+
+    return patch(
+        "rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name",
+        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(),
+    )
+
+
+def create_patch_docstore_insert():
+    """Create a patcher for docStoreConn.insert that always succeeds."""
+    return patch(
+        "common.settings.docStoreConn",
+        new_callable=MagicMock,
+    )
+
+
+def create_patch_storage_binary(binary_data=b"fake pdf content"):
+    """Create a patcher for storage retrieval."""
+    mock_async = AsyncMock(return_value=binary_data)
+    return patch(
+        "rag.svr.task_executor_refactor.task_handler.File2DocumentService.get_storage_address",
+        return_value=("bucket_test", "name_test"),
+    ), patch(
+        "rag.svr.task_executor_refactor.task_handler.thread_pool_exec",
+        new_callable=MagicMock,
+        return_value=mock_async,
+    )
+
+
+def create_patch_parser_chunking(chunks=None):
+    """Create a patcher for the parser chunking to return predefined chunks.
+
+    Args:
+        chunks: List of chunk dicts to return from the parser.
+                If None, returns a default single chunk.
+    """
+    if chunks is None:
+        chunks = [{
+            "content_with_weight": "This is a test chunk content.",
+            "page_num_int": [0],
+            "top_int": [0],
+            "position_int": [0, 0, 0, 0],
+        }]
+
+    mock_async = AsyncMock(return_value=chunks)
+    return patch(
+        "rag.svr.task_executor_refactor.chunk_service.thread_pool_exec",
+        new_callable=MagicMock,
+        return_value=mock_async,
+    )
+
+
+# =============================================================================
+# Shared Helper Functions for Integration Tests
+# =============================================================================
+
+
+def create_mock_embedding_model(vector_size: int = 128):
+    """Create a mock embedding model that returns deterministic vectors matching input size."""
+    mock_model = MagicMock()
+
+    def mock_encode(texts):
+        n = len(texts) if isinstance(texts, list) else 1
+        return (
+            np.random.rand(n, vector_size).astype(np.float32),
+            10 * n,
+        )
+
+    mock_model.encode = mock_encode
+    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)
+    return mock_model
+
+
+def create_mock_chat_model():
+    """Create a mock chat model."""
+    mock_model = MagicMock()
+    mock_model.llm_name = "mock_chat"
+    mock_model.__enter__ = MagicMock(return_value=mock_model)
+    mock_model.__exit__ = MagicMock(return_value=False)
+    return mock_model
+
+
+def create_mock_settings():
+    """Create a mock settings object with STORAGE_IMPL and docStoreConn."""
+    mock_settings = MagicMock()
+    mock_settings.STORAGE_IMPL = MagicMock()
+    mock_settings.STORAGE_IMPL.get = MagicMock(return_value=b"fake binary content")
+    mock_settings.docStoreConn = MagicMock()
+    mock_settings.docStoreConn.create_idx = MagicMock(return_value=None)
+    mock_settings.docStoreConn.insert = MagicMock(return_value=None)
+    mock_settings.docStoreConn.delete = MagicMock(return_value=None)
+    mock_settings.docStoreConn.index_exist = MagicMock(return_value=True)
+    mock_settings.docStoreConn.search = MagicMock(return_value={"hits": []})
+    mock_settings.DOC_MAXIMUM_SIZE = 100 * 1024 * 1024  # 100MB
+    mock_settings.DOC_BULK_SIZE = 100
+    mock_settings.retriever = MagicMock()
+    return mock_settings
+
+
+def create_default_chunks(count: int = 2) -> List[Dict[str, Any]]:
+    """Create default chunk dictionaries for testing."""
+    chunks = []
+    for i in range(count):
+        chunks.append({
+            "id": f"chunk_{i}_{uuid.uuid4().hex[:6]}",
+            "content_with_weight": f"This is test chunk content number {i}.",
+            "page_num_int": [i],
+            "top_int": [i * 100],
+            "position_int": [i, 0, i + 1, 0],
+            "doc_id": "doc_test",
+            "kb_id": "kb_test",
+            "docnm_kwd": "test_document.pdf",
+        })
+    return chunks
+
+
+def create_mock_chunk_service(chunks=None):
+    """Create a mock ChunkService instance."""
+    if chunks is None:
+        chunks = create_default_chunks(count=3)
+    mock_service = MagicMock()
+    mock_service.build_chunks = AsyncMock(return_value=chunks)
+    mock_service.insert_chunks = AsyncMock(return_value=True)
+    return mock_service
+
+
+@pytest.fixture
+def mock_embedding_model_factory():
+    """Provide a factory for mock embedding models."""
+    return create_mock_embedding_model
+
+
+@pytest.fixture
+def mock_chat_model_factory():
+    """Provide a factory for mock chat models."""
+    return create_mock_chat_model
+
+
+@pytest.fixture
+def mock_settings_factory():
+    """Provide a factory for mock settings."""
+    return create_mock_settings
+
+
+@pytest.fixture
+def mock_chunk_service_factory():
+    """Provide a factory for mock chunk services."""
+    return create_mock_chunk_service
+
+
+# =============================================================================
+# RaptorService Fixtures
+# =============================================================================
+
+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
+
+
+@pytest.fixture
+def mock_raptor_context():
+    """Provide a mock TaskContext for RaptorService tests."""
+    return create_mock_raptor_context()

test/unit_test/rag/svr/task_executor_refactor/test_chunk_builder.py

@@ -0,0 +1,219 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for ChunkBuilder module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.chunk_builder import (
+    get_parser,
+    run_chunking,
+    extract_outline,
+)
+
+
+class TestGetParser:
+    """Tests for get_parser function."""
+
+    @pytest.mark.parametrize("parser_id", [
+        "naive", "general", "table", "paper", "book",
+        "picture", "audio", "email", "presentation", "manual",
+        "laws", "qa", "resume", "one", "tag",
+    ])
+    def test_get_parser_returns_non_none(self, parser_id):
+        """Test that get_parser returns non-None for all parser types."""
+        parser = get_parser(parser_id)
+        assert parser is not None
+
+    def test_get_parser_kg(self):
+        """Test getting kg parser (maps to naive)."""
+        from common.constants import ParserType
+        parser = get_parser(ParserType.KG.value)
+        assert parser is not None
+
+
+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()
+
+        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.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:
+            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)
+
+    @pytest.mark.asyncio
+    async def test_run_chunking_exception(self):
+        """Test chunking handles exception."""
+        ctx = self._create_mock_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()
+        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()
+
+        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
+            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()
+
+        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)
+
+    @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
+
+        await extract_outline([], ctx)
+
+        mock_rec_ctx.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()
+
+        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"
+        )
+
+    @pytest.mark.asyncio
+    async def test_extract_outline_persistence_exception(self):
+        """Test outline extraction handles persistence exception."""
+        ctx = self._create_mock_context()
+
+        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)

test/unit_test/rag/svr/task_executor_refactor/test_chunk_post_processor.py

@@ -0,0 +1,460 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for ChunkPostProcessor module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.chunk_post_processor import (
+    extract_keywords,
+    generate_questions,
+    generate_metadata,
+    apply_tags,
+    count_with_key,
+    build_metadata_config,
+)
+
+
+class TestExtractKeywords:
+    """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()
+        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_by_type_and_name") 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]
+
+    @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)
+        docs = [{"content_with_weight": "This is test content"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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]
+
+    @pytest.mark.asyncio
+    async def test_extract_keywords_empty_docs(self):
+        """Test keyword extraction with empty docs list."""
+        ctx = self._create_mock_context()
+        docs = []
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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()
+
+
+class TestGenerateQuestions:
+    """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"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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]
+
+    @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)
+        docs = [{"content_with_weight": "This is test content"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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]
+
+
+class TestGenerateMetadata:
+    """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"}},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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()
+
+    @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"}},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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"
+                            )
+
+
+class TestApplyTags:
+    """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"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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
+
+    @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"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") 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}
+
+                    await apply_tags(docs, ctx)
+
+                    # Should return early due to cancellation
+
+
+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
+
+    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
+
+    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
+
+    def test_count_with_key_empty_docs(self):
+        """Test counting with empty docs list."""
+        result = count_with_key([], "tag")
+        assert result == 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
+
+    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
+
+
+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": {}}
+
+    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": [],
+        }
+        result = build_metadata_config(parser_config)
+        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"}]
+
+    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"}]
+
+    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"}]
+
+    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
+
+    def test_missing_metadata_key(self):
+        """When parser_config has no 'metadata' key, built_in_metadata alone is returned."""
+        parser_config = {"built_in_metadata": []}
+        result = build_metadata_config(parser_config)
+        assert result == []

test/unit_test/rag/svr/task_executor_refactor/test_chunk_service.py

@@ -0,0 +1,453 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for ChunkService module.
+
+Note: After refactoring, some functionality has been moved to:
+- chunk_builder.py: Parser factory, run_chunking, extract_outline
+- chunk_post_processor.py: Keyword extraction, question generation, metadata, tagging
+
+This test file now focuses on ChunkService-specific functionality:
+- build_chunks orchestration
+- _prepare_docs_and_upload
+- insert_chunks and related methods
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.chunk_service import ChunkService
+
+
+class TestChunkServiceInit:
+    """Tests for ChunkService initialization."""
+
+    def test_init_stores_task_context(self):
+        """Test that task context is stored."""
+        ctx = MagicMock()
+        service = ChunkService(ctx=ctx)
+        assert service._task_context is ctx
+
+
+class TestChunkServiceBuildChunks:
+    """Tests for build_chunks method."""
+
+    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
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 1000  # Small limit
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            result = await service.build_chunks(b"test binary")
+
+            assert result == []
+            mock_rec_ctx.record.assert_any_call("file_size_exceeded", True)
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_file_size_ok(self):
+        """Test build_chunks proceeds when file size is within limit."""
+        ctx = self._create_mock_context(size=1000)
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 10000000  # Large limit
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
+                mock_parser = MagicMock()
+                mock_get_parser.return_value = mock_parser
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
+                    mock_run_chunking.return_value = [{"content_with_weight": "test"}]
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
+                        with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
+                            mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
+
+                            await service.build_chunks(b"test binary")
+
+                            mock_rec_ctx.record.assert_any_call("file_size_exceeded", False)
+                            mock_rec_ctx.record.assert_any_call("parser_id", "naive")
+                            mock_get_parser.assert_called_once_with("naive")
+
+    @pytest.mark.asyncio
+    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})
+
+        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.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()
+
+
+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()
+        service = ChunkService(ctx=ctx)
+
+        cks = [{"content_with_weight": "test chunk"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.STORAGE_IMPL = MagicMock()
+            mock_settings.STORAGE_IMPL.put = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.image2id", new_callable=AsyncMock):
+
+                docs = await service._prepare_docs_and_upload(cks)
+
+                assert len(docs) == 1
+                assert docs[0]["doc_id"] == "doc_1"
+                assert docs[0]["kb_id"] == "kb_1"
+
+    @pytest.mark.asyncio
+    async def test_prepare_docs_and_upload_with_pagerank(self):
+        """Test document preparation with pagerank."""
+        ctx = self._create_mock_context()
+        ctx.pagerank = 5
+        service = ChunkService(ctx=ctx)
+
+        cks = [{"content_with_weight": "test chunk"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.STORAGE_IMPL = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.image2id", new_callable=AsyncMock):
+
+                docs = await service._prepare_docs_and_upload(cks)
+
+                assert docs[0].get("pagerank_fea") == 5
+
+
+class TestChunkServiceInsertChunks:
+    """Tests for insert_chunks method."""
+
+    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()
+        service = ChunkService(ctx=ctx)
+
+        chunks = [
+            {"id": "chunk_1", "content_with_weight": "test1"},
+            {"id": "chunk_2", "content_with_weight": "test2"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_BULK_SIZE = 100
+            mock_settings.docStoreConn = MagicMock()
+            mock_settings.docStoreConn.insert = MagicMock(return_value=None)
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
+                mock_index.return_value = "test_index"
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
+                    mock_thread.return_value = None
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.TaskService") as mock_task:
+                        mock_task.update_chunk_ids = MagicMock()
+
+                        result = await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
+
+                        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_insert_chunks_canceled(self):
+        """Test chunk insertion when task is canceled."""
+        ctx = self._create_mock_context()
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        service = ChunkService(ctx=ctx)
+
+        chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_BULK_SIZE = 100
+            mock_settings.docStoreConn = MagicMock()
+            mock_settings.docStoreConn.insert = MagicMock(return_value=None)
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
+                mock_index.return_value = "test_index"
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
+                    mock_thread.return_value = None
+
+                    result = await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
+
+                    assert result is False
+                    ctx.progress_cb.assert_called_with(-1, msg="Task has been canceled.")
+
+    @pytest.mark.asyncio
+    async def test_insert_chunks_doc_store_error(self):
+        """Test chunk insertion when doc store returns error."""
+        ctx = self._create_mock_context()
+        service = ChunkService(ctx=ctx)
+
+        chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_BULK_SIZE = 100
+            mock_settings.docStoreConn = MagicMock()
+            mock_settings.docStoreConn.insert = MagicMock(return_value="Error message")
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
+                mock_index.return_value = "test_index"
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
+                    mock_thread.return_value = "Error"
+
+                    with pytest.raises(Exception, match="Insert chunk error"):
+                        await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
+
+
+class TestChunkServiceCreateMotherChunks:
+    """Tests for _create_mother_chunks class method."""
+
+    def test_create_mother_chunks_with_mom_field(self):
+        """Test creating mother chunks from mom field."""
+        chunks = [
+            {"id": "chunk_1", "mom": "Summary text 1", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 1
+        assert mothers[0]["content_with_weight"] == "Summary text 1"
+        assert mothers[0]["available_int"] == 0
+
+    def test_create_mother_chunks_with_mom_with_weight_field(self):
+        """Test creating mother chunks from mom_with_weight field."""
+        chunks = [
+            {"id": "chunk_1", "mom_with_weight": "Summary text 2", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 1
+        assert mothers[0]["content_with_weight"] == "Summary text 2"
+
+    def test_create_mother_chunks_no_mom_field(self):
+        """Test creating mother chunks when no mom field present."""
+        chunks = [
+            {"id": "chunk_1", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 0
+
+    def test_create_mother_chunks_empty_mom(self):
+        """Test creating mother chunks with empty mom field."""
+        chunks = [
+            {"id": "chunk_1", "mom": "", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 0
+
+    def test_create_mother_chunks_deduplicates_ids(self):
+        """Test that mother chunks deduplicate by ID."""
+        chunks = [
+            {"id": "chunk_1", "mom": "Same summary", "content_with_weight": "test1"},
+            {"id": "chunk_2", "mom": "Same summary", "content_with_weight": "test2"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 1
+
+    def test_create_mother_chunks_filters_fields(self):
+        """Test that mother chunks only keep allowed fields."""
+        chunks = [
+            {"id": "chunk_1", "mom": "Summary", "extra_field": "should be removed", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert "extra_field" not in mothers[0]
+        assert "id" in mothers[0]
+        assert "content_with_weight" in mothers[0]

test/unit_test/rag/svr/task_executor_refactor/test_comparator.py

@@ -0,0 +1,598 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for Comparator module.
+"""
+
+from rag.svr.task_executor_refactor.report_generator import (
+    ComparisonResult,
+    ComparisonReport,
+)
+from rag.svr.task_executor_refactor.comparator import (
+    ContextComparator,
+)
+from rag.svr.task_executor_refactor.recording_context import RecordingContext
+
+
+class TestComparisonResult:
+    """Tests for ComparisonResult dataclass."""
+
+    def test_init_with_required_fields(self):
+        """Test initialization with required fields."""
+        result = ComparisonResult(key="test_key", match=True)
+        assert result.key == "test_key"
+        assert result.match is True
+        assert result.production_value is None
+        assert result.dry_run_value is None
+        assert result.diff_details is None
+
+    def test_init_with_all_fields(self):
+        """Test initialization with all fields."""
+        result = ComparisonResult(
+            key="test_key",
+            match=False,
+            production_value=100,
+            dry_run_value=200,
+            diff_details="Values differ"
+        )
+        assert result.key == "test_key"
+        assert result.match is False
+        assert result.production_value == 100
+        assert result.dry_run_value == 200
+        assert result.diff_details == "Values differ"
+
+    def test_to_dict_match(self):
+        """Test to_dict for matching result."""
+        result = ComparisonResult(key="key", match=True)
+        d = result.to_dict()
+        assert d == {"key": "key", "match": True, "diff_details": None}
+
+    def test_to_dict_mismatch(self):
+        """Test to_dict for mismatching result."""
+        result = ComparisonResult(
+            key="key",
+            match=False,
+            diff_details="Difference"
+        )
+        d = result.to_dict()
+        assert d == {"key": "key", "match": False, "diff_details": "Difference"}
+
+
+class TestComparisonReport:
+    """Tests for ComparisonReport dataclass."""
+
+    def test_init_with_required_fields(self):
+        """Test initialization with required fields."""
+        report = ComparisonReport(task_id="task_123")
+        assert report.task_id == "task_123"
+        assert report.total_keys == 0
+        assert report.matched_keys == 0
+        assert report.mismatched_keys == 0
+        assert report.missing_in_production == []
+        assert report.missing_in_dry_run == []
+        assert report.details == []
+
+    def test_summary_no_keys(self):
+        """Test summary when no keys to compare."""
+        report = ComparisonReport(task_id="task_123")
+        assert "No keys to compare" in report.summary()
+
+    def test_summary_with_keys(self):
+        """Test summary with keys."""
+        report = ComparisonReport(
+            task_id="task_123",
+            total_keys=10,
+            matched_keys=8,
+            mismatched_keys=2
+        )
+        summary = report.summary()
+        assert "8/10" in summary
+        assert "80.0%" in summary
+
+    def test_to_dict(self):
+        """Test to_dict serialization."""
+        report = ComparisonReport(
+            task_id="task_123",
+            total_keys=1,
+            matched_keys=1,
+            details=[ComparisonResult(key="k", match=True)]
+        )
+        d = report.to_dict()
+        assert d["task_id"] == "task_123"
+        assert d["total_keys"] == 1
+        assert len(d["details"]) == 1
+
+    def test_to_markdown(self):
+        """Test to_markdown serialization."""
+        report = ComparisonReport(
+            task_id="task_123",
+            total_keys=1,
+            matched_keys=1,
+            mismatched_keys=0,
+            missing_in_production=[],
+            missing_in_dry_run=[],
+            details=[ComparisonResult(key="k", match=True)]
+        )
+        md = report.to_markdown()
+        assert "# Comparison Report: task_123" in md
+        assert "## Summary" in md
+        assert "## Details" in md
+
+    def test_to_markdown_empty_details(self):
+        """Test to_markdown with no details."""
+        report = ComparisonReport(task_id="task_123")
+        md = report.to_markdown()
+        assert "No comparison details" in md
+
+
+class TestContextComparatorInit:
+    """Tests for ContextComparator initialization."""
+
+    def test_init_default_tolerance(self):
+        """Test initialization with default tolerance."""
+        comparator = ContextComparator()
+        assert comparator.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."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_compare_none_values(self):
+        """Test comparing None values."""
+        result = self.comparator.compare_value("key", None, None)
+        assert result.match is True
+
+    def test_compare_one_none(self):
+        """Test comparing when one value is None."""
+        result = self.comparator.compare_value("key", 1, None)
+        assert result.match is False
+        assert "None" in result.diff_details
+
+    def test_compare_equal_strings(self):
+        """Test comparing equal strings."""
+        result = self.comparator.compare_value("key", "hello", "hello")
+        assert result.match is True
+
+    def test_compare_different_strings(self):
+        """Test comparing different strings."""
+        result = self.comparator.compare_value("key", "hello", "world")
+        assert result.match is False
+
+    def test_compare_equal_booleans(self):
+        """Test comparing equal booleans."""
+        result = self.comparator.compare_value("key", True, True)
+        assert result.match is True
+
+    def test_compare_different_booleans(self):
+        """Test comparing different booleans."""
+        result = self.comparator.compare_value("key", True, False)
+        assert result.match is False
+
+    def test_compare_equal_integers(self):
+        """Test comparing equal integers."""
+        result = self.comparator.compare_value("key", 42, 42)
+        assert result.match is True
+
+    def test_compare_equal_floats_within_tolerance(self):
+        """Test comparing equal floats within tolerance."""
+        result = self.comparator.compare_value("key", 1.0000001, 1.0000002)
+        assert result.match is True
+
+    def test_compare_different_floats_exceeding_tolerance(self):
+        """Test comparing floats exceeding tolerance."""
+        result = self.comparator.compare_value("key", 1.0, 2.0)
+        assert result.match is False
+        assert "exceeds tolerance" in result.diff_details
+
+    def test_compare_equal_lists(self):
+        """Test comparing equal lists."""
+        result = self.comparator.compare_value("key", [1, 2, 3], [1, 2, 3])
+        assert result.match is True
+
+    def test_compare_different_length_lists(self):
+        """Test comparing lists with different lengths."""
+        result = self.comparator.compare_value("key", [1, 2], [1, 2, 3])
+        assert result.match is False
+        assert "Length differs" in result.diff_details
+
+    def test_compare_equal_dicts(self):
+        """Test comparing equal dicts."""
+        result = self.comparator.compare_value("key", {"a": 1}, {"a": 1})
+        assert result.match is True
+
+    def test_compare_different_dicts(self):
+        """Test comparing different dicts."""
+        result = self.comparator.compare_value("key", {"a": 1}, {"a": 2})
+        assert result.match is False
+
+    def test_compare_chunks_key_uses_chunk_comparison(self):
+        """Test that chunk keys use chunk comparison strategy."""
+        result = self.comparator.compare_value(
+            "raw_chunks",
+            [{"id": "1", "content_with_weight": "a"}],
+            [{"id": "1", "content_with_weight": "a"}]
+        )
+        assert result.match is True
+
+
+class TestContextComparatorCompareLists:
+    """Tests for _compare_lists method."""
+
+    def test_equal_lists(self):
+        """Test comparing equal lists."""
+        result = ContextComparator._compare_lists("key", [1, 2], [1, 2])
+        assert result.match is True
+
+    def test_different_length_lists(self):
+        """Test comparing lists with different lengths."""
+        result = ContextComparator._compare_lists("key", [1], [1, 2])
+        assert result.match is False
+
+    def test_different_elements(self):
+        """Test comparing lists with different elements."""
+        result = ContextComparator._compare_lists("key", [1, 2], [1, 3])
+        assert result.match is False
+
+
+class TestContextComparatorCompareDicts:
+    """Tests for _compare_dicts method."""
+
+    def test_equal_dicts(self):
+        """Test comparing equal dicts."""
+        result = ContextComparator._compare_dicts("key", {"a": 1}, {"a": 1})
+        assert result.match is True
+
+    def test_dicts_different_keys(self):
+        """Test comparing dicts with different keys."""
+        result = ContextComparator._compare_dicts("key", {"a": 1}, {"b": 1})
+        assert result.match is False
+        assert "Keys differ" in result.diff_details
+
+    def test_dicts_same_keys_different_values(self):
+        """Test comparing dicts with same keys but different values."""
+        result = ContextComparator._compare_dicts("key", {"a": 1}, {"a": 2})
+        assert result.match is False
+
+
+class TestContextComparatorCompareNumbers:
+    """Tests for _compare_numbers method."""
+
+    def test_equal_numbers(self):
+        """Test comparing equal numbers."""
+        comparator = ContextComparator()
+        result = comparator._compare_numbers("key", 1.0, 1.0)
+        assert result.match is True
+
+    def test_numbers_within_tolerance(self):
+        """Test comparing numbers within tolerance."""
+        comparator = ContextComparator(float_tolerance=0.1)
+        result = comparator._compare_numbers("key", 1.0, 1.05)
+        assert result.match is True
+
+    def test_numbers_exceeding_tolerance(self):
+        """Test comparing numbers exceeding tolerance."""
+        comparator = ContextComparator(float_tolerance=0.01)
+        result = comparator._compare_numbers("key", 1.0, 1.1)
+        assert result.match is False
+
+
+class TestContextComparatorCompareChunks:
+    """Tests for _compare_chunks method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_equal_chunks(self):
+        """Test comparing equal chunk lists."""
+        prod = [{"id": "1", "content_with_weight": "a"}]
+        dry = [{"id": "1", "content_with_weight": "a"}]
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is True
+
+    def test_different_count_chunks(self):
+        """Test comparing chunks with different counts."""
+        prod = [{"id": "1"}]
+        dry = [{"id": "1"}, {"id": "2"}]
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+        assert "Chunk count differs" in result.diff_details
+
+    def test_different_ids_chunks(self):
+        """Test comparing chunks with different IDs."""
+        prod = [{"id": "1"}]
+        dry = [{"id": "2"}]
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+        assert "Chunk IDs differ" in result.diff_details
+
+    def test_empty_chunks_lists(self):
+        """Test comparing empty chunk lists."""
+        result = self.comparator._compare_chunks("raw_chunks", [], [])
+        assert result.match is True
+
+    def test_all_chunks_compared_not_sampled(self):
+        """Test that ALL chunks are compared, not just samples.
+        
+        This test creates 10 chunks where only the middle one (index 5) differs.
+        With the old sampling strategy, this difference might be missed.
+        With full comparison, the difference should always be detected.
+        """
+        prod = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(10)]
+        dry = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(10)]
+        # Only modify chunk at index 5 (which might not be sampled in old strategy)
+        dry[5]["content_with_weight"] = "different_content"
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+        assert "Content differs" in result.diff_details
+
+    def test_all_chunks_detect_first_difference(self):
+        """Test that first chunk difference is detected."""
+        prod = [{"id": "1", "content_with_weight": "a"}, {"id": "2", "content_with_weight": "b"}]
+        dry = [{"id": "1", "content_with_weight": "different"}, {"id": "2", "content_with_weight": "b"}]
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+
+    def test_all_chunks_detect_last_difference(self):
+        """Test that last chunk difference is detected."""
+        prod = [{"id": "1", "content_with_weight": "a"}, {"id": "2", "content_with_weight": "b"}]
+        dry = [{"id": "1", "content_with_weight": "a"}, {"id": "2", "content_with_weight": "different"}]
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+
+    def test_all_chunks_large_list_all_match(self):
+        """Test that large list of chunks all match."""
+        prod = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        dry = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is True
+
+    def test_all_chunks_large_list_one_mismatch(self):
+        """Test that a single mismatch in a large list is detected."""
+        prod = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        dry = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        # Modify only the last chunk
+        dry[99]["content_with_weight"] = "different"
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+
+
+class TestContextComparatorExtractChunkIds:
+    """Tests for _extract_chunk_ids method."""
+
+    def test_extract_ids_from_valid_chunks(self):
+        """Test extracting IDs from valid chunks."""
+        chunks = [{"id": "1"}, {"id": "2"}, {"id": "3"}]
+        ids = ContextComparator._extract_chunk_ids(chunks)
+        assert ids == {"1", "2", "3"}
+
+    def test_extract_ids_from_empty_chunks(self):
+        """Test extracting IDs from empty list."""
+        ids = ContextComparator._extract_chunk_ids([])
+        assert ids == set()
+
+    def test_extract_ids_from_chunks_without_id(self):
+        """Test extracting IDs from chunks without id field."""
+        chunks = [{"content": "a"}, {"id": "1"}]
+        ids = ContextComparator._extract_chunk_ids(chunks)
+        assert ids == {"1"}
+
+
+class TestContextComparatorGetChunkId:
+    """Tests for _get_chunk_id method."""
+
+    def test_get_id_from_valid_chunk(self):
+        """Test getting ID from valid chunk."""
+        chunk = {"id": "123"}
+        assert ContextComparator._get_chunk_id(chunk) == "123"
+
+    def test_get_id_from_chunk_without_id(self):
+        """Test getting ID from chunk without id."""
+        chunk = {"content": "a"}
+        assert ContextComparator._get_chunk_id(chunk) == ""
+
+    def test_get_id_from_non_dict(self):
+        """Test getting ID from non-dict."""
+        assert ContextComparator._get_chunk_id("not a dict") == ""
+
+
+class TestContextComparatorCompare:
+    """Tests for compare method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_compare_empty_contexts(self):
+        """Test comparing empty contexts."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.total_keys == 0
+
+    def test_compare_matching_values(self):
+        """Test comparing contexts with matching values."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key", "value")
+        ctx2.record("key", "value")
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.matched_keys == 1
+        assert report.mismatched_keys == 0
+
+    def test_compare_mismatching_values(self):
+        """Test comparing contexts with mismatching values."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key1", "value1")
+        ctx2.record("key1", "value2")
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.mismatched_keys == 1
+
+    def test_compare_missing_key_in_one_context(self):
+        """Test comparing when key is missing in one context."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key1", "value1")
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert "key1" in report.missing_in_dry_run
+
+    def test_compare_with_specific_keys(self):
+        """Test comparing with specific keys list."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key1", "value1")
+        ctx1.record("key2", "value2")
+        ctx2.record("key1", "value1")
+        ctx2.record("key2", "value2")
+        report = self.comparator.compare("task_1", ctx1, ctx2, comparison_keys=["key1"])
+        assert report.total_keys == 1
+
+    def test_compare_filters_out_time_keys(self):
+        """Test that _time keys are filtered out."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("operation_time", 1.0)
+        ctx2.record("operation_time", 1.0)
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.total_keys == 0
+
+
+class TestContextComparatorStripNonDeterministicFields:
+    """Tests for _strip_non_deterministic_fields method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_strip_seconds_from_dict_value(self):
+        """Test that 'seconds' key is removed from dict values."""
+        data = {
+            "graphrag_result": {"seconds": 45.48, "status": "done"},
+            "other_key": "value"
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert "seconds" not in result["graphrag_result"]
+        assert result["graphrag_result"] == {"status": "done"}
+        assert result["other_key"] == "value"
+
+    def test_strip_seconds_from_multiple_dict_values(self):
+        """Test that 'seconds' is removed from multiple dict values."""
+        data = {
+            "result1": {"seconds": 10.0, "count": 5},
+            "result2": {"seconds": 20.0, "name": "test"},
+            "simple_key": 123
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert result["result1"] == {"count": 5}
+        assert result["result2"] == {"name": "test"}
+        assert result["simple_key"] == 123
+
+    def test_strip_does_not_modify_original_dict(self):
+        """Test that the original dict is not modified in place."""
+        data = {
+            "result": {"seconds": 1.0, "value": "test"}
+        }
+        _ = data["result"].copy()
+        self.comparator._strip_non_deterministic_fields(data)
+        # The original nested dict should still have seconds since we only do shallow copy
+        assert "seconds" in data["result"]
+
+    def test_strip_with_empty_dict_values(self):
+        """Test handling of empty dict values."""
+        data = {
+            "empty_dict": {},
+            "normal_key": "value"
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert result["empty_dict"] == {}
+        assert result["normal_key"] == "value"
+
+    def test_strip_with_non_dict_values(self):
+        """Test that non-dict values are not affected."""
+        data = {
+            "string_val": "test",
+            "int_val": 42,
+            "list_val": [1, 2, 3],
+            "dict_val": {"seconds": 1.0, "name": "test"}
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert result["string_val"] == "test"
+        assert result["int_val"] == 42
+        assert result["list_val"] == [1, 2, 3]
+        assert result["dict_val"] == {"name": "test"}
+
+    def test_strip_seconds_from_graphrag_result(self):
+        """Test the specific case from the bug report: graphrag_result with seconds."""
+        prod_data = {
+            "graphrag_result": {
+                "seconds": 45.48,
+                "status": "success",
+                "entity_count": 100
+            }
+        }
+        dry_run_data = {
+            "graphrag_result": {
+                "seconds": 0.99,
+                "status": "success",
+                "entity_count": 100
+            }
+        }
+        prod_stripped = self.comparator._strip_non_deterministic_fields(prod_data)
+        dry_run_stripped = self.comparator._strip_non_deterministic_fields(dry_run_data)
+        
+        # After stripping, both should be equal (except for seconds)
+        assert prod_stripped["graphrag_result"] == {"status": "success", "entity_count": 100}
+        assert dry_run_stripped["graphrag_result"] == {"status": "success", "entity_count": 100}
+        assert prod_stripped["graphrag_result"] == dry_run_stripped["graphrag_result"]
+
+    def test_compare_with_seconds_in_dict_values(self):
+        """Test that compare correctly handles dict values with 'seconds' field."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("graphrag_result", {"seconds": 45.48, "status": "success"})
+        ctx2.record("graphrag_result", {"seconds": 0.99, "status": "success"})
+        
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        # Should match because seconds is stripped
+        assert report.matched_keys == 1
+        assert report.mismatched_keys == 0
+
+    def test_compare_with_different_dict_values_excluding_seconds(self):
+        """Test that compare correctly detects differences in dict values (excluding seconds)."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("graphrag_result", {"seconds": 45.48, "status": "success", "count": 100})
+        ctx2.record("graphrag_result", {"seconds": 0.99, "status": "failed", "count": 50})
+        
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        # Should mismatch because status and count differ
+        assert report.mismatched_keys == 1
+        assert report.matched_keys == 0

test/unit_test/rag/svr/task_executor_refactor/test_constants.py

@@ -0,0 +1,43 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for constants module.
+"""
+
+import pytest
+from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID
+
+
+class TestConstants:
+    """Tests for constants module."""
+
+    def test_canvas_debug_doc_id_exists(self):
+        """Test that CANVAS_DEBUG_DOC_ID constant exists."""
+        assert CANVAS_DEBUG_DOC_ID is not None
+
+    @pytest.mark.parametrize("expected_type", [str])
+    def test_canvas_debug_doc_id_type(self, expected_type):
+        """Test that CANVAS_DEBUG_DOC_ID is a string."""
+        assert isinstance(CANVAS_DEBUG_DOC_ID, expected_type)
+
+    @pytest.mark.parametrize("expected_value", ["dataflow_x"])
+    def test_canvas_debug_doc_id_value(self, expected_value):
+        """Test that CANVAS_DEBUG_DOC_ID has expected value."""
+        assert CANVAS_DEBUG_DOC_ID == expected_value
+
+    def test_canvas_debug_doc_id_not_empty(self):
+        """Test that CANVAS_DEBUG_DOC_ID is not empty."""
+        assert len(CANVAS_DEBUG_DOC_ID) > 0

test/unit_test/rag/svr/task_executor_refactor/test_dataflow_service.py

@@ -0,0 +1,381 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for DataflowService module.
+
+Tests validate behavior through the public run_dataflow() entry point.
+Private orchestration helpers (_process_chunks, _encode_batch, _normalize_chunks,
+_get_output_type, _embed_chunks, _load_dsl, etc.) are exercised implicitly; no test
+reaches directly into those internals.
+"""
+
+import pytest
+from unittest.mock import MagicMock, AsyncMock, patch
+
+from rag.svr.task_executor_refactor.dataflow_service import DataflowService
+
+
+class TestDataflowServiceRunDataflow:
+    """Tests for the public run_dataflow() method.
+
+    Internal helpers (_load_dsl, _normalize_chunks, _get_output_type, _process_chunks,
+    _embed_chunks, _encode_batch) are exercised through this single entry point so
+    the suite stays resilient when internal method boundaries change.
+    """
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.PipelineOperationLogService")
+    async def test_run_dataflow_dsl_not_found(self, mock_pipeline_log, mock_canvas, task_context):
+        """Test run_dataflow returns early when DSL is not found."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        mock_canvas.get_by_id.return_value = (False, None)
+
+        service = DataflowService(ctx=task_context)
+        with pytest.raises(AssertionError, match="User pipeline not found"):
+            await service.run_dataflow()
+
+    @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_empty_chunks(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow handles empty pipeline output."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value={})
+        mock_pipeline_class.return_value = mock_pipeline
+
+        with patch.object(DataflowService, '_record_pipeline_log'):
+            service = DataflowService(ctx=task_context)
+            await service.run_dataflow()
+
+    @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_chunks_output(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow processes 'chunks' output type end-to-end."""
+        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 = {
+            "chunks": [
+                {"text": "Hello world", "content_with_weight": "Hello world"},
+            ],
+            "embedding_token_consumption": 5,
+        }
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value=chunks)
+        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()
+
+                            # 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()
+
+    @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_embedding_failure(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow handles embedding failure gracefully."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_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 = {
+            "chunks": [
+                {"text": "Hello"},
+            ],
+            "embedding_token_consumption": 1,
+        }
+        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=(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()
+
+    @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_billing_hook_success(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow calls billing hook on success."""
+        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 = {
+            "chunks": [
+                {"text": "Hello"},
+            ],
+            "embedding_token_consumption": 1,
+        }
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value=chunks)
+        mock_pipeline_class.return_value = mock_pipeline
+
+        billing_hook = MagicMock()
+        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()
+
+                            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")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_with_billing_hook_error(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow calls billing hook on error."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        task_context._task["name"] = "test.pdf"
+        task_context._write_interceptor = None
+
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(side_effect=Exception("Pipeline failure"))
+        mock_pipeline_class.return_value = mock_pipeline
+
+        billing_hook = MagicMock()
+        billing_hook.on_pipeline_success = AsyncMock()
+        billing_hook.on_pipeline_error = AsyncMock()
+
+        service = DataflowService(ctx=task_context, billing_hook=billing_hook)
+        with pytest.raises(Exception, match="Pipeline failure"):
+            await service.run_dataflow()
+
+        billing_hook.on_pipeline_error.assert_called_once()
+        billing_hook.on_pipeline_success.assert_not_called()
+
+
+class TestDataflowServiceNormalizeChunks:
+    """Tests for _normalize_chunks — stable pure helper for output-format normalization."""
+
+    def test_normalize_chunks_from_chunks_key(self):
+        """Test normalization from 'chunks' key."""
+        result = DataflowService._normalize_chunks({"chunks": [{"a": 1}]})
+        assert result == [{"a": 1}]
+
+    def test_normalize_chunks_from_json_key(self):
+        """Test normalization from 'json' key."""
+        result = DataflowService._normalize_chunks({"json": [{"a": 1}]})
+        assert result == [{"a": 1}]
+
+    def test_normalize_chunks_from_markdown_key(self):
+        """Test normalization from 'markdown' key."""
+        result = DataflowService._normalize_chunks({"markdown": "# Title"})
+        assert result == [{"text": ["# Title"]}]
+
+    def test_normalize_chunks_from_text_key(self):
+        """Test normalization from 'text' key."""
+        result = DataflowService._normalize_chunks({"text": "plain text"})
+        assert result == [{"text": ["plain text"]}]
+
+    def test_normalize_chunks_from_html_key(self):
+        """Test normalization from 'html' key."""
+        result = DataflowService._normalize_chunks({"html": "<p>content</p>"})
+        assert result == [{"text": ["<p>content</p>"]}]
+
+    def test_normalize_chunks_unknown_key(self):
+        """Test normalization with unknown key returns empty."""
+        result = DataflowService._normalize_chunks({"unknown": "data"})
+        assert result == []
+
+    def test_normalize_chunks_empty_markdown(self):
+        """Test normalization with empty markdown value returns empty."""
+        result = DataflowService._normalize_chunks({"markdown": ""})
+        assert result == []
+
+    def test_normalize_chunks_preserves_deepcopy(self):
+        """Test normalization returns a deepcopy so mutations don't leak."""
+        input_data = {"chunks": [{"key": "value"}]}
+        result = DataflowService._normalize_chunks(input_data)
+        result[0]["key"] = "modified"
+        assert input_data["chunks"][0]["key"] == "value"
+
+
+class TestDataflowServiceGetOutputType:
+    """Tests for _get_output_type — stable pure helper for output-type detection."""
+
+    def test_get_output_type_chunks(self):
+        assert DataflowService._get_output_type({"chunks": []}) == "chunks"
+
+    def test_get_output_type_json(self):
+        assert DataflowService._get_output_type({"json": []}) == "json"
+
+    def test_get_output_type_markdown(self):
+        assert DataflowService._get_output_type({"markdown": ""}) == "markdown"
+
+    def test_get_output_type_text(self):
+        assert DataflowService._get_output_type({"text": ""}) == "text"
+
+    def test_get_output_type_html(self):
+        assert DataflowService._get_output_type({"html": ""}) == "html"
+
+    def test_get_output_type_empty(self):
+        assert DataflowService._get_output_type({}) == "empty"
+
+
+class TestDataflowServiceProcessChunks:
+    """Tests for _process_chunks — stable pure helper for chunk metadata processing."""
+
+    def test_process_chunks_adds_doc_id_and_kb_id(self, task_context):
+        """Test _process_chunks adds doc_id, kb_id, and metadata."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert chunks[0]["doc_id"] == "doc_123"
+        assert "kb_id" in chunks[0]
+        assert "content_with_weight" in chunks[0]
+        assert "text" not in chunks[0]
+
+    def test_process_chunks_generates_id(self, task_context):
+        """Test _process_chunks auto-generates id."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "id" in chunks[0]
+
+    def test_process_chunks_questions_field(self, task_context):
+        """Test _process_chunks processes questions field."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content", "questions": "Q1\nQ2"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "questions" not in chunks[0]
+        assert "question_kwd" in chunks[0]
+
+    def test_process_chunks_summary_field(self, task_context):
+        """Test _process_chunks processes summary field."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content", "summary": "summary text"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "summary" not in chunks[0]
+        assert "content_ltks" in chunks[0]
+
+    def test_process_chunks_metadata_field(self, task_context):
+        """Test _process_chunks extracts metadata."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content", "metadata": {"key": "val"}}]
+        metadata = DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "metadata" not in chunks[0]
+        assert "key" in metadata
+
+
+class TestDataflowServiceInit:
+    """Tests for DataflowService initialization."""
+
+    @patch("rag.svr.task_executor_refactor.dataflow_service.settings")
+    def test_init_with_custom_batch_sizes(self, mock_settings):
+        """Test initialization with custom batch sizes."""
+        ctx = MagicMock()
+        service = DataflowService(ctx=ctx, embedding_batch_size=64, doc_bulk_size=50)
+        assert service._embedding_batch_size == 64
+        assert service._doc_bulk_size == 50
+
+    @patch("rag.svr.task_executor_refactor.dataflow_service.settings")
+    def test_init_with_default_sizes(self, mock_settings):
+        """Test initialization with default batch sizes."""
+        mock_settings.EMBEDDING_BATCH_SIZE = 32
+        mock_settings.DOC_BULK_SIZE = 100
+        ctx = MagicMock()
+        service = DataflowService(ctx=ctx)
+        assert service._embedding_batch_size == 32
+        assert service._doc_bulk_size == 100
+
+    def test_init_stores_context_and_hook(self):
+        """Test initialization stores context and billing hook."""
+        ctx = MagicMock()
+        hook = MagicMock()
+        service = DataflowService(ctx=ctx, billing_hook=hook)
+        assert service._task_context is ctx
+        assert service._billing_hook is hook

test/unit_test/rag/svr/task_executor_refactor/test_embedding_service.py

@@ -0,0 +1,145 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for 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.
+"""
+
+import numpy as np
+from unittest.mock import MagicMock, patch
+
+from rag.svr.task_executor_refactor.embedding_service import EmbeddingService
+
+
+class TestEmbeddingServiceInit:
+    """Tests for EmbeddingService initialization."""
+
+    @patch("rag.svr.task_executor_refactor.embedding_service.settings")
+    def test_init_with_default_batch_size(self, mock_settings):
+        """Test initialization with default batch size."""
+        mock_settings.EMBEDDING_BATCH_SIZE = 32
+        ctx = MagicMock()
+        service = EmbeddingService(ctx=ctx)
+        assert service._embedding_batch_size == 32
+
+    @patch("rag.svr.task_executor_refactor.embedding_service.settings")
+    def test_init_with_custom_batch_size(self, mock_settings):
+        """Test initialization with custom batch size."""
+        ctx = MagicMock()
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=64)
+        assert service._embedding_batch_size == 64
+
+    def test_init_stores_task_context(self):
+        """Test that task context is stored."""
+        ctx = MagicMock()
+        service = EmbeddingService(ctx=ctx)
+        assert service._task_context is ctx
+
+
+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.
+    """
+
+    @patch.object(EmbeddingService, '_run_encode')
+    def test_embed_chunks_basic(self, mock_run_encode):
+        """Test basic chunk embedding."""
+        mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        tk_count, vector_size = 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)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        service.embed_chunks(docs, model)
+
+        mock_run_encode.assert_called()
+
+    @patch.object(EmbeddingService, '_run_encode')
+    def test_embed_chunks_with_title_content_combination(self, mock_run_encode):
+        """Test that title and content vectors are combined."""
+        mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        _, vector_size = 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):
+        """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)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "a" * 200},
+        ]
+        tk_count, vector_size = 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]

test/unit_test/rag/svr/task_executor_refactor/test_embedding_utils.py

@@ -0,0 +1,331 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for EmbeddingUtils module.
+"""
+
+import numpy as np
+from unittest.mock import patch
+from rag.svr.task_executor_refactor.embedding_utils import EmbeddingUtils
+
+
+class TestEmbeddingUtilsPrepareTexts:
+    """Tests for prepare_texts_for_embedding class method."""
+
+    def test_prepare_texts_basic(self):
+        """Test basic text preparation."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+            {"docnm_kwd": "Title2", "content_with_weight": "Content2"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert titles == ["Title1", "Title2"]
+        assert contents == ["Content1", "Content2"]
+
+    def test_prepare_texts_with_question_kwd(self):
+        """Test text preparation with question_kwd."""
+        docs = [
+            {"docnm_kwd": "Title1", "question_kwd": ["Q1", "Q2"], "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert titles == ["Title1"]
+        assert contents == ["Q1\nQ2"]
+
+    def test_prepare_texts_with_empty_question_kwd(self):
+        """Test text preparation with empty question_kwd falls back to content."""
+        docs = [
+            {"docnm_kwd": "Title1", "question_kwd": [], "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert contents == ["Content1"]
+
+    def test_prepare_texts_with_missing_question_kwd(self):
+        """Test text preparation without question_kwd uses content."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert contents == ["Content1"]
+
+    def test_prepare_texts_normalizes_table_html(self):
+        """Test that table HTML tags are normalized."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "<table><tr><td>Cell</td></tr></table>"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        # Table tags should be replaced with spaces
+        assert "<table>" not in contents[0]
+
+    def test_prepare_texts_whitespace_only_becomes_none(self):
+        """Test that whitespace-only content becomes 'None'."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "   \n\n  "},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert contents == ["None"]
+
+    def test_prepare_texts_default_title(self):
+        """Test that missing docnm_kwd uses 'Title' as default."""
+        docs = [
+            {"content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert titles == ["Title"]
+
+    def test_prepare_texts_without_question_kwd(self):
+        """Test text preparation with use_question_kwd=False."""
+        docs = [
+            {"docnm_kwd": "Title1", "question_kwd": ["Q1"], "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs, use_question_kwd=False)
+        assert contents == ["Content1"]
+
+
+class TestEmbeddingUtilsPrepareDataflowTexts:
+    """Tests for prepare_texts_for_dataflow_embedding class method."""
+
+    def test_prepare_dataflow_texts_with_questions(self):
+        """Test dataflow text preparation with questions field."""
+        chunks = [
+            {"questions": "Q1\nQ2"},
+            {"questions": "Q3"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Q1\nQ2", "Q3"]
+
+    def test_prepare_dataflow_texts_with_summary(self):
+        """Test dataflow text preparation with summary field (no questions)."""
+        chunks = [
+            {"summary": "Summary1"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Summary1"]
+
+    def test_prepare_dataflow_texts_with_text(self):
+        """Test dataflow text preparation with text field (no questions/summary)."""
+        chunks = [
+            {"text": "Text content"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Text content"]
+
+    def test_prepare_dataflow_texts_priority(self):
+        """Test field priority: questions > summary > text."""
+        chunks = [
+            {"questions": "Q", "summary": "S", "text": "T"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Q"]
+
+        chunks = [
+            {"summary": "S", "text": "T"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["S"]
+
+
+class TestEmbeddingUtilsTruncateTexts:
+    """Tests for truncate_texts class method."""
+
+    @patch("rag.svr.task_executor_refactor.embedding_utils.truncate")
+    def test_truncate_texts_calls_truncate(self, mock_truncate):
+        """Test truncate_texts calls truncate with correct max_length."""
+        mock_truncate.return_value = "truncated"
+        texts = ["long text 1", "long text 2"]
+        max_length = 100
+
+        _ = EmbeddingUtils.truncate_texts(texts, max_length)
+
+        assert mock_truncate.call_count == 2
+        # Should subtract 10 for safety margin
+        mock_truncate.assert_called_with("long text 2", 90)
+
+    @patch("rag.svr.task_executor_refactor.embedding_utils.truncate")
+    def test_truncate_texts_returns_list(self, mock_truncate):
+        """Test truncate_texts returns a list of same length."""
+        mock_truncate.return_value = "truncated"
+        texts = ["text1", "text2", "text3"]
+        result = EmbeddingUtils.truncate_texts(texts, 50)
+        assert len(result) == 3
+
+
+class TestEmbeddingUtilsStackVectors:
+    """Tests for stack_vectors class method."""
+
+    def test_stack_vectors_with_multiple_batches(self):
+        """Test stacking multiple vector batches."""
+        batch1 = np.array([[1.0, 2.0], [3.0, 4.0]])
+        batch2 = np.array([[5.0, 6.0]])
+        result = EmbeddingUtils.stack_vectors([batch1, batch2])
+        expected = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        np.testing.assert_array_equal(result, expected)
+
+    def test_stack_vectors_with_empty_batches(self):
+        """Test stacking empty batches returns empty array."""
+        result = EmbeddingUtils.stack_vectors([])
+        assert result.size == 0
+
+    def test_stack_vectors_with_single_batch(self):
+        """Test stacking a single batch."""
+        batch = np.array([[1.0, 2.0]])
+        result = EmbeddingUtils.stack_vectors([batch])
+        np.testing.assert_array_equal(result, batch)
+
+
+class TestEmbeddingUtilsAttachVectors:
+    """Tests for attach_vectors class method."""
+
+    def test_attach_vectors_basic(self):
+        """Test attaching vectors to docs."""
+        docs = [{"id": 1}, {"id": 2}]
+        vectors = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+
+        vector_size = EmbeddingUtils.attach_vectors(docs, vectors)
+
+        assert vector_size == 3
+        assert "q_3_vec" in docs[0]
+        assert "q_3_vec" in docs[1]
+        assert docs[0]["q_3_vec"] == [1.0, 2.0, 3.0]
+        assert docs[1]["q_3_vec"] == [4.0, 5.0, 6.0]
+
+    def test_attach_vectors_custom_key_template(self):
+        """Test attaching vectors with custom key template."""
+        docs = [{"id": 1}]
+        vectors = np.array([[1.0, 2.0]])
+
+        EmbeddingUtils.attach_vectors(docs, vectors, vector_key_template="vec_%d")
+
+        assert "vec_2" in docs[0]
+
+    def test_attach_vectors_modifies_in_place(self):
+        """Test that attach_vectors modifies docs in place."""
+        docs = [{"id": 1}]
+        vectors = np.array([[1.0, 2.0]])
+        original_id = id(docs)
+
+        EmbeddingUtils.attach_vectors(docs, vectors)
+
+        assert id(docs) == original_id
+
+
+class TestEmbeddingUtilsCombineVectors:
+    """Tests for combine_title_content_vectors class method."""
+
+    def test_combine_vectors_with_title_and_content(self):
+        """Test combining title and content vectors with weight."""
+        title_vecs = np.array([[1.0, 2.0], [3.0, 4.0]])
+        content_vecs = np.array([[5.0, 6.0], [7.0, 8.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs, title_weight=0.3)
+
+        # Expected: 0.3 * title + 0.7 * content
+        expected = 0.3 * title_vecs + 0.7 * content_vecs
+        np.testing.assert_array_almost_equal(result, expected)
+
+    def test_combine_vectors_with_default_weight(self):
+        """Test combining with default weight when not specified."""
+        title_vecs = np.array([[1.0, 2.0]])
+        content_vecs = np.array([[5.0, 6.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs)
+
+        # Expected: 0.1 * title + 0.9 * content (default weight is 0.1)
+        expected = 0.1 * title_vecs + 0.9 * content_vecs
+        np.testing.assert_array_almost_equal(result, expected)
+
+    def test_combine_vectors_with_none_title(self):
+        """Test combining when title vectors is None returns content."""
+        content_vecs = np.array([[5.0, 6.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(None, content_vecs, title_weight=0.3)
+
+        np.testing.assert_array_equal(result, content_vecs)
+
+    def test_combine_vectors_with_mismatched_shapes(self):
+        """Test combining when shapes don't match returns content."""
+        title_vecs = np.array([[1.0, 2.0]])
+        content_vecs = np.array([[5.0, 6.0], [7.0, 8.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs, title_weight=0.3)
+
+        # Should return content_vecs when shapes don't match
+        np.testing.assert_array_equal(result, content_vecs)
+
+    def test_combine_vectors_with_zero_weight(self):
+        """Test combining when weight is 0 uses default 0.1."""
+        title_vecs = np.array([[1.0, 2.0]])
+        content_vecs = np.array([[5.0, 6.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs, title_weight=0)
+
+        # Should use default weight of 0.1
+        expected = 0.1 * title_vecs + 0.9 * content_vecs
+        np.testing.assert_array_almost_equal(result, expected)
+
+
+class TestEmbeddingUtilsInternals:
+    """Tests for internal helper methods."""
+
+    def test_extract_content_with_question_kwd(self):
+        """Test _extract_content with question_kwd."""
+        doc = {"question_kwd": ["Q1", "Q2"], "content_with_weight": "Content"}
+        result = EmbeddingUtils._extract_content(doc, use_question_kwd=True)
+        assert result == "Q1\nQ2"
+
+    def test_extract_content_without_question_kwd(self):
+        """Test _extract_content without question_kwd."""
+        doc = {"content_with_weight": "Content"}
+        result = EmbeddingUtils._extract_content(doc, use_question_kwd=True)
+        assert result == "Content"
+
+    def test_extract_content_with_use_question_false(self):
+        """Test _extract_content with use_question_kwd=False."""
+        doc = {"question_kwd": ["Q1"], "content_with_weight": "Content"}
+        result = EmbeddingUtils._extract_content(doc, use_question_kwd=False)
+        assert result == "Content"
+
+    def test_normalize_table_html(self):
+        """Test _normalize_table_html removes table tags."""
+        html = "<table><tr><td>Cell</td></tr></table>"
+        result = EmbeddingUtils._normalize_table_html(html)
+        assert "<table>" not in result
+        assert "<tr>" not in result
+        assert "<td>" not in result
+
+    def test_handle_whitespace(self):
+        """Test _handle_whitespace replaces whitespace-only with placeholder."""
+        assert EmbeddingUtils._handle_whitespace("   \n  ") == "None"
+        assert EmbeddingUtils._handle_whitespace("  text  ") == "  text  "
+
+    def test_handle_whitespace_with_empty_string(self):
+        """Test _handle_whitespace with empty string."""
+        assert EmbeddingUtils._handle_whitespace("") == "None"
+
+
+class TestEmbeddingUtilsConstants:
+    """Tests for class constants."""
+
+    def test_default_title_weight(self):
+        """Test DEFAULT_TITLE_WEIGHT value."""
+        assert EmbeddingUtils.DEFAULT_TITLE_WEIGHT == 0.1
+
+    def test_default_title_placeholder(self):
+        """Test DEFAULT_TITLE_PLACEHOLDER value."""
+        assert EmbeddingUtils.DEFAULT_TITLE_PLACEHOLDER == "Title"
+
+    def test_content_placeholder_for_whitespace(self):
+        """Test CONTENT_PLACEHOLDER_FOR_WHITESPACE value."""
+        assert EmbeddingUtils.CONTENT_PLACEHOLDER_FOR_WHITESPACE == "None"

test/unit_test/rag/svr/task_executor_refactor/test_post_processor.py

@@ -0,0 +1,130 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for PostProcessor module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.post_processor import PostProcessor
+
+
+class TestPostProcessorInit:
+    """Tests for PostProcessor initialization."""
+
+    def test_init_stores_task_context(self):
+        """Test that task context is stored."""
+        ctx = MagicMock()
+        service = PostProcessor(ctx=ctx)
+        assert service._task_context is ctx
+
+
+class TestPostProcessorProcessTableParserMetadata:
+    """Tests for process_table_parser_metadata method."""
+
+    @pytest.mark.asyncio
+    async def test_skips_non_table_parser(self):
+        """Test that processing is skipped for non-table parser."""
+        ctx = MagicMock()
+        ctx.parser_id = "naive"
+        service = PostProcessor(ctx=ctx)
+
+        await service.process_table_parser_metadata("doc_1", [])
+
+        # Should return early without any further processing
+
+    @pytest.mark.asyncio
+    async def test_skips_when_not_manual_column_mode(self):
+        """Test that processing is skipped when not in manual column mode."""
+        ctx = MagicMock()
+        ctx.parser_id = "table"
+        ctx.raw_task = {}
+        service = PostProcessor(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge:
+            mock_merge.return_value = {"table_column_mode": "auto"}
+            await service.process_table_parser_metadata("doc_1", [])
+
+            mock_merge.assert_called_once()
+
+
+class TestPostProcessorInsertTocChunk:
+    """Tests for insert_toc_chunk method."""
+
+    @pytest.mark.asyncio
+    async def test_returns_false_for_none_chunk(self):
+        """Test that method returns False when chunk is None."""
+        ctx = MagicMock()
+        service = PostProcessor(ctx=ctx)
+        chunk_service = MagicMock()
+
+        result = await service.insert_toc_chunk(None, chunk_service)
+
+        assert result is False
+        chunk_service.insert_chunks.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_checks_cancellation(self):
+        """Test that cancellation is checked."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        ctx.progress_cb = MagicMock()
+        service = PostProcessor(ctx=ctx)
+        chunk_service = MagicMock()
+        toc_chunk = {"id": "toc_1"}
+
+        result = await service.insert_toc_chunk(toc_chunk, chunk_service)
+
+        assert result is False
+        ctx.progress_cb.assert_called_with(-1, msg="Task has been canceled.")
+
+    @pytest.mark.asyncio
+    async def test_inserts_toc_chunk_successfully(self):
+        """Test successful TOC chunk insertion."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        service = PostProcessor(ctx=ctx)
+        chunk_service = AsyncMock()
+        chunk_service.insert_chunks = AsyncMock(return_value=True)
+        toc_chunk = {"id": "toc_1"}
+
+        result = await service.insert_toc_chunk(toc_chunk, chunk_service)
+
+        assert result is True
+        chunk_service.insert_chunks.assert_called_once_with(
+            "task_1", "tenant_1", "kb_1", [toc_chunk]
+        )
+
+    @pytest.mark.asyncio
+    async def test_handles_insert_failure(self):
+        """Test handling of insert failure."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        service = PostProcessor(ctx=ctx)
+        chunk_service = AsyncMock()
+        chunk_service.insert_chunks = AsyncMock(return_value=False)
+        toc_chunk = {"id": "toc_1"}
+
+        result = await service.insert_toc_chunk(toc_chunk, chunk_service)
+
+        assert result is False

test/unit_test/rag/svr/task_executor_refactor/test_raptor_service.py

@@ -0,0 +1,452 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Tests for RaptorService.
+
+Coverage is driven through the public entry point `run_raptor_for_kb()`.
+
+Design principles:
+- All orchestration behavior is validated through the public API.
+- Only stable pure helpers (`_collect_doc_info`, `_schedule_raptor_cleanup`)
+  are tested directly.
+- Internal methods (`_run_file_level_raptor`, `_run_dataset_level_raptor`,
+  `_should_skip_raptor`, `_load_doc_chunks`, `_load_all_doc_chunks`,
+  `_generate_raptor`, `_get_raptor_chunk_methods`) are NOT tested directly —
+  their behavior is covered by exercising `run_raptor_for_kb()` with
+  appropriate mocks.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from rag.svr.task_executor_refactor.raptor_service import RaptorService
+
+
+# =============================================================================
+# Stable Pure Helpers (tested directly)
+# =============================================================================
+
+
+class TestRaptorServiceInit:
+    """Tests for RaptorService initialization."""
+
+    def test_init_stores_task_context(self, mock_raptor_context):
+        svc = RaptorService(mock_raptor_context)
+        assert svc._task_context is mock_raptor_context
+
+    def test_init_uses_provided_kb_id(self, mock_raptor_context):
+        mock_raptor_context.kb_id = "custom_kb"
+        svc = RaptorService(mock_raptor_context)
+        assert svc._task_context.kb_id == "custom_kb"
+
+
+class TestRaptorServiceCollectDocInfo:
+    """Tests for _collect_doc_info — stable pure data aggregation (classmethod)."""
+
+    def _make_mock_doc(self, name, type, parser_id, parser_config):
+        """Create a mock document with accessible attributes."""
+        mock_doc = MagicMock()
+        mock_doc.name = name
+        mock_doc.type = type
+        mock_doc.parser_id = parser_id
+        mock_doc.parser_config = parser_config
+        return mock_doc
+
+    def test_collect_doc_info_success(self):
+        doc_ids = ["doc_1", "doc_2"]
+
+        mock_doc_1 = self._make_mock_doc(name="", type="pdf", parser_id="naive", parser_config={})
+        mock_doc_2 = self._make_mock_doc(name="doc2.txt", type="txt", parser_id="manual", parser_config={"chunk_token_num": 512})
+
+        def get_by_id_side_effect(doc_id):
+            if doc_id == "doc_1":
+                return True, mock_doc_1
+            if doc_id == "doc_2":
+                return True, mock_doc_2
+            return False, None
+
+        with patch("rag.svr.task_executor_refactor.raptor_service.DocumentService") as mock_ds:
+            mock_ds.get_by_id = MagicMock(side_effect=get_by_id_side_effect)
+            result = RaptorService._collect_doc_info(doc_ids)
+
+        assert len(result) == 2
+        assert result["doc_1"]["name"] == ""
+        assert result["doc_1"]["type"] == "pdf"
+        assert result["doc_1"]["parser_id"] == "naive"
+        assert result["doc_2"]["name"] == "doc2.txt"
+        assert result["doc_2"]["type"] == "txt"
+        assert result["doc_2"]["parser_id"] == "manual"
+        assert result["doc_2"]["parser_config"] == {"chunk_token_num": 512}
+
+    def test_collect_doc_info_empty_input(self):
+        result = RaptorService._collect_doc_info([])
+        assert result == {}
+
+    def test_collect_doc_info_deduplicates_doc_ids(self):
+        """Duplicate doc_ids should be deduplicated."""
+        doc_ids = ["doc_1", "doc_1", "doc_2"]
+
+        mock_doc = self._make_mock_doc(name="test.pdf", type="pdf", parser_id="naive", parser_config={})
+
+        called_ids = []
+
+        def get_by_id_side_effect(doc_id):
+            called_ids.append(doc_id)
+            return True, mock_doc
+
+        with patch("rag.svr.task_executor_refactor.raptor_service.DocumentService") as mock_ds:
+            mock_ds.get_by_id = MagicMock(side_effect=get_by_id_side_effect)
+            result = RaptorService._collect_doc_info(doc_ids)
+
+        assert sorted(called_ids) == ["doc_1", "doc_2"]
+        assert len(result) == 2
+
+    def test_collect_doc_info_missing_document(self):
+        doc_ids = ["doc_1", "missing_doc"]
+
+        mock_doc = self._make_mock_doc(name="test.pdf", type="pdf", parser_id="naive", parser_config={})
+
+        def get_by_id_side_effect(doc_id):
+            if doc_id == "doc_1":
+                return True, mock_doc
+            return False, None
+
+        with patch("rag.svr.task_executor_refactor.raptor_service.DocumentService") as mock_ds:
+            mock_ds.get_by_id = MagicMock(side_effect=get_by_id_side_effect)
+            result = RaptorService._collect_doc_info(doc_ids)
+
+        assert len(result) == 1
+        assert "doc_1" in result
+        assert "missing_doc" not in result
+
+
+class TestRaptorServiceScheduleRaptorCleanup:
+    """Tests for _schedule_raptor_cleanup — stable pure data operation (classmethod)."""
+
+    def test_schedule_cleanup_adds_entry(self):
+        cleanup_list = []
+        RaptorService._schedule_raptor_cleanup("doc_1", "tree_builder_a", cleanup_list)
+        assert cleanup_list == [("doc_1", "tree_builder_a")]
+
+    def test_schedule_cleanup_deduplicates(self):
+        cleanup_list = [("doc_1", "tree_builder_a")]
+        RaptorService._schedule_raptor_cleanup("doc_1", "tree_builder_a", cleanup_list)
+        assert len(cleanup_list) == 1
+
+    def test_schedule_cleanup_keep_method_none(self):
+        cleanup_list = []
+        RaptorService._schedule_raptor_cleanup("doc_1", None, cleanup_list)
+        assert cleanup_list == [("doc_1", None)]
+
+    def test_schedule_cleanup_multiple_docs(self):
+        cleanup_list = []
+        RaptorService._schedule_raptor_cleanup("doc_1", "t1", cleanup_list)
+        RaptorService._schedule_raptor_cleanup("doc_2", "t2", cleanup_list)
+        RaptorService._schedule_raptor_cleanup("doc_3", None, cleanup_list)
+        assert len(cleanup_list) == 3
+        assert ("doc_1", "t1") in cleanup_list
+        assert ("doc_2", "t2") in cleanup_list
+        assert ("doc_3", None) in cleanup_list
+
+
+# =============================================================================
+# Public Entry Point Tests
+# =============================================================================
+
+
+class TestRaptorServiceRunRaptorForKb:
+    """Tests for run_raptor_for_kb() — the public entry point.
+
+    All orchestration behavior (file-level vs dataset-level dispatch,
+    chunk loading, skip logic, cleanup scheduling) is validated through
+    this method by mocking internal helpers and observing:
+    - Return values (chunks, token_count, cleanup_raptor_chunks)
+    - Mock call patterns (which internal method was invoked, with what args)
+    """
+
+    @pytest.fixture
+    def sample_chunks(self):
+        """Sample RAPTOR summary chunks returned by internal methods."""
+        return [{"id": "chunk_1", "content_with_weight": "Summary 1"}]
+
+    @pytest.fixture
+    def raptor_config_file_scope(self):
+        """RAPTOR config with file-level scope."""
+        return {
+            "raptor": {
+                "tree_builder": "raptor",
+                "clustering_method": "gmm",
+                "scope": "file",
+                "prompt": "summarize",
+                "max_token": 512,
+                "threshold": 0.5,
+                "max_cluster": 64,
+                "random_seed": 42,
+            }
+        }
+
+    @pytest.fixture
+    def raptor_config_dataset_scope(self):
+        """RAPTOR config with dataset-level scope."""
+        return {
+            "raptor": {
+                "tree_builder": "raptor",
+                "clustering_method": "gmm",
+                "scope": "dataset",
+                "prompt": "summarize",
+                "max_token": 512,
+                "threshold": 0.5,
+                "max_cluster": 64,
+                "random_seed": 42,
+            }
+        }
+
+    # ---- Basic dispatch (file-level scope) ----
+
+    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."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1", "doc_2"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        vector_size = 128
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+            "doc_2": {"name": "b.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           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()
+
+        mock_file.assert_called_once()
+        mock_dataset.assert_not_called()
+        assert chunks == sample_chunks
+        assert tk_count == 42
+
+    # ---- Basic dispatch (dataset-level scope) ----
+
+    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."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        vector_size = 128
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           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()
+
+        mock_dataset.assert_called_once()
+        mock_file.assert_not_called()
+        assert chunks == sample_chunks
+        assert tk_count == 99
+
+    # ---- Empty / no documents ----
+
+    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()
+        embd_mdl = MagicMock()
+
+        with patch.object(svc, "_collect_doc_info", return_value={}), \
+           patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           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()
+
+        assert chunks == []
+        assert tk_count == 0
+        assert cleanup == []
+
+    # ---- Cleanup scheduling through the public API ----
+
+    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.
+        """
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+
+        expected_cleanup = [("doc_1", "tree_builder_a")]
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), 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()
+
+        assert cleanup == expected_cleanup
+
+    # ---- Dispatch with missing raptor config key ----
+
+    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."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        config = {}  # No raptor key at all
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           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()
+
+        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(
+        self, mock_raptor_context, sample_chunks, raptor_config_file_scope
+    ):
+        """Vector size is used to construct vctr_nm and passed to internal method."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        vector_size = 256
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), 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()
+
+        # Verify _run_file_level_raptor received vctr_nm with the correct vector size
+        # Positional args: 0=raptor_config, 1=tree_builder, 2=clustering_method,
+        #   3=chat_mdl, 4=embd_mdl, 5=vctr_nm
+        positional_args = mock_file.call_args[0]
+        assert positional_args[5] == "q_256_vec"
+
+    # ---- Document info collection through public API ----
+
+    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."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_a"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+
+        expected_info = {"doc_a": {"name": "file.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}}}
+
+        with patch.object(svc, "_collect_doc_info", return_value=expected_info) as mock_collect, \
+           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()
+
+        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

test/unit_test/rag/svr/task_executor_refactor/test_recording_context.py

@@ -0,0 +1,357 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for RecordingContext module.
+"""
+
+import time
+import pytest
+from rag.svr.task_executor_refactor.recording_context import (
+    RecordingContext,
+    get_recording_context,
+    set_recording_context,
+    recording_context_manager,
+    timed_with_recording,
+)
+
+
+class TestRecordingContextInit:
+    """Tests for RecordingContext initialization."""
+
+    def test_init_creates_empty_data(self):
+        """Test that __init__ creates empty _data dict."""
+        ctx = RecordingContext()
+        assert ctx._data == {}
+
+    def test_init_creates_empty_records(self):
+        """Test that __init__ creates empty records list."""
+        ctx = RecordingContext()
+        assert ctx.records == []
+
+
+class TestRecordingContextRecord:
+    """Tests for RecordingContext.record method."""
+
+    def test_record_single_value(self):
+        """Test recording a single value."""
+        ctx = RecordingContext()
+        ctx.record("chunk_count", 100)
+        assert ctx.get("chunk_count") == 100
+
+    def test_record_overwrites_existing_value(self):
+        """Test that recording with same key overwrites previous value."""
+        ctx = RecordingContext()
+        ctx.record("key", "value1")
+        ctx.record("key", "value2")
+        assert ctx.get("key") == "value2"
+
+    def test_record_none_value(self):
+        """Test recording None value."""
+        ctx = RecordingContext()
+        ctx.record("key", None)
+        assert ctx.get("key") is None
+
+    def test_record_complex_object(self):
+        """Test recording a complex object like list or dict."""
+        ctx = RecordingContext()
+        ctx.record("chunks", [{"id": 1}, {"id": 2}])
+        assert ctx.get("chunks") == [{"id": 1}, {"id": 2}]
+
+
+class TestRecordingContextFuncReturnValues:
+    """Tests for function return value recording."""
+
+    def test_save_func_return_value_first_call(self):
+        """Test saving first return value for a function."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("test_func", 42)
+        assert ctx.get_func_return_values("test_func") == [42]
+
+    def test_save_func_return_value_multiple_calls(self):
+        """Test saving multiple return values for same function."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("test_func", 1)
+        ctx.save_func_return_value("test_func", 2)
+        ctx.save_func_return_value("test_func", 3)
+        assert ctx.get_func_return_values("test_func") == [1, 2, 3]
+
+    def test_get_func_return_values_nonexistent_function(self):
+        """Test getting return values for nonexistent function returns empty list."""
+        ctx = RecordingContext()
+        assert ctx.get_func_return_values("nonexistent") == []
+
+    def test_get_func_return_values_multiple_functions(self):
+        """Test getting return values for different functions."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("func_a", "a1")
+        ctx.save_func_return_value("func_b", "b1")
+        ctx.save_func_return_value("func_a", "a2")
+        assert ctx.get_func_return_values("func_a") == ["a1", "a2"]
+        assert ctx.get_func_return_values("func_b") == ["b1"]
+
+
+class TestRecordingContextGet:
+    """Tests for RecordingContext.get method."""
+
+    def test_get_existing_key(self):
+        """Test getting an existing key."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        assert ctx.get("key") == "value"
+
+    def test_get_nonexistent_key_returns_none(self):
+        """Test getting nonexistent key returns None."""
+        ctx = RecordingContext()
+        assert ctx.get("missing") is None
+
+    def test_get_nonexistent_key_returns_default(self):
+        """Test getting nonexistent key returns provided default."""
+        ctx = RecordingContext()
+        assert ctx.get("missing", "default") == "default"
+
+    def test_get_with_none_default(self):
+        """Test getting with None as default."""
+        ctx = RecordingContext()
+        assert ctx.get("missing", None) is None
+
+
+class TestRecordingContextGetAllFuncReturnValues:
+    """Tests for get_all_func_return_values method."""
+
+    def test_get_all_func_return_values_empty(self):
+        """Test getting all values when none recorded."""
+        ctx = RecordingContext()
+        assert ctx.get_all_func_return_values() == {}
+
+    def test_get_all_func_return_values_with_data(self):
+        """Test getting all values with some data."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("func_a", 1)
+        ctx.save_func_return_value("func_b", 2)
+        result = ctx.get_all_func_return_values()
+        assert result == {"func_a": [1], "func_b": [2]}
+
+    def test_get_all_func_return_values_returns_copy(self):
+        """Test that returned dict is a copy, not the original."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("func", 1)
+        result = ctx.get_all_func_return_values()
+        result["func"] = []
+        # Original should be unchanged
+        assert ctx.get_func_return_values("func") == [1]
+
+
+class TestRecordingContextHas:
+    """Tests for RecordingContext.has method."""
+
+    def test_has_existing_key(self):
+        """Test has returns True for existing key."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        assert ctx.has("key") is True
+
+    def test_has_nonexistent_key(self):
+        """Test has returns False for nonexistent key."""
+        ctx = RecordingContext()
+        assert ctx.has("missing") is False
+
+    def test_has_after_clear(self):
+        """Test has returns False after clear."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        ctx.clear()
+        assert ctx.has("key") is False
+
+
+class TestRecordingContextClear:
+    """Tests for RecordingContext.clear method."""
+
+    def test_clear_removes_all_data(self):
+        """Test that clear removes all recorded data."""
+        ctx = RecordingContext()
+        ctx.record("key1", "value1")
+        ctx.record("key2", "value2")
+        ctx.clear()
+        assert ctx._data == {}
+
+    def test_clear_removes_all_records(self):
+        """Test that clear removes all timing records."""
+        ctx = RecordingContext()
+        with ctx.measure("op1"):
+            pass
+        ctx.clear()
+        assert ctx.records == []
+
+
+class TestRecordingContextMeasure:
+    """Tests for RecordingContext.measure context manager."""
+
+    def test_measure_records_elapsed_time(self):
+        """Test that measure records elapsed time."""
+        ctx = RecordingContext()
+        with ctx.measure("test_op"):
+            time.sleep(0.01)
+        assert len(ctx.records) == 1
+        assert ctx.records[0][0] == "test_op"
+        assert ctx.records[0][1] >= 0.01
+
+    def test_measure_multiple_operations(self):
+        """Test measuring multiple operations."""
+        ctx = RecordingContext()
+        with ctx.measure("op1"):
+            time.sleep(0.01)
+        with ctx.measure("op2"):
+            time.sleep(0.02)
+        assert len(ctx.records) == 2
+        assert ctx.records[0][0] == "op1"
+        assert ctx.records[1][0] == "op2"
+
+    def test_measure_preserves_context_on_exception(self):
+        """Test that measure still records time on exception."""
+        ctx = RecordingContext()
+        with pytest.raises(ValueError):
+            with ctx.measure("failing_op"):
+                raise ValueError("test error")
+        assert len(ctx.records) == 1
+        assert ctx.records[0][0] == "failing_op"
+
+
+class TestRecordingContextReset:
+    """Tests for RecordingContext.reset method."""
+
+    def test_reset_clears_data(self):
+        """Test that reset clears all data."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        ctx.reset()
+        assert ctx._data == {}
+
+    def test_reset_clears_records(self):
+        """Test that reset clears all records."""
+        ctx = RecordingContext()
+        with ctx.measure("op"):
+            pass
+        ctx.reset()
+        assert ctx.records == []
+
+
+class TestRecordingContextRepr:
+    """Tests for RecordingContext.__repr__ method."""
+
+    def test_repr_empty_context(self):
+        """Test repr of empty context."""
+        ctx = RecordingContext()
+        assert "RecordingContext" in repr(ctx)
+
+    def test_repr_with_data(self):
+        """Test repr with some data."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        r = repr(ctx)
+        assert "RecordingContext" in r
+        assert "key" in r
+
+
+class TestContextVariableFunctions:
+    """Tests for context variable functions."""
+
+    def test_set_and_get_recording_context(self):
+        """Test set and get recording context."""
+        ctx = RecordingContext()
+        set_recording_context(ctx)
+        assert get_recording_context() is ctx
+
+    def test_set_recording_context_none_unbinds(self):
+        """Test setting None unbinds the context."""
+        ctx = RecordingContext()
+        set_recording_context(ctx)
+        set_recording_context(None)
+        # After unbinding, get should raise RuntimeError
+        with pytest.raises(RuntimeError, match="no context"):
+            get_recording_context()
+
+
+class TestRecordingContextManager:
+    """Tests for recording_context_manager context manager."""
+
+    def test_context_manager_with_provided_context(self):
+        """Test context manager with provided context."""
+        ctx = RecordingContext()
+        with recording_context_manager(ctx) as mgr:
+            assert mgr is ctx
+            mgr.record("key", "value")
+        assert ctx.get("key") == "value"
+
+    def test_context_manager_creates_new_context(self):
+        """Test context manager creates new context when none provided."""
+        with recording_context_manager() as ctx:
+            assert isinstance(ctx, RecordingContext)
+            ctx.record("key", "value")
+            assert ctx.get("key") == "value"
+
+    def test_context_manager_restores_previous_context(self):
+        """Test context manager restores previous context after exit."""
+        outer_ctx = RecordingContext()
+        set_recording_context(outer_ctx)
+        
+        inner_ctx = RecordingContext()
+        with recording_context_manager(inner_ctx):
+            assert get_recording_context() is inner_ctx
+        
+        # After exiting, should restore outer_ctx
+        assert get_recording_context() is outer_ctx
+
+
+class TestTimedWithRecordingDecorator:
+    """Tests for timed_with_recording decorator."""
+
+    def test_decorator_without_parentheses(self):
+        """Test decorator used without parentheses."""
+        ctx = RecordingContext()
+        set_recording_context(ctx)
+        
+        @timed_with_recording
+        def test_func():
+            time.sleep(0.01)
+            return 42
+        
+        result = test_func()
+        assert result == 42
+
+    def test_decorator_with_parentheses_and_context(self):
+        """Test decorator with explicit context."""
+        ctx = RecordingContext()
+        
+        @timed_with_recording(recording_context=ctx)
+        def test_func():
+            time.sleep(0.01)
+            return "hello"
+        
+        result = test_func()
+        assert result == "hello"
+
+    def test_decorator_without_context_raises_error(self):
+        """Test decorator raises RuntimeError when no context is available."""
+        # Ensure no context is set
+        set_recording_context(None)
+        
+        @timed_with_recording
+        def test_func():
+            return 123
+        
+        # Should raise RuntimeError because no context is available
+        with pytest.raises(RuntimeError, match="no context"):
+            test_func()

test/unit_test/rag/svr/task_executor_refactor/test_task_context.py

@@ -0,0 +1,417 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for TaskContext module.
+"""
+
+from unittest.mock import MagicMock
+from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+
+
+def _make_ctx(task, **kwargs):
+    """Helper to create TaskContext with default limiters and callbacks."""
+    return TaskContext(
+        task=task,
+        limiters=kwargs.get("limiters", TaskLimiters()),
+        callbacks=kwargs.get("callbacks", TaskCallbacks()),
+        write_interceptor=kwargs.get("write_interceptor", None),
+    )
+
+
+class TestTaskContextInit:
+    """Tests for TaskContext initialization."""
+
+    def test_init_with_minimal_task(self):
+        """Test initialization with minimal task dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.id == "task_1"
+
+    def test_init_with_all_parameters(self):
+        """Test initialization with all parameters."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        chat_limiter = MagicMock()
+        minio_limiter = MagicMock()
+        chunk_limiter = MagicMock()
+        embed_limiter = MagicMock()
+        kg_limiter = MagicMock()
+        write_interceptor = MagicMock()
+        progress_callback = MagicMock()
+        has_canceled_func = MagicMock()
+
+        ctx = TaskContext(
+            task=task,
+            limiters=TaskLimiters(
+                chat=chat_limiter,
+                minio=minio_limiter,
+                chunk=chunk_limiter,
+                embed=embed_limiter,
+                kg=kg_limiter,
+            ),
+            callbacks=TaskCallbacks(
+                progress=progress_callback,
+                has_canceled=has_canceled_func,
+            ),
+            write_interceptor=write_interceptor,
+        )
+
+        assert ctx.chat_limiter is chat_limiter
+        assert ctx.minio_limiter is minio_limiter
+        assert ctx.chunk_limiter is chunk_limiter
+        assert ctx.embed_limiter is embed_limiter
+        assert ctx.kg_limiter is kg_limiter
+        assert ctx.write_interceptor is write_interceptor
+        assert ctx.callbacks.progress is progress_callback
+        assert ctx.has_canceled_func is has_canceled_func
+
+    def test_init_defaults_for_callbacks(self):
+        """Test that callbacks default to no-op functions."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        # Should not raise
+        ctx.callbacks.progress()
+        assert ctx.has_canceled_func("task_1") is False
+
+
+class TestTaskContextIdentityProperties:
+    """Tests for task identity properties."""
+
+    def test_id(self):
+        """Test id property."""
+        task = {"id": "task_123", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.id == "task_123"
+
+    def test_tenant_id(self):
+        """Test tenant_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.tenant_id == "tenant_1"
+
+    def test_kb_id_default(self):
+        """Test kb_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_id == ""
+
+    def test_kb_id(self):
+        """Test kb_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "kb_id": "kb_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_id == "kb_1"
+
+    def test_doc_id_default(self):
+        """Test doc_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_id == ""
+
+    def test_doc_id(self):
+        """Test doc_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "doc_id": "doc_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_id == "doc_1"
+
+    def test_doc_ids_default(self):
+        """Test doc_ids property defaults to empty list."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_ids == []
+
+    def test_doc_ids(self):
+        """Test doc_ids property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "doc_ids": ["doc_1", "doc_2"]}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_ids == ["doc_1", "doc_2"]
+
+
+class TestTaskContextDocumentMetadataProperties:
+    """Tests for document metadata properties."""
+
+    def test_name_default(self):
+        """Test name property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.name == ""
+
+    def test_name(self):
+        """Test name property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "name": "test.pdf"}
+        ctx = _make_ctx(task=task)
+        assert ctx.name == "test.pdf"
+
+    def test_location_default(self):
+        """Test location property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.location == ""
+
+    def test_size_default(self):
+        """Test size property defaults to 0."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.size == 0
+
+    def test_size(self):
+        """Test size property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "size": 1024}
+        ctx = _make_ctx(task=task)
+        assert ctx.size == 1024
+
+
+class TestTaskContextParserProperties:
+    """Tests for parser configuration properties."""
+
+    def test_parser_id_default(self):
+        """Test parser_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_id == ""
+
+    def test_parser_id(self):
+        """Test parser_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "parser_id": "naive"}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_id == "naive"
+
+    def test_parser_config_default(self):
+        """Test parser_config property defaults to empty dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_config == {}
+
+    def test_parser_config(self):
+        """Test parser_config property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "parser_config": {"chunk_size": 512}}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_config == {"chunk_size": 512}
+
+    def test_kb_parser_config_default(self):
+        """Test kb_parser_config property defaults to empty dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_parser_config == {}
+
+    def test_kb_parser_config(self):
+        """Test kb_parser_config property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "kb_parser_config": {"language": "en"}}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_parser_config == {"language": "en"}
+
+
+class TestTaskContextLanguageAndModelProperties:
+    """Tests for language and model properties."""
+
+    def test_language_default(self):
+        """Test language property defaults to 'en'."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.language == "en"
+
+    def test_language(self):
+        """Test language property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "language": "zh"}
+        ctx = _make_ctx(task=task)
+        assert ctx.language == "zh"
+
+    def test_llm_id_default(self):
+        """Test llm_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.llm_id == ""
+
+    def test_llm_id(self):
+        """Test llm_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "llm_id": "gpt-4"}
+        ctx = _make_ctx(task=task)
+        assert ctx.llm_id == "gpt-4"
+
+    def test_embd_id_default(self):
+        """Test embd_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.embd_id == ""
+
+    def test_embd_id(self):
+        """Test embd_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "embd_id": "text-embedding-ada-002"}
+        ctx = _make_ctx(task=task)
+        assert ctx.embd_id == "text-embedding-ada-002"
+
+
+class TestTaskContextPageRangeProperties:
+    """Tests for page range properties."""
+
+    def test_from_page_default(self):
+        """Test from_page property defaults to 0."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.from_page == 0
+
+    def test_from_page(self):
+        """Test from_page property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "from_page": 10}
+        ctx = _make_ctx(task=task)
+        assert ctx.from_page == 10
+
+    def test_to_page_default(self):
+        """Test to_page property defaults to -1."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.to_page == -1
+
+    def test_to_page(self):
+        """Test to_page property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "to_page": 100}
+        ctx = _make_ctx(task=task)
+        assert ctx.to_page == 100
+
+
+class TestTaskContextTaskTypeAndRoutingProperties:
+    """Tests for task type and routing properties."""
+
+    def test_task_type_default(self):
+        """Test task_type property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.task_type == ""
+
+    def test_task_type(self):
+        """Test task_type property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "task_type": "raptor"}
+        ctx = _make_ctx(task=task)
+        assert ctx.task_type == "raptor"
+
+    def test_dataflow_id_default(self):
+        """Test dataflow_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.dataflow_id == ""
+
+    def test_dataflow_id(self):
+        """Test dataflow_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "dataflow_id": "flow_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.dataflow_id == "flow_1"
+
+
+class TestTaskContextAdditionalProperties:
+    """Tests for additional properties."""
+
+    def test_pagerank_default(self):
+        """Test pagerank property defaults to 0."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.pagerank == 0
+
+    def test_pagerank(self):
+        """Test pagerank property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "pagerank": 10}
+        ctx = _make_ctx(task=task)
+        assert ctx.pagerank == 10
+
+    def test_file_default(self):
+        """Test file property defaults to None."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.file is None
+
+    def test_file(self):
+        """Test file property."""
+        file_obj = MagicMock()
+        task = {"id": "task_1", "tenant_id": "tenant_1", "file": file_obj}
+        ctx = _make_ctx(task=task)
+        assert ctx.file is file_obj
+
+
+class TestTaskContextMemoryProperties:
+    """Tests for memory task properties."""
+
+    def test_memory_id_default(self):
+        """Test memory_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.memory_id == ""
+
+    def test_memory_id(self):
+        """Test memory_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "memory_id": "mem_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.memory_id == "mem_1"
+
+    def test_source_id_default(self):
+        """Test source_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.source_id == ""
+
+    def test_source_id(self):
+        """Test source_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "source_id": "src_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.source_id == "src_1"
+
+    def test_message_dict_default(self):
+        """Test message_dict property defaults to empty dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.message_dict == {}
+
+    def test_message_dict(self):
+        """Test message_dict property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "message_dict": {"key": "value"}}
+        ctx = _make_ctx(task=task)
+        assert ctx.message_dict == {"key": "value"}
+
+
+class TestTaskContextRawTask:
+    """Tests for raw_task property and get method."""
+
+    def test_raw_task_returns_original_dict(self):
+        """Test raw_task returns the original task dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "custom_key": "value"}
+        ctx = _make_ctx(task=task)
+        assert ctx.raw_task is task
+
+    def test_get_existing_key(self):
+        """Test get method with existing key."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "custom_key": "value"}
+        ctx = _make_ctx(task=task)
+        assert ctx.get("custom_key") == "value"
+
+    def test_get_nonexistent_key_returns_none(self):
+        """Test get method with nonexistent key returns None."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.get("missing") is None
+
+    def test_get_with_default(self):
+        """Test get method with default value."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.get("missing", "default") == "default"
+
+
+class TestTaskContextProgressCallback:
+    """Tests for progress callback functionality."""
+
+    def test_progress_cb_is_set_in_init(self):
+        """Test that _progress_cb is set during initialization."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        # _progress_cb should be set in __init__
+        assert hasattr(ctx, '_progress_cb')
+        assert ctx._progress_cb is not None

test/unit_test/rag/svr/task_executor_refactor/test_task_handler.py

@@ -0,0 +1,300 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for TaskHandler module.
+
+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.
+
+Stable pure helpers (_build_toc, _get_vector_size) are tested directly since they
+are side-effect-free data transformations.
+"""
+
+import pytest
+import numpy as np
+from unittest.mock import MagicMock, AsyncMock, patch
+
+from rag.svr.task_executor_refactor.task_handler import TaskHandler
+
+
+class TestTaskHandlerHandleTask:
+    """Tests for the public handle_task() entry point."""
+
+    @pytest.mark.asyncio
+    async def test_handle_task_calls_handle(self):
+        """Test handle_task delegates to handle()."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.doc_id = "doc_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        handler = TaskHandler(ctx=ctx)
+        handler.handle = AsyncMock()
+        await handler.handle_task()
+        handler.handle.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_task_cleanup_on_cancel(self):
+        """Test handle_task cleans up docStore when canceled."""
+        from common import settings
+        mock_doc_store = MagicMock()
+        mock_doc_store.index_exist = MagicMock(return_value=True)
+        mock_doc_store.delete = MagicMock(return_value=None)
+        orig = settings.docStoreConn
+        settings.docStoreConn = mock_doc_store
+        try:
+            ctx = MagicMock()
+            ctx.id = "task_1"
+            ctx.tenant_id = "tenant_1"
+            ctx.kb_id = "kb_1"
+            ctx.doc_id = "doc_1"
+            ctx.has_canceled_func = MagicMock(return_value=True)
+            ctx.recording_context = MagicMock()
+            handler = TaskHandler(ctx=ctx)
+            handler.handle = AsyncMock(side_effect=Exception("test error"))
+            # 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
+
+
+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.
+    """
+
+    @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)
+
+        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()
+            await handler.handle()
+            mock_handle.assert_called_once_with(ctx.raw_task)
+
+    @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)
+
+        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()
+
+        handler = TaskHandler(ctx=ctx)
+        await handler.handle()
+        ctx.progress_cb.assert_called_once_with(-1, msg="Task has been canceled.")
+
+    @pytest.mark.asyncio
+    async def test_handle_standard_chunking(self):
+        """Test handle 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
+
+        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()
+
+        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()
+
+        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()
+
+        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()
+
+        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()
+
+        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)
+
+        handler = TaskHandler(ctx=ctx)
+        handler._bind_embedding_model = AsyncMock(return_value=None)
+
+        await handler.handle()
+        # Should not call _run_standard_chunking when model is None
+        assert not hasattr(handler, '_run_standard_chunking_called')
+
+
+class TestTaskHandlerGetVectorSize:
+    """Tests for _get_vector_size — stable pure helper."""
+
+    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
+
+
+class TestTaskHandlerBuildToc:
+    """Tests for _build_toc — stable pure helper (requires LLM mocking)."""
+
+    def test_build_toc_with_empty_docs(self):
+        """Test _build_toc returns None when run_toc_from_text returns empty."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+
+        docs = [{"id": "chunk_1", "content_with_weight": "text", "page_num_int": [1], "top_int": [0]}]
+
+        def mock_asyncio_run(coro):
+            # Close the coroutine to prevent "never awaited" warnings
+            coro.close()
+            return []
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_cfg:
+            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
+
+    def test_build_toc_with_results(self):
+        """Test _build_toc builds TOC chunk when results exist."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+
+        docs = [{"id": "chunk_0", "content_with_weight": "text", "doc_id": "doc_1", "page_num_int": [1], "top_int": [0]}]
+        toc_result = [{"chunk_id": "0", "title": "Section 1"}]
+
+        def mock_asyncio_run(coro):
+            # 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_by_type_and_name") as mock_cfg:
+            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
+
+
+class TestTaskHandlerInit:
+    """Tests for TaskHandler initialization."""
+
+    def test_init_stores_context_and_hook(self):
+        ctx = MagicMock()
+        hook = MagicMock()
+        handler = TaskHandler(ctx=ctx, billing_hook=hook)
+        assert handler._task_context is ctx
+        assert handler._billing_hook is hook
+
+    def test_init_default_hook_none(self):
+        ctx = MagicMock()
+        handler = TaskHandler(ctx=ctx)
+        assert handler._billing_hook is None

test/unit_test/rag/svr/task_executor_refactor/test_task_handler_integration.py

@@ -0,0 +1,993 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Integration tests for TaskHandler orchestration.
+"""
+
+import asyncio
+import gc
+import uuid
+from typing import Any, Dict
+from unittest.mock import MagicMock, AsyncMock, patch
+
+import pytest
+
+from rag.svr.task_executor_refactor.task_handler import TaskHandler
+from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+from rag.svr.task_executor_refactor.recording_context import BaseRecordingContext, RecordingContext
+from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID
+
+# Import shared helpers from conftest
+from test.unit_test.rag.svr.task_executor_refactor.conftest import (
+    AsyncMockLimiter,
+    create_mock_embedding_model,
+    create_default_chunks,
+    create_mock_settings,
+    create_mock_chunk_service,
+)
+
+
+def create_task_context(
+    task_dict: Dict[str, Any],
+    is_canceled: bool = False,
+    recording_context: BaseRecordingContext | None = None,
+) -> TaskContext:
+    """Create a real TaskContext with mocked limiters and callbacks.
+
+    Args:
+        task_dict: Task dictionary with all task attributes.
+        is_canceled: If True, has_canceled_func returns True.
+        recording_context: RecordingContext to inject. If None, a new one
+            is created automatically so that recording_context access works.
+
+    Returns:
+        TaskContext with all required dependencies injected.
+    """
+    if recording_context is None:
+        recording_context = RecordingContext()
+    limiter = AsyncMockLimiter()
+    progress_callback = MagicMock()
+    ctx = TaskContext(
+        task=task_dict,
+        limiters=TaskLimiters(
+            chat=limiter,
+            minio=limiter,
+            chunk=limiter,
+            embed=limiter,
+            kg=limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=progress_callback,
+            has_canceled=MagicMock(return_value=is_canceled),
+        ),
+        recording_context=recording_context,
+    )
+    # Add progress_callback property for task_handler compatibility
+    ctx.progress_callback = progress_callback
+    # Add set_progress_cb method for task_handler compatibility
+    ctx.set_progress_cb = lambda cb: setattr(ctx.callbacks, 'progress_cb', cb)
+    return ctx
+
+
+# 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()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            task_status = recording_ctx.get("task_status")
+            assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_insertion_result(self):
+        """Verify that insertion_result is recorded as 'success'."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            insertion_result = recording_ctx.get("insertion_result")
+            assert insertion_result == "success", f"Expected insertion_result='success', got {insertion_result}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_chunk_ids(self):
+        """Verify that chunk_ids_count is recorded after build_chunks."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunks = create_default_chunks(count=3)
+        mock_chunk_service = create_mock_chunk_service(chunks=mock_chunks)
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+            mock_run_toc.return_value = []  # TOC returns empty when not enabled
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            chunk_ids_count = recording_ctx.get("chunk_ids_count")
+            assert chunk_ids_count is not None, "chunk_ids_count should be recorded"
+            assert chunk_ids_count == 3, f"Expected chunk_ids_count=3, got {chunk_ids_count}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_token_count(self):
+        """Verify that token_count and vector_size are recorded after embedding."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            token_count = recording_ctx.get("token_count")
+            vector_size = recording_ctx.get("vector_size")
+
+            assert token_count is not None, "token_count should be recorded"
+            assert vector_size is not None, "vector_size should be recorded"
+            assert vector_size == 128, f"Expected vector_size=128, got {vector_size}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_progress_callback_invoked(self):
+        """Verify that progress_callback is invoked multiple times during pipeline."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            ctx.progress_callback.assert_called()
+            call_count = ctx.progress_callback.call_count
+            assert call_count > 0, "progress_callback should have been invoked at least once"
+
+
+class TestTaskCancellationCleanupIntegration:
+    """P0: Integration tests for task cancellation cleanup flow."""
+
+    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()
+        ctx = create_task_context(task_dict, is_canceled=True)
+        mock_settings = create_mock_settings()
+
+        call_log = []
+
+        def mock_thread_impl(func, *args, **kwargs):
+            # Get the actual method name from the mock
+            func_repr = repr(func)
+            call_log.append(func_repr)
+            if 'index_exist' in func_repr:
+                return True
+            if 'delete' in func_repr:
+                return {"result": "deleted"}
+            return {"result": "deleted"}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_index"), \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec", side_effect=mock_thread_impl):
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle_task()
+
+            # Verify delete was called by checking the call log
+            delete_calls = [c for c in call_log if 'delete' in c]
+            assert len(delete_calls) >= 1, f"Expected at least one delete call, got: {call_log}"
+
+    @pytest.mark.asyncio
+    async def test_canceled_task_progress_callback_with_negative_one(self):
+        """Verify that progress_callback is called with -1 when task is canceled."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict, is_canceled=True)
+        mock_settings = create_mock_settings()
+
+        def mock_thread_impl(func, *args, **kwargs):
+            func_repr = repr(func)
+            if 'index_exist' in func_repr:
+                return True
+            if 'delete' in func_repr:
+                return {"result": "deleted"}
+            return {"result": "deleted"}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_index"), \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec", side_effect=mock_thread_impl):
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle_task()
+
+            ctx.progress_callback.assert_called()
+            call_args_list = ctx.progress_callback.call_args_list
+            # Check for -1 in any position of the call arguments
+            has_negative_progress = False
+            for call in call_args_list:
+                # Check positional args
+                for arg in call[0]:
+                    if arg == -1:
+                        has_negative_progress = True
+                        break
+                # Check keyword args
+                if call[1].get("prog") == -1:
+                    has_negative_progress = True
+                if has_negative_progress:
+                    break
+            assert has_negative_progress, f"progress_callback should have been called with -1 progress. Calls: {call_args_list}"
+
+    @pytest.mark.asyncio
+    async def test_canceled_task_does_not_proceed_to_chunking(self):
+        """Verify that canceled task does not proceed to embedding model binding."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict, is_canceled=True)
+        mock_settings = create_mock_settings()
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
+
+            mock_index_name.return_value = "test_index"
+            mock_settings.docStoreConn.index_exist.return_value = True
+            mock_settings.docStoreConn.delete.return_value = {"result": "deleted"}
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return {"result": "deleted"}
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle_task()
+
+            mock_bundle.assert_not_called()
+
+
+class TestRaptorPipelineIntegration:
+    """P1: Integration tests for the RAPTOR pipeline."""
+
+    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()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_kb = MagicMock()
+        mock_kb.id = "kb_test"
+        mock_kb.parser_config = {"raptor": {"use_raptor": False}}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.RaptorService") as mock_raptor_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_index_name.return_value = "test_index"
+            mock_kb_service.get_by_id.return_value = (True, mock_kb)
+            mock_kb_service.update_by_id.return_value = True
+            mock_raptor_service.return_value.run_raptor_for_kb = AsyncMock(return_value=([], 0, []))
+            mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
+            mock_doc_service.increment_chunk_num = MagicMock()
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return None
+
+            mock_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            task_status = recording_ctx.get("task_status")
+            assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
+
+    @pytest.mark.asyncio
+    async def test_raptor_pipeline_enables_raptor_if_not_configured(self):
+        """Verify that RAPTOR is enabled if not already configured."""
+        task_dict = self._create_raptor_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_kb = MagicMock()
+        mock_kb.id = "kb_test"
+        mock_kb.parser_config = {"raptor": {"use_raptor": False}}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.RaptorService") as mock_raptor_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_index_name.return_value = "test_index"
+            mock_kb_service.get_by_id.return_value = (True, mock_kb)
+            mock_kb_service.update_by_id.return_value = True
+            mock_raptor_service.return_value.run_raptor_for_kb = AsyncMock(return_value=([], 0, []))
+            mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
+            mock_doc_service.increment_chunk_num = MagicMock()
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return None
+
+            mock_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            # Check that the kb parser_config was updated
+            mock_kb_service.update_by_id.assert_called_once()
+            call_args = mock_kb_service.update_by_id.call_args
+            update_dict = call_args[0][1]
+            assert update_dict.get("parser_config", {}).get("raptor", {}).get("use_raptor") is True, \
+                "RAPTOR should be enabled in parser_config after running"
+
+
+class TestEmbeddingModelBindingFailureIntegration:
+    """P1: Integration tests for embedding model binding failure."""
+
+    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()
+        ctx = create_task_context(task_dict)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
+
+            mock_get_config.side_effect = Exception("Model not found")
+            mock_get_default.side_effect = Exception("Model not found")
+
+            handler = TaskHandler(ctx=ctx)
+
+            with pytest.raises(Exception, match="Model not found"):
+                await handler.handle()
+
+    @pytest.mark.asyncio
+    async def test_embedding_binding_failure_calls_progress_callback(self):
+        """Verify that embedding model binding failure calls progress_callback."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
+
+            mock_get_config.side_effect = Exception("Model not found")
+            mock_get_default.side_effect = Exception("Model not found")
+
+            handler = TaskHandler(ctx=ctx)
+
+            with pytest.raises(Exception):
+                await handler.handle()
+
+            ctx.progress_callback.assert_called()
+
+
+class TestDataflowPipelineIntegration:
+    """P2: Integration tests for the dataflow pipeline."""
+
+    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()
+        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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_dataflow_service.assert_called_once()
+            mock_instance.run_dataflow.assert_called_once()
+
+
+class TestTocAsyncFlowIntegration:
+    """P2: Integration tests for TOC async flow."""
+
+    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()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls, \
+             patch("rag.svr.task_executor_refactor.post_processor.DocumentService") as mock_post_doc_service:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+            mock_run_toc.return_value = [{"title": "Test TOC", "level": 1}]
+            mock_post_doc_service.increment_chunk_num = MagicMock()
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_run_toc.assert_called()
+
+        # Explicit cleanup to prevent resource leaks
+        del mock_embedding, mock_settings, mock_chunk_service
+        del mock_get_config, mock_get_default, mock_bundle, mock_file_service
+        del mock_index_name, mock_doc_service, mock_chunk_service_cls, mock_run_toc, mock_post_doc_service
+        del mock_thread_exec, mock_chunk_thread_exec
+        # Allow pending callbacks to execute
+        await asyncio.sleep(0)
+        gc.collect()
+
+    @pytest.mark.asyncio(loop_scope="function")
+    @pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
+    async def test_toc_async_flow_does_not_create_thread_when_disabled(self):
+        """Verify that TOC async flow does not create a thread when disabled.
+        
+        Note: This test has a known issue with resource leaks (unclosed sockets and
+        event loops) when run as part of the full test suite. The warning filter
+        above suppresses these warnings temporarily. The root cause is related to
+        asyncio.to_thread creating new event loops that are not properly cleaned up
+        by pytest-asyncio.
+        """
+
+        task_dict = self._create_toc_enabled_task_dict()
+        task_dict["parser_config"]["toc_extraction"] = False
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_run_toc.assert_not_called()
+
+        # Explicit cleanup to prevent resource leaks
+        del mock_embedding, mock_settings, mock_chunk_service
+        del mock_get_config, mock_get_default, mock_bundle, mock_file_service
+        del mock_index_name, mock_doc_service, mock_chunk_service_cls, mock_run_toc
+        del mock_thread_exec, mock_chunk_thread_exec
+        # Allow pending callbacks to execute and close event loop
+        await asyncio.sleep(0)
+        # Cancel all pending tasks
+        current_task = asyncio.current_task()
+        pending = [t for t in asyncio.all_tasks() if t is not current_task and not t.done()]
+        for task in pending:
+            task.cancel()
+        if pending:
+            await asyncio.gather(*pending, return_exceptions=True)
+        gc.collect()
+
+
+class TestRecordingContextDataFlowAssertions:
+    """P2: Integration tests for RecordingContext data flow assertions."""
+
+    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()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            file_size_exceeded = recording_ctx.get("file_size_exceeded")
+            assert file_size_exceeded is None or file_size_exceeded is False, \
+                f"Expected file_size_exceeded to be False/None for small file, got {file_size_exceeded}"
+
+    @pytest.mark.asyncio
+    async def test_recording_context_captures_parser_id(self):
+        """Verify that RecordingContext captures parser_id from task context."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            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
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            # parser_id is available in the task context, verify task completion
+            task_status = recording_ctx.get("task_status")
+            assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
+            # Verify the parser_id is accessible from the task context
+            assert ctx.parser_id == "naive", f"Expected parser_id='naive', got {ctx.parser_id}"

test/unit_test/rag/svr/task_executor_refactor/test_task_raptor_utils.py

@@ -0,0 +1,219 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for rag/svr/task_executor_refactor/raptor_utils.py module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch
+from rag.svr.task_executor_refactor.raptor_utils import (
+    get_raptor_chunk_field_map,
+    delete_raptor_chunks,
+)
+
+
+class TestGetRaptorChunkFieldMap:
+    """Tests for get_raptor_chunk_field_map function."""
+
+    @pytest.mark.asyncio
+    async def test_returns_primary_result_when_raptor_chunks_exist(self):
+        """Test that primary result is returned when RAPTOR chunks exist."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        mock_doc_store.search.return_value = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}}
+        mock_doc_store.get_fields.return_value = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}}
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                async def mock_exec(*args, **kwargs):
+                    return {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}}
+                mock_thread.side_effect = mock_exec
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = {"chunk_1"}
+
+                    result = await get_raptor_chunk_field_map("doc_1", "tenant_1", "kb_1")
+
+                    assert "chunk_1" in result
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_falls_back_to_secondary_search_when_no_raptor_chunks(self):
+        """Test that fallback search is used when no RAPTOR chunks found."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            call_count = 0
+            async def mock_exec(*args, **kwargs):
+                nonlocal call_count
+                call_count += 1
+                if call_count == 1:
+                    return {}  # Primary returns empty
+                else:
+                    return {"chunk_1": {"raptor_kwd": "raptor"}}  # Fallback
+
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                mock_thread.side_effect = mock_exec
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = set()  # Primary has no RAPTOR chunks
+
+                    _ = await get_raptor_chunk_field_map("doc_1", "tenant_1", "kb_1")
+
+                    # Should have called thread_pool_exec twice (primary + fallback)
+                    assert mock_thread.call_count == 2
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_handles_fallback_search_exception(self):
+        """Test that exception in fallback search is handled gracefully."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        mock_doc_store.get_fields.return_value = {}
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            call_count = 0
+            async def mock_exec(*args, **kwargs):
+                nonlocal call_count
+                call_count += 1
+                if call_count == 1:
+                    return {}  # Primary returns empty
+                else:
+                    raise Exception("Fallback search failed")  # Fallback will raise exception
+
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                mock_thread.side_effect = mock_exec
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = set()  # Primary has no RAPTOR chunks
+
+                    # Fallback will raise exception, but it should be caught
+                    result = await get_raptor_chunk_field_map("doc_1", "tenant_1", "kb_1")
+
+                    # Should return primary result (empty)
+                    assert result == {}
+        finally:
+            settings.docStoreConn = original_retriever
+
+
+class TestDeleteRaptorChunks:
+    """Tests for delete_raptor_chunks function."""
+
+    @pytest.mark.asyncio
+    async def test_deletes_all_chunks_when_keep_method_is_none(self):
+        """Test that all RAPTOR chunks are deleted when keep_method is None."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                mock_thread.return_value = 0
+
+                _ = await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method=None)
+
+                mock_thread.assert_called_once()
+                # Verify delete was called with correct condition
+                call_args = mock_thread.call_args
+                assert call_args[0][0] == settings.docStoreConn.delete
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_returns_0_when_no_stale_chunks(self):
+        """Test that 0 is returned when no stale chunks to delete."""
+        with patch("rag.svr.task_executor_refactor.raptor_utils.get_raptor_chunk_field_map") as mock_get_map:
+            mock_get_map.return_value = {}
+
+            with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                mock_collect.return_value = set()  # No stale chunks
+
+                result = await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method="raptor")
+
+                assert result == 0
+                mock_collect.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_deletes_stale_chunks_when_keep_method_specified(self):
+        """Test that stale chunks are deleted when keep_method is specified."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.get_raptor_chunk_field_map") as mock_get_map:
+                mock_get_map.return_value = {
+                    "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}},
+                    "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}
+                }
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = {"chunk_1"}  # Only chunk_1 is stale (psi, not raptor)
+
+                    with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                        mock_thread.return_value = 0
+
+                        _ = await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method="raptor")
+
+                        # Should have called delete for stale chunks
+                        mock_thread.assert_called_once()
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_logs_info_when_removing_stale_chunks(self):
+        """Test that info is logged when removing stale chunks."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.get_raptor_chunk_field_map") as mock_get_map:
+                mock_get_map.return_value = {
+                    "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}
+                }
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = {"chunk_1"}
+
+                    with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                        mock_thread.return_value = 0
+
+                        with patch("rag.svr.task_executor_refactor.raptor_utils.logging.info") as mock_log:
+                            await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method="raptor")
+
+                            # Should have logged the removal
+                            mock_log.assert_called()
+        finally:
+            settings.docStoreConn = original_retriever

test/unit_test/rag/svr/task_executor_refactor/test_write_operation_interceptor.py

@@ -0,0 +1,228 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Unit tests for WriteOperationInterceptor module.
+"""
+
+import pytest
+from rag.svr.task_executor_refactor.write_operation_interceptor import (
+    WriteOperationInterceptor,
+    ALLOWED_METHOD_NAMES,
+)
+
+
+def _create_valid_recorded_values():
+    """Helper to create valid recorded_values dict."""
+    return {method: [] for method in ALLOWED_METHOD_NAMES}
+
+
+@pytest.fixture
+def valid_recorded_values():
+    """Provide a valid recorded_values dict for testing."""
+    return _create_valid_recorded_values()
+
+
+class TestAllowedMethodNames:
+    """Tests for ALLOWED_METHOD_NAMES constant."""
+
+    def test_allowed_method_names_count(self):
+        """Test that ALLOWED_METHOD_NAMES contains exactly 8 methods."""
+        assert len(ALLOWED_METHOD_NAMES) == 10
+
+    def test_allowed_method_names_contains_expected_methods(self):
+        """Test that ALLOWED_METHOD_NAMES contains all expected methods."""
+        expected_methods = {
+            "KnowledgebaseService.update_by_id",
+            "TaskService.update_chunk_ids",
+            "DocumentService.increment_chunk_num",
+            "DocMetadataService.update_document_metadata",
+            "PipelineOperationLogService.record_pipeline_operation",
+            "handle_save_to_memory_task",
+            "PipelineOperationLogService.create",
+            "delete_raptor_chunks",
+            "docStoreConn.insert",
+            "docStoreConn.delete"
+        }
+        assert ALLOWED_METHOD_NAMES == expected_methods
+
+
+class TestWriteOperationInterceptorInit:
+    """Tests for WriteOperationInterceptor.__init__."""
+
+    def test_init_with_valid_empty_values(self, valid_recorded_values):
+        """Test initialization with valid but empty values for all methods."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor is not None
+
+    def test_init_with_valid_values(self, valid_recorded_values):
+        """Test initialization with valid recorded values."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 0]
+        valid_recorded_values["handle_save_to_memory_task"] = [None]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor is not None
+
+    def test_init_with_extra_keys_ignored(self, valid_recorded_values):
+        """Test that extra keys in recorded_values are ignored."""
+        valid_recorded_values["invalid_method_name"] = [1, 2, 3]
+        # Should not raise an error, extra keys are simply ignored
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor is not None
+        # The extra key should not be accessible
+        assert "invalid_method_name" not in interceptor._recorded_values
+
+
+class TestWriteOperationInterceptorIntercept:
+    """Tests for WriteOperationInterceptor.intercept."""
+
+    def test_intercept_returns_first_value(self, valid_recorded_values):
+        """Test that intercept returns the first value in the list."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 0, 2]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert result == 1
+
+    def test_intercept_returns_subsequent_values(self, valid_recorded_values):
+        """Test that intercept returns subsequent values on each call."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 0, 2]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.intercept("KnowledgebaseService.update_by_id") == 1
+        assert interceptor.intercept("KnowledgebaseService.update_by_id") == 0
+        assert interceptor.intercept("KnowledgebaseService.update_by_id") == 2
+
+    def test_intercept_invalid_method_raises_value_error(self, valid_recorded_values):
+        """Test that intercepting an invalid method raises ValueError."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        with pytest.raises(ValueError, match="Cannot intercept method"):
+            interceptor.intercept("invalid_method_name")
+
+    def test_intercept_empty_list_raises_index_error(self, valid_recorded_values):
+        """Test that intercepting when list is empty raises IndexError."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        with pytest.raises(IndexError, match="No more recorded values"):
+            interceptor.intercept("KnowledgebaseService.update_by_id")
+
+    def test_intercept_pops_value(self, valid_recorded_values):
+        """Test that intercept pops the value from the internal list."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [42]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        # Check internal state, not the original input list (which is copied)
+        assert len(interceptor._recorded_values["KnowledgebaseService.update_by_id"]) == 0
+
+    def test_intercept_with_none_value(self, valid_recorded_values):
+        """Test that intercept can return None values."""
+        valid_recorded_values["handle_save_to_memory_task"] = [None]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("handle_save_to_memory_task")
+        assert result is None
+
+    def test_intercept_with_default_value_when_empty(self, valid_recorded_values):
+        """Test that intercept returns default_value when list is empty."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=42)
+        assert result == 42
+
+    def test_intercept_with_default_value_none_when_empty(self, valid_recorded_values):
+        """Test that intercept returns None when default_value is None and list is empty."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        # When default_value is None, it should return None (not raise IndexError)
+        # because None is a valid default value (different from _NO_DEFAULT sentinel)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=None)
+        assert result is None
+
+    def test_intercept_default_value_does_not_affect_existing_values(self, valid_recorded_values):
+        """Test that default_value is only used when list is empty."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [100]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        # Should return the recorded value, not the default_value
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=999)
+        assert result == 100
+
+    @pytest.mark.parametrize("default_value", [
+        "default_string",
+        {"status": "success", "data": [1, 2, 3]},
+        [1, 2, 3, 4, 5],
+        (1, "two", 3.0),
+        True,
+        False,
+        0,
+        "",
+        [],
+        {},
+    ])
+    def test_intercept_with_various_default_values(self, valid_recorded_values, default_value):
+        """Test that intercept returns various default_value types when list is empty."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=default_value)
+        assert result == default_value
+
+    def test_intercept_with_complex_values(self, valid_recorded_values):
+        """Test that intercept can return complex values like dicts and tuples."""
+        complex_value = {"key": "value", "nested": [1, 2, 3]}
+        valid_recorded_values["DocMetadataService.update_document_metadata"] = [complex_value]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("DocMetadataService.update_document_metadata")
+        assert result == complex_value
+
+class TestWriteOperationInterceptorRemainingCount:
+    """Tests for WriteOperationInterceptor.remaining_count."""
+
+    def test_remaining_count_with_values(self, valid_recorded_values):
+        """Test remaining_count returns correct count."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 2, 3]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 3
+
+    def test_remaining_count_empty_list(self, valid_recorded_values):
+        """Test remaining_count returns 0 for empty list."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 0
+
+        with pytest.raises(IndexError):
+            interceptor.intercept("KnowledgebaseService.update_by_id")
+
+    def test_remaining_count_after_intercept(self, valid_recorded_values):
+        """Test remaining_count decreases after intercept calls."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 2, 3]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 3
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 2
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 1
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 0
+
+    def test_remaining_count_invalid_method(self, valid_recorded_values):
+        """Test remaining_count returns 0 for invalid method names."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("invalid_method") == 0
+
+
+class TestWriteOperationInterceptorRepr:
+    """Tests for WriteOperationInterceptor.__repr__."""
+
+    def test_repr_contains_class_name(self, valid_recorded_values):
+        """Test that repr contains the class name."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        repr_str = repr(interceptor)
+        assert "WriteOperationInterceptor" in repr_str
+
+    def test_repr_contains_total_recorded(self, valid_recorded_values):
+        """Test that repr contains total_recorded."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        repr_str = repr(interceptor)
+        assert "total_recorded=" in repr_str

test/unit_test/rag/utils/test_raptor_utils.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -12,395 +12,441 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-#
 
 """
-Unit tests for Raptor utility functions.
+Unit tests for rag/utils/raptor_utils.py module.
 """
 
-import logging
-
 import pytest
 from rag.utils.raptor_utils import (
-    CSV_EXTENSIONS,
-    EXCEL_EXTENSIONS,
-    STRUCTURED_EXTENSIONS,
-    collect_raptor_chunk_ids,
-    collect_raptor_methods,
-    get_raptor_clustering_method,
+    RAPTOR_TREE_BUILDER,
+    PSI_TREE_BUILDER,
+    GMM_CLUSTERING_METHOD,
+    AHC_CLUSTERING_METHOD,
     get_raptor_tree_builder,
-    get_skip_reason,
+    get_raptor_clustering_method,
+    _as_extra_dict,
+    _has_raptor_marker,
+    _raptor_methods_from_fields,
+    collect_raptor_methods,
+    collect_raptor_chunk_ids,
+    make_raptor_summary_chunk_id,
     is_structured_file_type,
     is_tabular_pdf,
-    make_raptor_summary_chunk_id,
     should_skip_raptor,
+    get_skip_reason,
 )
 
 
+class TestGetRaptorTreeBuilder:
+    """Tests for get_raptor_tree_builder function."""
+
+    def test_returns_default_raptor_tree_builder(self):
+        """Test that default tree builder is 'raptor'."""
+        result = get_raptor_tree_builder(None)
+        assert result == RAPTOR_TREE_BUILDER
+
+    def test_returns_default_with_empty_config(self):
+        """Test that empty config returns default."""
+        result = get_raptor_tree_builder({})
+        assert result == RAPTOR_TREE_BUILDER
+
+    def test_returns_configured_tree_builder(self):
+        """Test that configured tree builder is returned."""
+        config = {"tree_builder": PSI_TREE_BUILDER}
+        result = get_raptor_tree_builder(config)
+        assert result == PSI_TREE_BUILDER
+
+    def test_returns_ext_tree_builder(self):
+        """Test that ext.tree_builder takes precedence."""
+        config = {"tree_builder": "old", "ext": {"tree_builder": PSI_TREE_BUILDER}}
+        result = get_raptor_tree_builder(config)
+        assert result == PSI_TREE_BUILDER
+
+    def test_raises_error_for_unsupported_tree_builder(self):
+        """Test that unsupported tree builder raises ValueError."""
+        config = {"tree_builder": "unknown"}
+        with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
+            get_raptor_tree_builder(config)
+
+
+class TestGetRaptorClusteringMethod:
+    """Tests for get_raptor_clustering_method function."""
+
+    def test_returns_default_gmm(self):
+        """Test that default clustering method is 'gmm'."""
+        result = get_raptor_clustering_method(None)
+        assert result == GMM_CLUSTERING_METHOD
+
+    def test_returns_configured_clustering_method(self):
+        """Test that configured clustering method is returned."""
+        config = {"clustering_method": AHC_CLUSTERING_METHOD}
+        result = get_raptor_clustering_method(config)
+        assert result == AHC_CLUSTERING_METHOD
+
+    def test_returns_ext_clustering_method(self):
+        """Test that ext.clustering_method takes precedence."""
+        config = {"clustering_method": "old", "ext": {"clustering_method": AHC_CLUSTERING_METHOD}}
+        result = get_raptor_clustering_method(config)
+        assert result == AHC_CLUSTERING_METHOD
+
+    def test_raises_error_for_unsupported_clustering_method(self):
+        """Test that unsupported clustering method raises ValueError."""
+        config = {"clustering_method": "unknown"}
+        with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
+            get_raptor_clustering_method(config)
+
+
+class TestAsExtraDict:
+    """Tests for _as_extra_dict function."""
+
+    def test_returns_dict_as_is(self):
+        """Test that dict input is returned as-is."""
+        input_dict = {"key": "value"}
+        result = _as_extra_dict(input_dict)
+        assert result == input_dict
+
+    def test_returns_empty_dict_for_none(self):
+        """Test that None input returns empty dict."""
+        result = _as_extra_dict(None)
+        assert result == {}
+
+    def test_returns_empty_dict_for_empty_string(self):
+        """Test that empty string input returns empty dict."""
+        result = _as_extra_dict("")
+        assert result == {}
+
+    def test_parses_valid_json_string(self):
+        """Test that valid JSON string is parsed correctly."""
+        input_str = '{"key": "value"}'
+        result = _as_extra_dict(input_str)
+        assert result == {"key": "value"}
+
+    def test_returns_empty_dict_for_non_dict_json(self):
+        """Test that non-dict JSON returns empty dict."""
+        input_str = '[1, 2, 3]'
+        result = _as_extra_dict(input_str)
+        assert result == {}
+
+    def test_parses_python_dict_literal(self):
+        """Test that Python dict literal is parsed."""
+        input_str = "{'key': 'value'}"
+        result = _as_extra_dict(input_str)
+        assert result == {"key": "value"}
+
+    def test_returns_empty_dict_for_malformed_string(self):
+        """Test that malformed string returns empty dict."""
+        input_str = "{invalid json}"
+        result = _as_extra_dict(input_str)
+        assert result == {}
+
+
+class TestHasRaptorMarker:
+    """Tests for _has_raptor_marker function."""
+
+    def test_returns_true_for_raptor_string(self):
+        """Test that 'raptor' string returns True."""
+        assert _has_raptor_marker("raptor") is True
+
+    def test_returns_true_for_raptor_in_list(self):
+        """Test that 'raptor' in list returns True."""
+        assert _has_raptor_marker(["raptor", "other"]) is True
+
+    def test_returns_false_for_other_string(self):
+        """Test that other string returns False."""
+        assert _has_raptor_marker("other") is False
+
+    def test_returns_false_for_empty_list(self):
+        """Test that empty list returns False."""
+        assert _has_raptor_marker([]) is False
+
+    def test_returns_false_for_list_without_raptor(self):
+        """Test that list without 'raptor' returns False."""
+        assert _has_raptor_marker(["psi", "other"]) is False
+
+
+class TestRaptorMethodsFromFields:
+    """Tests for _raptor_methods_from_fields function."""
+
+    def test_returns_default_raptor_method(self):
+        """Test that default method is 'raptor'."""
+        result = _raptor_methods_from_fields({})
+        assert result == {RAPTOR_TREE_BUILDER}
+
+    def test_returns_method_from_extra_dict(self):
+        """Test that method is extracted from extra dict."""
+        fields = {"extra": {"raptor_method": PSI_TREE_BUILDER}}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {PSI_TREE_BUILDER}
+
+    def test_returns_method_from_extra_field(self):
+        """Test that method is extracted from extra field directly."""
+        fields = {"extra": "{'raptor_method': 'psi'}"}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {PSI_TREE_BUILDER}
+
+    def test_handles_list_method(self):
+        """Test that list method is converted to set."""
+        fields = {"extra": {"raptor_method": ["raptor", "psi"]}}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {RAPTOR_TREE_BUILDER, PSI_TREE_BUILDER}
+
+    def test_handles_empty_method(self):
+        """Test that empty method returns default."""
+        fields = {"extra": {"raptor_method": ""}}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {RAPTOR_TREE_BUILDER}
+
+
+class TestCollectRaptorMethods:
+    """Tests for collect_raptor_methods function."""
+
+    def test_returns_empty_set_for_empty_map(self):
+        """Test that empty field map returns empty set."""
+        result = collect_raptor_methods({})
+        assert result == set()
+
+    def test_collects_methods_from_raptor_chunks(self):
+        """Test that methods are collected from RAPTOR chunks."""
+        field_map = {
+            "chunk_1": {
+                "raptor_kwd": "raptor",
+                "extra": {"raptor_method": PSI_TREE_BUILDER}
+            }
+        }
+        result = collect_raptor_methods(field_map)
+        assert result == {PSI_TREE_BUILDER}
+
+    def test_skips_non_raptor_chunks(self):
+        """Test that non-RAPTOR chunks are skipped."""
+        field_map = {
+            "chunk_1": {
+                "raptor_kwd": "other",
+                "extra": {"raptor_method": PSI_TREE_BUILDER}
+            }
+        }
+        result = collect_raptor_methods(field_map)
+        assert result == set()
+
+    def test_collects_multiple_methods(self):
+        """Test that multiple methods are collected."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
+            "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}
+        }
+        result = collect_raptor_methods(field_map)
+        assert result == {RAPTOR_TREE_BUILDER, PSI_TREE_BUILDER}
+
+
+class TestCollectRaptorChunkIds:
+    """Tests for collect_raptor_chunk_ids function."""
+
+    def test_returns_empty_set_for_empty_map(self):
+        """Test that empty field map returns empty set."""
+        result = collect_raptor_chunk_ids({})
+        assert result == set()
+
+    def test_collects_ids_of_raptor_chunks(self):
+        """Test that IDs of RAPTOR chunks are collected."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor"},
+            "chunk_2": {"raptor_kwd": "raptor"}
+        }
+        result = collect_raptor_chunk_ids(field_map)
+        assert result == {"chunk_1", "chunk_2"}
+
+    def test_excludes_specified_methods(self):
+        """Test that specified methods are excluded."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
+            "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}
+        }
+        result = collect_raptor_chunk_ids(field_map, exclude_methods={"raptor"})
+        assert result == {"chunk_2"}
+
+    def test_skips_non_raptor_chunks(self):
+        """Test that non-RAPTOR chunks are skipped."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor"},
+            "chunk_2": {"raptor_kwd": "other"}
+        }
+        result = collect_raptor_chunk_ids(field_map)
+        assert result == {"chunk_1"}
+
+
+class TestMakeRaptorSummaryChunkId:
+    """Tests for make_raptor_summary_chunk_id function."""
+
+    def test_generates_consistent_id(self):
+        """Test that same input generates same ID."""
+        id1 = make_raptor_summary_chunk_id("content", "doc_1")
+        id2 = make_raptor_summary_chunk_id("content", "doc_1")
+        assert id1 == id2
+
+    def test_generates_different_ids_for_different_content(self):
+        """Test that different content generates different ID."""
+        id1 = make_raptor_summary_chunk_id("content1", "doc_1")
+        id2 = make_raptor_summary_chunk_id("content2", "doc_1")
+        assert id1 != id2
+
+    def test_generates_different_ids_for_different_doc(self):
+        """Test that different doc_id generates different ID."""
+        id1 = make_raptor_summary_chunk_id("content", "doc_1")
+        id2 = make_raptor_summary_chunk_id("content", "doc_2")
+        assert id1 != id2
+
+    def test_returns_string(self):
+        """Test that result is a string."""
+        result = make_raptor_summary_chunk_id("content", "doc_1")
+        assert isinstance(result, str)
+
+
 class TestIsStructuredFileType:
-    """Test file type detection for structured data"""
+    """Tests for is_structured_file_type function."""
 
-    @pytest.mark.parametrize("file_type,expected", [
-        (".xlsx", True),
-        (".xls", True),
-        (".xlsm", True),
-        (".xlsb", True),
-        (".csv", True),
-        (".tsv", True),
-        ("xlsx", True),  # Without leading dot
-        ("XLSX", True),  # Uppercase
-        (".pdf", False),
-        (".docx", False),
-        (".txt", False),
-        ("", False),
-        (None, False),
-    ])
-    def test_file_type_detection(self, file_type, expected):
-        """Test detection of various file types"""
-        assert is_structured_file_type(file_type) == expected
+    def test_returns_true_for_xlsx(self):
+        """Test that .xlsx is recognized as structured."""
+        assert is_structured_file_type(".xlsx") is True
 
-    def test_excel_extensions_defined(self):
-        """Test that Excel extensions are properly defined"""
-        assert ".xlsx" in EXCEL_EXTENSIONS
-        assert ".xls" in EXCEL_EXTENSIONS
-        assert len(EXCEL_EXTENSIONS) >= 4
+    def test_returns_true_for_xls(self):
+        """Test that .xls is recognized as structured."""
+        assert is_structured_file_type(".xls") is True
 
-    def test_csv_extensions_defined(self):
-        """Test that CSV extensions are properly defined"""
-        assert ".csv" in CSV_EXTENSIONS
-        assert ".tsv" in CSV_EXTENSIONS
+    def test_returns_true_for_csv(self):
+        """Test that .csv is recognized as structured."""
+        assert is_structured_file_type(".csv") is True
 
-    def test_structured_extensions_combined(self):
-        """Test that structured extensions include both Excel and CSV"""
-        assert EXCEL_EXTENSIONS.issubset(STRUCTURED_EXTENSIONS)
-        assert CSV_EXTENSIONS.issubset(STRUCTURED_EXTENSIONS)
+    def test_returns_true_for_tsv(self):
+        """Test that .tsv is recognized as structured."""
+        assert is_structured_file_type(".tsv") is True
+
+    def test_returns_false_for_pdf(self):
+        """Test that .pdf is not structured."""
+        assert is_structured_file_type(".pdf") is False
+
+    def test_returns_false_for_txt(self):
+        """Test that .txt is not structured."""
+        assert is_structured_file_type(".txt") is False
+
+    def test_returns_false_for_none(self):
+        """Test that None is not structured."""
+        assert is_structured_file_type(None) is False
+
+    def test_returns_false_for_empty_string(self):
+        """Test that empty string is not structured."""
+        assert is_structured_file_type("") is False
+
+    def test_handles_case_insensitive(self):
+        """Test that case is handled insensitively."""
+        assert is_structured_file_type(".XLSX") is True
+        assert is_structured_file_type("xlsx") is True
+
+    def test_handles_missing_dot(self):
+        """Test that missing dot is handled."""
+        assert is_structured_file_type("xlsx") is True
 
 
-class TestIsTabularPDF:
-    """Test tabular PDF detection"""
+class TestIsTabularPdf:
+    """Tests for is_tabular_pdf function."""
 
-    def test_table_parser_detected(self):
-        """Test that table parser is detected as tabular"""
+    def test_returns_true_for_table_parser(self):
+        """Test that table parser returns True."""
         assert is_tabular_pdf("table", {}) is True
-        assert is_tabular_pdf("TABLE", {}) is True
 
-    def test_html4excel_detected(self):
-        """Test that html4excel config is detected as tabular"""
+    def test_returns_true_for_html4excel(self):
+        """Test that html4excel enabled returns True."""
         assert is_tabular_pdf("naive", {"html4excel": True}) is True
-        assert is_tabular_pdf("", {"html4excel": True}) is True
 
-    def test_non_tabular_pdf(self):
-        """Test that non-tabular PDFs are not detected"""
+    def test_returns_false_for_naive_parser(self):
+        """Test that naive parser returns False."""
         assert is_tabular_pdf("naive", {}) is False
-        assert is_tabular_pdf("naive", {"html4excel": False}) is False
+
+    def test_returns_false_for_empty_parser_id(self):
+        """Test that empty parser_id returns False."""
         assert is_tabular_pdf("", {}) is False
 
-    def test_combined_conditions(self):
-        """Test combined table parser and html4excel"""
-        assert is_tabular_pdf("table", {"html4excel": True}) is True
-        assert is_tabular_pdf("table", {"html4excel": False}) is True
+    def test_returns_false_for_html4excel_false(self):
+        """Test that html4excel=False returns False."""
+        assert is_tabular_pdf("naive", {"html4excel": False}) is False
+
+    def test_handles_case_insensitive_parser_id(self):
+        """Test that parser_id case is handled."""
+        assert is_tabular_pdf("TABLE", {}) is True
+        assert is_tabular_pdf("Table", {}) is True
 
 
 class TestShouldSkipRaptor:
-    """Test Raptor skip logic"""
+    """Tests for should_skip_raptor function."""
 
-    def test_skip_excel_files(self):
-        """Test that Excel files skip Raptor"""
-        assert should_skip_raptor(".xlsx") is True
-        assert should_skip_raptor(".xls") is True
-        assert should_skip_raptor(".xlsm") is True
+    def test_skips_for_xlsx_file(self):
+        """Test that .xlsx file skips Raptor."""
+        assert should_skip_raptor(file_type=".xlsx") is True
 
-    def test_skip_csv_files(self):
-        """Test that CSV files skip Raptor"""
-        assert should_skip_raptor(".csv") is True
-        assert should_skip_raptor(".tsv") is True
+    def test_skips_for_csv_file(self):
+        """Test that .csv file skips Raptor."""
+        assert should_skip_raptor(file_type=".csv") is True
 
-    def test_skip_tabular_pdf_with_table_parser(self):
-        """Test that tabular PDFs skip Raptor"""
-        assert should_skip_raptor(".pdf", parser_id="table") is True
-        assert should_skip_raptor("pdf", parser_id="TABLE") is True
+    def test_skips_for_tabular_pdf(self):
+        """Test that tabular PDF skips Raptor."""
+        assert should_skip_raptor(file_type=".pdf", parser_id="table") is True
 
-    def test_skip_tabular_pdf_with_html4excel(self):
-        """Test that PDFs with html4excel skip Raptor"""
-        assert should_skip_raptor(".pdf", parser_config={"html4excel": True}) is True
+    def test_does_not_skip_for_normal_pdf(self):
+        """Test that normal PDF does not skip Raptor."""
+        assert should_skip_raptor(file_type=".pdf", parser_id="naive") is False
 
-    def test_dont_skip_regular_pdf(self):
-        """Test that regular PDFs don't skip Raptor"""
-        assert should_skip_raptor(".pdf", parser_id="naive") is False
-        assert should_skip_raptor(".pdf", parser_config={}) is False
+    def test_does_not_skip_for_txt_file(self):
+        """Test that .txt file does not skip Raptor."""
+        assert should_skip_raptor(file_type=".txt") is False
 
-    def test_dont_skip_text_files(self):
-        """Test that text files don't skip Raptor"""
-        assert should_skip_raptor(".txt") is False
-        assert should_skip_raptor(".docx") is False
-        assert should_skip_raptor(".md") is False
+    def test_respects_auto_disable_config_false(self):
+        """Test that auto_disable_for_structured_data=False disables skipping."""
+        assert should_skip_raptor(
+            file_type=".xlsx",
+            raptor_config={"auto_disable_for_structured_data": False}
+        ) is False
 
-    def test_override_with_config(self):
-        """Test that auto-disable can be overridden"""
-        raptor_config = {"auto_disable_for_structured_data": False}
-        
-        # Should not skip even for Excel files
-        assert should_skip_raptor(".xlsx", raptor_config=raptor_config) is False
-        assert should_skip_raptor(".csv", raptor_config=raptor_config) is False
-        assert should_skip_raptor(".pdf", parser_id="table", raptor_config=raptor_config) is False
+    def test_respects_auto_disable_config_true(self):
+        """Test that auto_disable_for_structured_data=True enables skipping."""
+        assert should_skip_raptor(
+            file_type=".xlsx",
+            raptor_config={"auto_disable_for_structured_data": True}
+        ) is True
 
-    def test_default_auto_disable_enabled(self):
-        """Test that auto-disable is enabled by default"""
-        # Empty raptor_config should default to auto_disable=True
-        assert should_skip_raptor(".xlsx", raptor_config={}) is True
-        assert should_skip_raptor(".xlsx", raptor_config=None) is True
+    def test_default_auto_disable_is_true(self):
+        """Test that default auto_disable is True."""
+        assert should_skip_raptor(file_type=".xlsx") is True
 
-    def test_explicit_auto_disable_enabled(self):
-        """Test explicit auto-disable enabled"""
-        raptor_config = {"auto_disable_for_structured_data": True}
-        assert should_skip_raptor(".xlsx", raptor_config=raptor_config) is True
+    def test_returns_false_for_none_file_type(self):
+        """Test that None file_type does not skip."""
+        assert should_skip_raptor(file_type=None) is False
 
 
 class TestGetSkipReason:
-    """Test skip reason generation"""
+    """Tests for get_skip_reason function."""
 
-    def test_excel_skip_reason(self):
-        """Test skip reason for Excel files"""
-        reason = get_skip_reason(".xlsx")
+    def test_returns_reason_for_structured_file(self):
+        """Test that reason is returned for structured file."""
+        reason = get_skip_reason(file_type=".xlsx")
         assert "Structured data file" in reason
         assert ".xlsx" in reason
-        assert "auto-disabled" in reason.lower()
 
-    def test_csv_skip_reason(self):
-        """Test skip reason for CSV files"""
-        reason = get_skip_reason(".csv")
-        assert "Structured data file" in reason
-        assert ".csv" in reason
-
-    def test_tabular_pdf_skip_reason(self):
-        """Test skip reason for tabular PDFs"""
-        reason = get_skip_reason(".pdf", parser_id="table")
+    def test_returns_reason_for_tabular_pdf(self):
+        """Test that reason is returned for tabular PDF."""
+        reason = get_skip_reason(file_type=".pdf", parser_id="table")
         assert "Tabular PDF" in reason
-        assert "table" in reason.lower()
-        assert "auto-disabled" in reason.lower()
+        assert "table" in reason
 
-    def test_html4excel_skip_reason(self):
-        """Test skip reason for html4excel PDFs"""
-        reason = get_skip_reason(".pdf", parser_config={"html4excel": True})
-        assert "Tabular PDF" in reason
-
-    def test_no_skip_reason_for_regular_files(self):
-        """Test that regular files have no skip reason"""
-        assert get_skip_reason(".txt") == ""
-        assert get_skip_reason(".docx") == ""
-        assert get_skip_reason(".pdf", parser_id="naive") == ""
-
-
-class TestEdgeCases:
-    """Test edge cases and error handling"""
-
-    def test_none_values(self):
-        """Test handling of None values"""
-        assert should_skip_raptor(None) is False
-        assert should_skip_raptor("") is False
-        assert get_skip_reason(None) == ""
-
-    def test_empty_strings(self):
-        """Test handling of empty strings"""
-        assert should_skip_raptor("") is False
-        assert get_skip_reason("") == ""
-
-    def test_case_insensitivity(self):
-        """Test case insensitive handling"""
-        assert is_structured_file_type("XLSX") is True
-        assert is_structured_file_type("XlSx") is True
-        assert is_tabular_pdf("TABLE", {}) is True
-        assert is_tabular_pdf("TaBlE", {}) is True
-
-    def test_with_and_without_dot(self):
-        """Test file extensions with and without leading dot"""
-        assert should_skip_raptor(".xlsx") is True
-        assert should_skip_raptor("xlsx") is True
-        assert should_skip_raptor(".CSV") is True
-        assert should_skip_raptor("csv") is True
-
-
-class TestIntegrationScenarios:
-    """Test real-world integration scenarios"""
-
-    def test_financial_excel_report(self):
-        """Test scenario: Financial quarterly Excel report"""
-        file_type = ".xlsx"
-        parser_id = "naive"
-        parser_config = {}
-        raptor_config = {"use_raptor": True}
-        
-        # Should skip Raptor
-        assert should_skip_raptor(file_type, parser_id, parser_config, raptor_config) is True
-        reason = get_skip_reason(file_type, parser_id, parser_config)
-        assert "Structured data file" in reason
-
-    def test_scientific_csv_data(self):
-        """Test scenario: Scientific experimental CSV results"""
-        file_type = ".csv"
-        
-        # Should skip Raptor
-        assert should_skip_raptor(file_type) is True
-        reason = get_skip_reason(file_type)
-        assert ".csv" in reason
-
-    def test_legal_contract_with_tables(self):
-        """Test scenario: Legal contract PDF with tables"""
-        file_type = ".pdf"
-        parser_id = "table"
-        parser_config = {}
-        
-        # Should skip Raptor
-        assert should_skip_raptor(file_type, parser_id, parser_config) is True
-        reason = get_skip_reason(file_type, parser_id, parser_config)
-        assert "Tabular PDF" in reason
-
-    def test_text_heavy_pdf_document(self):
-        """Test scenario: Text-heavy PDF document"""
-        file_type = ".pdf"
-        parser_id = "naive"
-        parser_config = {}
-        
-        # Should NOT skip Raptor
-        assert should_skip_raptor(file_type, parser_id, parser_config) is False
-        reason = get_skip_reason(file_type, parser_id, parser_config)
+    def test_returns_empty_for_normal_pdf(self):
+        """Test that empty reason is returned for normal PDF."""
+        reason = get_skip_reason(file_type=".pdf", parser_id="naive")
         assert reason == ""
 
-    def test_mixed_dataset_processing(self):
-        """Test scenario: Mixed dataset with various file types"""
-        files = [
-            (".xlsx", "naive", {}, True),  # Excel - skip
-            (".csv", "naive", {}, True),   # CSV - skip
-            (".pdf", "table", {}, True),   # Tabular PDF - skip
-            (".pdf", "naive", {}, False),  # Regular PDF - don't skip
-            (".docx", "naive", {}, False), # Word doc - don't skip
-            (".txt", "naive", {}, False),  # Text file - don't skip
-        ]
-        
-        for file_type, parser_id, parser_config, expected_skip in files:
-            result = should_skip_raptor(file_type, parser_id, parser_config)
-            assert result == expected_skip, f"Failed for {file_type}"
+    def test_returns_empty_for_txt_file(self):
+        """Test that empty reason is returned for .txt file."""
+        reason = get_skip_reason(file_type=".txt")
+        assert reason == ""
 
-    def test_override_for_special_excel(self):
-        """Test scenario: Override auto-disable for special Excel processing"""
-        file_type = ".xlsx"
-        raptor_config = {"auto_disable_for_structured_data": False}
-        
-        # Should NOT skip when explicitly disabled
-        assert should_skip_raptor(file_type, raptor_config=raptor_config) is False
-
-
-class TestRaptorTreeBuilderConfig:
-    """Test RAPTOR tree builder config resolution"""
-
-    def test_defaults_to_original_raptor_builder(self):
-        assert get_raptor_tree_builder({}) == "raptor"
-        assert get_raptor_tree_builder(None) == "raptor"
-
-    def test_reads_top_level_tree_builder(self):
-        assert get_raptor_tree_builder({"tree_builder": "psi"}) == "psi"
-
-    def test_reads_legacy_ext_tree_builder(self):
-        assert get_raptor_tree_builder({"ext": {"tree_builder": "psi"}}) == "psi"
-
-    def test_ext_tree_builder_overrides_stale_top_level_value(self):
-        assert get_raptor_tree_builder({"tree_builder": "psi", "ext": {"tree_builder": "raptor"}}) == "raptor"
-
-    def test_rejects_unknown_tree_builder(self):
-        with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
-            get_raptor_tree_builder({"tree_builder": "ahc"})
-
-
-class TestRaptorClusteringMethodConfig:
-    """Test RAPTOR clustering method config resolution"""
-
-    def test_defaults_to_gmm(self):
-        assert get_raptor_clustering_method({}) == "gmm"
-        assert get_raptor_clustering_method(None) == "gmm"
-
-    def test_reads_top_level_clustering_method(self):
-        assert get_raptor_clustering_method({"clustering_method": "gmm"}) == "gmm"
-        assert get_raptor_clustering_method({"clustering_method": "ahc"}) == "ahc"
-
-    def test_reads_legacy_ext_clustering_method(self):
-        assert get_raptor_clustering_method({"ext": {"clustering_method": "ahc"}}) == "ahc"
-
-    def test_ext_clustering_method_overrides_stale_top_level_value(self):
-        assert get_raptor_clustering_method({"clustering_method": "gmm", "ext": {"clustering_method": "ahc"}}) == "ahc"
-
-    def test_rejects_unknown_clustering_method(self):
-        with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
-            get_raptor_clustering_method({"clustering_method": "unknown"})
-
-
-class TestRaptorMethodCollection:
-    """Test RAPTOR summary method extraction from doc-store fields"""
-
-    def test_legacy_summary_without_method_is_original_raptor(self):
-        field_map = {"chunk_1": {"raptor_kwd": "raptor"}}
-
-        assert collect_raptor_methods(field_map) == {"raptor"}
-        assert collect_raptor_chunk_ids(field_map) == {"chunk_1"}
-
-    def test_extra_method_is_preserved(self):
-        field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}}
-
-        assert collect_raptor_methods(field_map) == {"psi"}
-        assert collect_raptor_chunk_ids(field_map) == {"chunk_1"}
-
-    def test_extra_field_supports_oceanbase_legacy_rows(self):
-        field_map = {
-            "chunk_1": {
-                "extra": {
-                    "raptor_kwd": "raptor",
-                    "raptor_method": "psi",
-                }
-            },
-            "chunk_2": {
-                "extra": "{\"raptor_kwd\": \"raptor\"}",
-            },
-            "chunk_3": {
-                "extra": {"raptor_kwd": ""},
-            },
-        }
-
-        assert collect_raptor_methods(field_map) == {"psi", "raptor"}
-        assert collect_raptor_chunk_ids(field_map) == {"chunk_1", "chunk_2"}
-
-    def test_non_raptor_rows_are_ignored(self):
-        field_map = {
-            "chunk_1": {"raptor_kwd": ""},
-            "chunk_2": {"extra": {"raptor_kwd": "graph"}},
-            "chunk_3": {},
-        }
-
-        assert collect_raptor_methods(field_map) == set()
-        assert collect_raptor_chunk_ids(field_map) == set()
-
-    def test_malformed_extra_payload_is_logged_and_ignored(self, caplog):
-        field_map = {"chunk_1": {"extra": "{bad json"}}
-
-        with caplog.at_level(logging.WARNING):
-            assert collect_raptor_methods(field_map) == set()
-            assert collect_raptor_chunk_ids(field_map) == set()
-
-        assert "Ignoring malformed RAPTOR extra payload" in caplog.text
-
-    def test_chunk_id_collection_can_preserve_current_method(self):
-        field_map = {
-            "legacy": {"raptor_kwd": "raptor"},
-            "old": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
-            "current": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}},
-        }
-
-        assert collect_raptor_chunk_ids(field_map, exclude_methods={"psi"}) == {"legacy", "old"}
-        assert collect_raptor_chunk_ids(field_map, exclude_methods={"raptor"}) == {"current"}
-
-    def test_summary_chunk_ids_include_real_document_id(self):
-        content = "same generated summary"
-
-        assert make_raptor_summary_chunk_id(content, "doc-a") != make_raptor_summary_chunk_id(content, "doc-b")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
+    def test_returns_empty_for_none_file_type(self):
+        """Test that empty reason is returned for None file_type."""
+        reason = get_skip_reason(file_type=None)
+        assert reason == ""