mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-08 00:18:12 +08:00
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>
This commit is contained in:
Jack
@@ -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 == ""
|
||||
|
||||
Reference in New Issue
Block a user