Files
2026-07-13 12:08:54 +08:00

156 lines
4.7 KiB
Python

"""Unit tests for ``NativeMarkdownIRBuilder`` (markers + side tables → IRDoc)."""
from __future__ import annotations
from lightrag.parser.markdown.extract import (
drawing_marker,
equation_marker,
table_marker,
)
from lightrag.parser.markdown.ir_builder import NativeMarkdownIRBuilder
def _normalize(blocks, meta, *, document_name="doc.md"):
return NativeMarkdownIRBuilder().normalize(
blocks,
document_name=document_name,
asset_dir_name="doc.blocks.assets",
parse_metadata=meta,
)
def test_pipe_table_becomes_json_irtable_with_header_grid():
blocks = [
{
"heading": "H",
"level": 1,
"parent_headings": [],
"content": table_marker("t1"),
}
]
meta = {
"md_tables": {
"t1": {"kind": "pipe", "rows": [["a", "b"]], "header": [["H1", "H2"]]}
}
}
ir = _normalize(blocks, meta)
(table,) = ir.blocks[0].tables
assert "{{TBL:" in ir.blocks[0].content_template
assert table.rows == [["a", "b"]]
assert table.html is None
assert table.table_header == [["H1", "H2"]]
assert (table.num_rows, table.num_cols) == (1, 2)
def test_html_table_becomes_html_irtable_with_thead_and_dims():
html = "<table><thead><tr><th>K</th><th>V</th></tr></thead><tbody><tr><td>a</td><td>1</td></tr></tbody></table>"
blocks = [
{
"heading": "H",
"level": 1,
"parent_headings": [],
"content": table_marker("t1"),
}
]
ir = _normalize(blocks, {"md_tables": {"t1": {"kind": "html", "html": html}}})
(table,) = ir.blocks[0].tables
assert table.rows is None
assert table.html and table.html.startswith("<table>")
assert isinstance(table.table_header, str) and "<thead>" in table.table_header
assert table.body_override and "<thead>" in table.body_override
assert (table.num_rows, table.num_cols) == (2, 2)
def test_equation_is_block_level():
blocks = [
{
"heading": "H",
"level": 1,
"parent_headings": [],
"content": equation_marker("e1"),
}
]
ir = _normalize(blocks, {"md_equations": {"e1": "E = mc^2"}})
(eq,) = ir.blocks[0].equations
assert eq.is_block is True
assert eq.latex == "E = mc^2"
assert "{{EQ:" in ir.blocks[0].content_template
def test_local_drawing_dedups_assetspec_across_occurrences():
content = f"{drawing_marker('d1')} {drawing_marker('d2')}"
blocks = [{"heading": "H", "level": 1, "parent_headings": [], "content": content}]
meta = {
"md_drawings": {
"d1": {"kind": "local", "asset_ref": "sha:1", "fmt": "png", "src": "a.png"},
"d2": {"kind": "local", "asset_ref": "sha:1", "fmt": "png", "src": "a.png"},
},
"md_assets": {"sha:1": {"suggested_name": "a.png", "data": b"BYTES"}},
}
ir = _normalize(blocks, meta)
assert len(ir.blocks[0].drawings) == 2
assert all(d.asset_ref == "sha:1" for d in ir.blocks[0].drawings)
# One shared on-disk asset, carried as bytes for the writer to materialize.
assert len(ir.assets) == 1
assert ir.assets[0].ref == "sha:1"
assert ir.assets[0].source == b"BYTES"
def test_external_drawing_uses_path_override_and_no_asset():
blocks = [
{
"heading": "H",
"level": 1,
"parent_headings": [],
"content": drawing_marker("d1"),
}
]
meta = {
"md_drawings": {
"d1": {
"kind": "external",
"url": "http://x/y.png",
"fmt": "png",
"src": "http://x/y.png",
}
}
}
ir = _normalize(blocks, meta)
(drawing,) = ir.blocks[0].drawings
assert drawing.path_override == "http://x/y.png"
assert drawing.asset_ref == ""
assert ir.assets == []
def test_doc_title_from_first_h1_else_filename():
ir = _normalize(
[{"heading": "First", "level": 1, "parent_headings": [], "content": "# First"}],
{},
)
assert ir.doc_title == "First"
ir2 = _normalize(
[{"heading": "Sub", "level": 2, "parent_headings": [], "content": "## Sub"}],
{},
document_name="myfile.md",
)
assert ir2.doc_title == "myfile"
def test_positions_use_heading_anchor():
ir = _normalize(
[{"heading": "Sec", "level": 2, "parent_headings": [], "content": "## Sec"}],
{},
)
(pos,) = ir.blocks[0].positions
assert pos.type == "heading"
assert pos.anchor == "Sec"
def test_document_format_from_suffix():
ir = _normalize(
[{"heading": "X", "level": 1, "parent_headings": [], "content": "# X"}],
{},
document_name="a.textpack",
)
assert ir.document_format == "textpack"