532 lines
21 KiB
Python
532 lines
21 KiB
Python
"""
|
||
布局后处理单元测试（纯逻辑，不依赖引擎和服务）。
|
||
|
||
测试 `src.engine.layout` 中的纯函数：
|
||
- sort_and_group：阅读顺序排序 + 行归组
|
||
- group_paragraphs：段落分组
|
||
- build_markdown_text：Markdown 格式化
|
||
- 辅助函数
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import pytest
|
||
|
||
from src.engine.layout import (
|
||
_block_center_x,
|
||
_block_center_y,
|
||
_block_height,
|
||
_detect_columns,
|
||
_detect_heading_level,
|
||
_detect_line_type,
|
||
_sort_and_group_single,
|
||
build_markdown_text,
|
||
group_paragraphs,
|
||
sort_and_group,
|
||
)
|
||
from src.schemas.ocr import BBox, TextBlock
|
||
|
||
|
||
# ── 工厂函数 ──────────────────────────────────────────────
|
||
|
||
def _make_block(
|
||
text: str,
|
||
x1: float = 0,
|
||
y1: float = 0,
|
||
x2: float = 100,
|
||
y2: float = 30,
|
||
confidence: float = 0.95,
|
||
) -> TextBlock:
|
||
"""创建一个测试用的 TextBlock。"""
|
||
return TextBlock(
|
||
text=text,
|
||
confidence=confidence,
|
||
bbox=BBox(points=[[x1, y1], [x2, y1], [x2, y2], [x1, y2]]),
|
||
)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 辅助函数
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestHelpers:
|
||
"""辅助函数测试。"""
|
||
|
||
def test_block_height(self) -> None:
|
||
block = _make_block("测试", y1=10, y2=50)
|
||
assert _block_height(block) == 40.0
|
||
|
||
def test_block_height_empty(self) -> None:
|
||
block = _make_block("测试", x1=0, y1=0, x2=0, y2=0)
|
||
assert _block_height(block) == 0.0
|
||
|
||
def test_block_center_y(self) -> None:
|
||
block = _make_block("测试", y1=10, y2=50)
|
||
assert _block_center_y(block) == 30.0
|
||
|
||
def test_block_center_x(self) -> None:
|
||
block = _make_block("测试", x1=20, x2=100)
|
||
assert _block_center_x(block) == 60.0
|
||
|
||
def test_detect_heading_chinese(self) -> None:
|
||
"""中文短行 → heading。"""
|
||
assert _detect_line_type("第一章 绪论") == "heading"
|
||
|
||
def test_detect_heading_short_no_punct(self) -> None:
|
||
"""短行不以句号结尾 → heading。"""
|
||
assert _detect_line_type("摘要") == "heading"
|
||
|
||
def test_detect_not_heading_with_punct(self) -> None:
|
||
"""以句号结尾的短行 → text。"""
|
||
assert _detect_line_type("这是一个句子。") == "text"
|
||
|
||
def test_detect_long_text(self) -> None:
|
||
"""长行（>= 40 字符）→ text。"""
|
||
long = "一二三四五六七八九十" * 4
|
||
assert _detect_line_type(long) == "text"
|
||
|
||
def test_detect_list_item_dash(self) -> None:
|
||
assert _detect_line_type("- 列表项") == "list_item"
|
||
|
||
def test_detect_list_item_bullet(self) -> None:
|
||
assert _detect_line_type("• 项目") == "list_item"
|
||
|
||
def test_detect_numbered(self) -> None:
|
||
assert _detect_line_type("1. 首先") == "numbered"
|
||
assert _detect_line_type("2) 其次") == "numbered"
|
||
|
||
def test_detect_separator(self) -> None:
|
||
assert _detect_line_type("---") == "separator"
|
||
assert _detect_line_type("===") == "separator"
|
||
assert _detect_line_type("***") == "separator"
|
||
|
||
def test_detect_code(self) -> None:
|
||
assert _detect_line_type(" def foo():") == "code"
|
||
assert _detect_line_type("\t\tindented") == "code"
|
||
|
||
def test_detect_blockquote_as_text(self) -> None:
|
||
"""> 前缀在 OCR 语境中通常是排版伪影，不作为代码块。"""
|
||
assert _detect_line_type("> 引用") == "text"
|
||
|
||
def test_detect_empty_text(self) -> None:
|
||
assert _detect_line_type("") == "text"
|
||
assert _detect_line_type(" ") == "text"
|
||
|
||
def test_detect_centered_heading(self) -> None:
|
||
"""居中短行 → 即使非中文也应视为 heading。"""
|
||
assert _detect_line_type("Hello World", center_x_ratio=0.5) == "heading"
|
||
|
||
def test_detect_not_centered(self) -> None:
|
||
"""不居中短行 + 非中文锁 → 不回退。"""
|
||
# "hello world" 非中文、非大写、不居中 → text
|
||
assert _detect_line_type("hello world", center_x_ratio=0.1) == "text"
|
||
|
||
def test_detect_uppercase_heading(self) -> None:
|
||
"""全大写短行 → heading。"""
|
||
assert _detect_line_type("INTRODUCTION") == "heading"
|
||
|
||
# ── 新增：编号变体 ────────────────────────────────
|
||
|
||
def test_detect_numbered_without_space(self) -> None:
|
||
"""编号后无空格（"1.内容"）→ numbered。"""
|
||
assert _detect_line_type("1.内容") == "numbered"
|
||
assert _detect_line_type("1)内容") == "numbered"
|
||
|
||
def test_detect_numbered_parenthesized(self) -> None:
|
||
"""括号编号 → numbered。"""
|
||
assert _detect_line_type("(1) 内容") == "numbered"
|
||
assert _detect_line_type("(2)内容") == "numbered"
|
||
|
||
def test_detect_numbered_chinese_numeral(self) -> None:
|
||
"""中文数字序号 → numbered。"""
|
||
assert _detect_line_type("① 第一项") == "numbered"
|
||
assert _detect_line_type("②第二项") == "numbered"
|
||
|
||
# ── 新增：TOC 目录行 ──────────────────────────────
|
||
|
||
def test_detect_toc_line(self) -> None:
|
||
"""TOC 目录行（含连续点号）→ numbered。"""
|
||
assert _detect_line_type("第一章 .............. 1") == "numbered"
|
||
assert _detect_line_type("1.1 概述 ............ 10") == "numbered"
|
||
|
||
# ── 新增：更多 bullet 变体 ────────────────────────
|
||
|
||
def test_detect_bullet_variants(self) -> None:
|
||
"""更多 bullet 变体 → list_item。"""
|
||
assert _detect_line_type("▪ 项目") == "list_item"
|
||
assert _detect_line_type("▸ 子项") == "list_item"
|
||
assert _detect_line_type("→ 链接") == "list_item"
|
||
assert _detect_line_type("· 条目") == "list_item"
|
||
|
||
def test_detect_standalone_bullet(self) -> None:
|
||
"""独立 bullet 字符 → list_item。"""
|
||
assert _detect_line_type("•") == "list_item"
|
||
assert _detect_line_type("▪") == "list_item"
|
||
assert _detect_line_type("→") == "list_item"
|
||
|
||
# ── 新增：字体大小辅助标题检测 ────────────────────
|
||
|
||
def test_detect_larger_font_heading(self) -> None:
|
||
"""大字（rel_height > 1.3）→ heading。"""
|
||
assert _detect_line_type("这是一段长文本但字号大", rel_height=1.5) == "heading"
|
||
# 中号字（rel_height < 1.3）且不居中 → text
|
||
assert _detect_line_type("这是一段长文本字号普通", rel_height=1.1) == "text"
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 标题级别检测
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestHeadingLevel:
|
||
"""标题级别检测。"""
|
||
|
||
def test_h1_centered_large(self) -> None:
|
||
assert _detect_heading_level(rel_height=1.5, is_centered=True) == 1
|
||
|
||
def test_h1_large_alone(self) -> None:
|
||
assert _detect_heading_level(rel_height=1.8) == 1
|
||
|
||
def test_h2_large(self) -> None:
|
||
assert _detect_heading_level(rel_height=1.4) == 2
|
||
|
||
def test_h2_centered(self) -> None:
|
||
assert _detect_heading_level(is_centered=True) == 2
|
||
|
||
def test_h3_default(self) -> None:
|
||
assert _detect_heading_level() == 3
|
||
assert _detect_heading_level(rel_height=1.1) == 3
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# sort_and_group
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestSortAndGroup:
|
||
"""阅读顺序排序 + 行归组。"""
|
||
|
||
def test_empty(self) -> None:
|
||
assert sort_and_group([]) == []
|
||
|
||
def test_single_block(self) -> None:
|
||
blocks = [_make_block("你好")]
|
||
lines = sort_and_group(blocks)
|
||
assert len(lines) == 1
|
||
assert len(lines[0]) == 1
|
||
assert lines[0][0].text == "你好"
|
||
|
||
def test_top_to_bottom_sort(self) -> None:
|
||
"""按 Y 坐标从上到下排序。"""
|
||
blocks = [
|
||
_make_block("第二行", y1=100, y2=130),
|
||
_make_block("第一行", y1=10, y2=40),
|
||
]
|
||
lines = sort_and_group(blocks)
|
||
assert len(lines) == 2
|
||
assert lines[0][0].text == "第一行"
|
||
assert lines[1][0].text == "第二行"
|
||
|
||
def test_same_line_grouping(self) -> None:
|
||
"""Y 坐标接近的 blocks 归为同一行，按 X 排序。"""
|
||
blocks = [
|
||
_make_block("右", x1=100, x2=130, y1=10, y2=30),
|
||
_make_block("左", x1=0, x2=30, y1=10, y2=30),
|
||
]
|
||
lines = sort_and_group(blocks)
|
||
assert len(lines) == 1
|
||
assert len(lines[0]) == 2
|
||
assert lines[0][0].text == "左"
|
||
assert lines[0][1].text == "右"
|
||
|
||
def test_multi_line_grouping(self) -> None:
|
||
"""多行正确归组。"""
|
||
blocks = [
|
||
_make_block("B1", x1=0, x2=30, y1=100, y2=130),
|
||
_make_block("A1", x1=0, x2=30, y1=10, y2=40),
|
||
_make_block("A2", x1=40, x2=70, y1=10, y2=40),
|
||
]
|
||
lines = sort_and_group(blocks)
|
||
assert len(lines) == 2
|
||
# 第一行：A1、A2
|
||
assert len(lines[0]) == 2
|
||
assert lines[0][0].text == "A1"
|
||
assert lines[0][1].text == "A2"
|
||
# 第二行：B1
|
||
assert len(lines[1]) == 1
|
||
assert lines[1][0].text == "B1"
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 多栏检测
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestMultiColumn:
|
||
"""多栏检测。"""
|
||
|
||
def assert_one_column(self, blocks: list, page_width: float = 800) -> None:
|
||
cols = _detect_columns(blocks, page_width)
|
||
assert len(cols) == 1, f"应检测为单栏，实际 {len(cols)} 栏"
|
||
|
||
def assert_two_columns(self, blocks: list, page_width: float = 800) -> None:
|
||
cols = _detect_columns(blocks, page_width)
|
||
assert len(cols) == 2, f"应检测为双栏，实际 {len(cols)} 栏"
|
||
assert len(cols[0]) > 0
|
||
assert len(cols[1]) > 0
|
||
|
||
def test_single_column_few_blocks(self) -> None:
|
||
"""少于 4 个 block 不触发多栏检测。"""
|
||
self.assert_one_column([_make_block("A")])
|
||
|
||
def test_single_column_no_gap(self) -> None:
|
||
"""无间隙 → 单栏。"""
|
||
blocks = [
|
||
_make_block("左1", x1=0, x2=100, y1=10, y2=30),
|
||
_make_block("右1", x1=100, x2=200, y1=10, y2=30),
|
||
_make_block("左2", x1=0, x2=100, y1=40, y2=60),
|
||
_make_block("右2", x1=100, x2=200, y1=40, y2=60),
|
||
]
|
||
self.assert_one_column(blocks)
|
||
|
||
def test_two_columns_clear_gap(self) -> None:
|
||
"""有明显水平间隙 → 双栏。"""
|
||
blocks = [
|
||
_make_block("左栏 A", x1=0, x2=150, y1=10, y2=30),
|
||
_make_block("左栏 B", x1=0, x2=150, y1=40, y2=60),
|
||
_make_block("右栏 A", x1=500, x2=650, y1=15, y2=35),
|
||
_make_block("右栏 B", x1=500, x2=650, y1=50, y2=70),
|
||
]
|
||
self.assert_two_columns(blocks, page_width=800)
|
||
|
||
def test_sort_and_group_two_columns(self) -> None:
|
||
"""双栏排版：左栏两行，右栏两行。"""
|
||
blocks = [
|
||
_make_block("右A", x1=500, x2=650, y1=50, y2=70),
|
||
_make_block("左B", x1=0, x2=150, y1=40, y2=60),
|
||
_make_block("左A", x1=0, x2=150, y1=10, y2=30),
|
||
_make_block("右B", x1=500, x2=650, y1=80, y2=100),
|
||
]
|
||
lines = sort_and_group(blocks, page_width=800)
|
||
# 左栏两行 → 右栏两行
|
||
assert len(lines) == 4
|
||
# 左栏第一行
|
||
assert lines[0][0].text == "左A"
|
||
# 左栏第二行
|
||
assert lines[1][0].text == "左B"
|
||
# 右栏第一行
|
||
assert lines[2][0].text == "右A"
|
||
# 右栏第二行
|
||
assert lines[3][0].text == "右B"
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# group_paragraphs
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestGroupParagraphs:
|
||
"""段落分组。"""
|
||
|
||
def test_empty(self) -> None:
|
||
assert group_paragraphs([]) == []
|
||
|
||
def test_single_para(self) -> None:
|
||
"""多行同段落（用长文本避免触发标题检测）。"""
|
||
lines = [
|
||
[_make_block("第一行正文内容这里是长文本")],
|
||
[_make_block("第二行正文内容再接一些字")],
|
||
]
|
||
paras = group_paragraphs(lines)
|
||
assert len(paras) == 1
|
||
assert paras[0]["type"] == "text"
|
||
assert paras[0]["type"] == "text"
|
||
|
||
def test_paragraph_break_by_gap(self) -> None:
|
||
"""行间距大 → 分段（需要 3+ 行，单行间距无法超过自身中位值的 1.8 倍）。"""
|
||
lines = [
|
||
[_make_block("第一段正文在这里第一行啊", y1=10, y2=30)],
|
||
[_make_block("第一段正文在这里第二行", y1=40, y2=60)],
|
||
[_make_block("第二段正文这里间距很大", y1=200, y2=230)],
|
||
]
|
||
paras = group_paragraphs(lines, page_width=800)
|
||
assert len(paras) == 2
|
||
# 第一段：两行
|
||
assert paras[0]["type"] == "text"
|
||
assert "第一段正文在这里" in paras[0]["text"]
|
||
# 第二段：一行
|
||
assert "第二段正文" in paras[1]["text"]
|
||
|
||
def test_heading_paragraph(self) -> None:
|
||
"""短行 → 标题段。"""
|
||
lines = [[_make_block("第一章")], [_make_block("正文内容")]]
|
||
paras = group_paragraphs(lines)
|
||
assert len(paras) >= 1
|
||
# 标题段
|
||
assert paras[0]["type"] == "heading"
|
||
|
||
def test_list_paragraph(self) -> None:
|
||
"""列表项段落。"""
|
||
lines = [
|
||
[_make_block("- 项目一", y1=10, y2=30)],
|
||
[_make_block("- 项目二", y1=40, y2=60)],
|
||
]
|
||
paras = group_paragraphs(lines)
|
||
assert len(paras) == 1
|
||
assert paras[0]["type"] == "list_item"
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# build_markdown_text
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestBuildMarkdown:
|
||
"""完整的 Markdown 转换。"""
|
||
|
||
def test_empty(self) -> None:
|
||
md, lines, paras = build_markdown_text([])
|
||
assert md == ""
|
||
assert lines == []
|
||
assert paras == []
|
||
|
||
def test_simple_text(self) -> None:
|
||
"""简单段落转换。"""
|
||
blocks = [
|
||
_make_block("你好世界"),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "你好世界" in md
|
||
|
||
def test_heading_conversion(self) -> None:
|
||
"""标题应加 # 前缀（级别基于字体大小）。"""
|
||
# 标题 block 更大 → 推断为 h2
|
||
blocks = [
|
||
_make_block("第一章", x1=0, x2=100, y1=10, y2=50), # 40px 高
|
||
_make_block("正文内容长一些避免被检测为标题", x1=0, x2=200, y1=60, y2=80), # 20px 高
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "### 第一章" in md or "## 第一章" in md
|
||
assert "正文内容长一些避免被检测为标题" in md
|
||
|
||
def test_list_conversion(self) -> None:
|
||
"""列表项应加 - 前缀。"""
|
||
blocks = [
|
||
_make_block("- 苹果"),
|
||
_make_block("- 香蕉"),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "- 苹果" in md
|
||
assert "- 香蕉" in md
|
||
|
||
def test_centered_heading(self) -> None:
|
||
"""居中的短行 → heading。"""
|
||
# 中心 X 在页面宽度中间
|
||
blocks = [
|
||
_make_block("欢迎光临", x1=300, x2=500, y1=10, y2=40),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks, page_width=800)
|
||
assert "欢迎光临" in md
|
||
|
||
def test_heading_level_h1(self) -> None:
|
||
"""居中大字 → h1（# 前缀）。"""
|
||
blocks = [
|
||
_make_block("年度报告", x1=300, x2=500, y1=10, y2=50),
|
||
_make_block("正文内容较小", x1=0, x2=100, y1=60, y2=75),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks, page_width=800)
|
||
assert "# 年度报告" in md or "## 年度报告" in md
|
||
|
||
def test_heading_level_h3(self) -> None:
|
||
"""等大字体短行 → h3（### 前缀）。"""
|
||
blocks = [
|
||
_make_block("摘要", x1=0, x2=100, y1=10, y2=30),
|
||
_make_block("第一章内容文字", x1=0, x2=100, y1=40, y2=70),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks, page_width=800)
|
||
# "摘要" ≤ 8 字无标点 → heading, 但字体与正文同大 → h3
|
||
assert "### 摘要" in md
|
||
|
||
def test_separator_conversion(self) -> None:
|
||
"""分隔线 → ---。"""
|
||
blocks = [
|
||
_make_block("---"),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "---" in md
|
||
|
||
def test_multi_line_same_para(self) -> None:
|
||
"""多行近距离内容合并在一个段落（长文本避免标题检测）。"""
|
||
blocks = [
|
||
_make_block("第一行长文本避免被检测为标题行", y1=10, y2=30),
|
||
_make_block("第二行避免被检测为标题", y1=35, y2=55),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "第一行长文本避免被检测为标题行" in md
|
||
assert "第二行避免被检测为标题" in md
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════
|
||
# 边缘情况
|
||
# ═══════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestEdgeCases:
|
||
"""边缘情况。"""
|
||
|
||
def test_single_block_only(self) -> None:
|
||
blocks = [_make_block("Hello")]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert md == "Hello"
|
||
|
||
def test_blocks_with_same_y(self) -> None:
|
||
"""完全同一 Y 位置的多个 blocks（罕见但需处理）。"""
|
||
blocks = [
|
||
_make_block("A", x1=0, x2=30, y1=10, y2=30),
|
||
_make_block("B", x1=40, x2=70, y1=10, y2=30),
|
||
]
|
||
lines = sort_and_group(blocks)
|
||
assert len(lines) == 1
|
||
assert len(lines[0]) == 2
|
||
|
||
def test_blocks_with_negligible_height(self) -> None:
|
||
"""高度为 0 的 block 不应崩。"""
|
||
block = TextBlock(
|
||
text="边缘",
|
||
confidence=0.5,
|
||
bbox=BBox(points=[[0, 0], [0, 0], [0, 0], [0, 0]]),
|
||
)
|
||
assert _block_height(block) == 0.0
|
||
lines = sort_and_group([block])
|
||
assert len(lines) == 1
|
||
|
||
def test_confidence_range(self) -> None:
|
||
"""置信度在 0~1 范围内。"""
|
||
block = _make_block("test", confidence=1.0)
|
||
assert 0 <= block.confidence <= 1
|
||
|
||
def test_long_line_content(self) -> None:
|
||
"""长行（> 40 字符）→ text 而非 heading。"""
|
||
long_text = "A" * 50
|
||
blocks = [_make_block(long_text)]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert long_text in md
|
||
assert "##" not in md
|
||
|
||
def test_empty_text_blocks(self) -> None:
|
||
"""空文本 block 不应造成影响。"""
|
||
blocks = [
|
||
_make_block(""),
|
||
_make_block("实际内容"),
|
||
]
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "实际内容" in md
|
||
|
||
def test_numeric_only(self) -> None:
|
||
"""纯数字行 → numbered。"""
|
||
blocks = [_make_block("42")]
|
||
# 短行数字 → numbered（先于 heading 触发）
|
||
md, lines, paras = build_markdown_text(blocks)
|
||
assert "42" in md
|