陈辅元
8466ed2fbe
- Added logic to set `table_title` from `caption` if it is under 160 characters and not already set. - Updated `_blocks_from_content_list` to assign `table_title` based on `caption` length. - Introduced new regex patterns for better detection of TOC entries and noise. - Enhanced heading detection to differentiate between numbered instructions and actual headings. - Added tests to verify that table captions are correctly assigned as titles and that numbered instructions are treated as body text.
182 lines
7.0 KiB
Python
182 lines
7.0 KiB
Python
"""Tests for MinerU content_list → Block mapping."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import subprocess
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from rag_cut.models import BlockType
|
|
from rag_cut.parsers.mineru_adapter import (
|
|
_blocks_from_content_list,
|
|
_find_content_list,
|
|
_run_mineru,
|
|
)
|
|
|
|
|
|
class MineruAdapterProcessTest(unittest.TestCase):
|
|
def test_default_timeout_allows_large_manual_to_finish(self) -> None:
|
|
with patch.dict("os.environ", {}, clear=True):
|
|
from rag_cut.parsers.mineru_adapter import _mineru_timeout_sec
|
|
|
|
self.assertEqual(_mineru_timeout_sec(), 540.0)
|
|
|
|
def test_configured_api_url_is_passed_to_mineru(self) -> None:
|
|
process = Mock()
|
|
process.communicate.return_value = ("", "")
|
|
process.returncode = 0
|
|
|
|
with (
|
|
tempfile.TemporaryDirectory() as tmp,
|
|
patch("rag_cut.parsers.mineru_adapter._mineru_command", return_value="mineru"),
|
|
patch.dict("os.environ", {"RAG_CUT_MINERU_API_URL": "http://127.0.0.1:30000"}),
|
|
patch("rag_cut.parsers.mineru_adapter.subprocess.Popen", return_value=process) as popen,
|
|
):
|
|
ok = _run_mineru(Path(tmp) / "manual.pdf", Path(tmp) / "output")
|
|
|
|
self.assertTrue(ok)
|
|
args = popen.call_args.args[0]
|
|
self.assertEqual(args[-2:], ["--api-url", "http://127.0.0.1:30000"])
|
|
|
|
def test_timeout_terminates_windows_process_tree(self) -> None:
|
|
process = Mock()
|
|
process.pid = 4321
|
|
process.communicate.side_effect = subprocess.TimeoutExpired("mineru", 30)
|
|
process.poll.return_value = None
|
|
|
|
with (
|
|
tempfile.TemporaryDirectory() as tmp,
|
|
patch("rag_cut.parsers.mineru_adapter._mineru_command", return_value="mineru"),
|
|
patch("rag_cut.parsers.mineru_adapter._mineru_timeout_sec", return_value=30),
|
|
patch("rag_cut.parsers.mineru_adapter.os.name", "nt"),
|
|
patch("rag_cut.parsers.mineru_adapter.subprocess.Popen", return_value=process),
|
|
patch("rag_cut.parsers.mineru_adapter.subprocess.run") as run,
|
|
):
|
|
ok = _run_mineru(Path(tmp) / "manual.pdf", Path(tmp) / "output")
|
|
|
|
self.assertFalse(ok)
|
|
run.assert_called_once_with(
|
|
["taskkill", "/PID", "4321", "/T", "/F"],
|
|
check=False,
|
|
capture_output=True,
|
|
timeout=10,
|
|
)
|
|
|
|
|
|
class MineruAdapterMappingTest(unittest.TestCase):
|
|
def test_chart_becomes_image_block(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
assets = Path(tmp)
|
|
img = assets / "chart.jpg"
|
|
img.write_bytes(b"fake")
|
|
items = [
|
|
{
|
|
"type": "text",
|
|
"text": "第二十節 TradingView",
|
|
"text_level": 2,
|
|
"page_idx": 30,
|
|
"bbox": [127, 86, 363, 104],
|
|
},
|
|
{
|
|
"type": "chart",
|
|
"img_path": "chart.jpg",
|
|
"page_idx": 30,
|
|
"bbox": [129, 171, 878, 428],
|
|
"image_caption": [],
|
|
},
|
|
]
|
|
with patch("rag_cut.parsers.mineru_adapter._mineru_ocr_enabled", return_value=False):
|
|
blocks = _blocks_from_content_list(items, assets / "content_list.json", assets)
|
|
|
|
types = [b.type for b in blocks]
|
|
self.assertEqual(types, [BlockType.HEADING, BlockType.IMAGE])
|
|
self.assertEqual(blocks[1].image_id, "chart.jpg")
|
|
self.assertEqual(blocks[1].meta.get("mineru_type"), "chart")
|
|
self.assertEqual(blocks[1].meta.get("page"), 31)
|
|
|
|
def test_table_keeps_markdown_and_screenshot(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
assets = Path(tmp)
|
|
img = assets / "table.png"
|
|
img.write_bytes(b"png")
|
|
items = [
|
|
{
|
|
"type": "table",
|
|
"img_path": "table.png",
|
|
"page_idx": 1,
|
|
"bbox": [100, 100, 400, 300],
|
|
"table_body": "<table><tr><td>A</td><td>B</td></tr></table>",
|
|
}
|
|
]
|
|
with patch("rag_cut.parsers.mineru_adapter._mineru_ocr_enabled", return_value=False):
|
|
blocks = _blocks_from_content_list(items, assets / "content_list.json", assets)
|
|
|
|
self.assertEqual(len(blocks), 1)
|
|
self.assertEqual(blocks[0].type, BlockType.TABLE)
|
|
self.assertIn("<table>", blocks[0].markdown or "")
|
|
self.assertEqual(blocks[0].image_id, "table.png")
|
|
self.assertTrue(blocks[0].image_path)
|
|
|
|
def test_page_footnote_kept_as_paragraph(self) -> None:
|
|
items = [
|
|
{
|
|
"type": "page_footnote",
|
|
"text": "<sup>1</sup> See Market Master.",
|
|
"page_idx": 8,
|
|
"bbox": [85, 889, 912, 916],
|
|
}
|
|
]
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
blocks = _blocks_from_content_list(items, Path(tmp) / "content_list.json", Path(tmp))
|
|
self.assertEqual(len(blocks), 1)
|
|
self.assertEqual(blocks[0].type, BlockType.PARAGRAPH)
|
|
self.assertTrue(blocks[0].meta.get("is_footnote"))
|
|
|
|
def test_table_caption_becomes_table_title(self) -> None:
|
|
items = [
|
|
{
|
|
"type": "table",
|
|
"table_body": "| Field | Description |\n| --- | --- |\n| id | Identifier |",
|
|
"table_caption": ["第一节 综合资讯栏"],
|
|
"page_idx": 7,
|
|
}
|
|
]
|
|
|
|
blocks = _blocks_from_content_list(items, Path("content_list.json"), Path("assets"))
|
|
|
|
self.assertEqual(len(blocks), 1)
|
|
self.assertEqual(blocks[0].meta.get("table_title"), "第一节 综合资讯栏")
|
|
|
|
def test_list_items_become_paragraph(self) -> None:
|
|
items = [
|
|
{
|
|
"type": "list",
|
|
"list_items": ["第一點說明", {"text": "第二點說明"}],
|
|
"page_idx": 2,
|
|
"bbox": [100, 200, 400, 260],
|
|
}
|
|
]
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
blocks = _blocks_from_content_list(items, Path(tmp) / "content_list.json", Path(tmp))
|
|
self.assertEqual(len(blocks), 1)
|
|
self.assertEqual(blocks[0].type, BlockType.PARAGRAPH)
|
|
self.assertIn("第一點說明", blocks[0].text or "")
|
|
self.assertIn("第二點說明", blocks[0].text or "")
|
|
|
|
def test_find_content_list_prefers_non_v2(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp) / "doc" / "auto"
|
|
root.mkdir(parents=True)
|
|
v1 = root / "doc_content_list.json"
|
|
v2 = root / "doc_content_list_v2.json"
|
|
v1.write_text("[]", encoding="utf-8")
|
|
v2.write_text("[]", encoding="utf-8")
|
|
found = _find_content_list(Path(tmp))
|
|
self.assertEqual(found, v1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|