ProbabilityRAG / tests /test_chunker.py
Piero7's picture
Add unit test suite (40 tests) + fix requirements for fresh installs
3d62c28
Raw
History Blame Contribute Delete
3.32 kB
"""Unit tests for src/chunker.py — the one rule that matters: statements start
chunks, Proof/Solution glue to the statement before them."""
from src.chunker import (
CHILD_MAX_TOK, CHARS_PER_TOKEN, Chunk, chunk_markdown, est_tokens,
split_children, _paragraphs,
)
def _mk(md: str):
return chunk_markdown(md, doc_id="doc", source="test.md")
def test_statement_starts_new_chunk_and_proof_glues():
md = (
"# 6.2 Variance of Discrete Random Variables\n\n"
"Theorem 6.9 The variance is additive for independent variables.\n\n"
"**Proof.** Expand the square and use linearity.\n\n"
"Example 6.1 Compute the variance of one die roll.\n"
)
chunks = _mk(md)
assert len(chunks) == 2
assert chunks[0].type == "theorem"
assert "Proof." in chunks[0].text # glued, never its own chunk
assert chunks[1].type == "example"
def test_chapter_pinned_by_numbering_not_heading_depth():
md = (
"# Chapter 6\n\nIntro prose.\n\n"
"# 6.2 Variance\n\nSection prose.\n\n"
"# Examples\n\nExample 6.5 Something.\n"
)
chunks = _mk(md)
assert all(c.chapter == "Chapter 6: Expected Value and Variance" for c in chunks)
# "Examples" is a generic header: it must NOT overwrite the numbered section.
assert chunks[-1].section == "6.2 Variance"
def test_front_matter_does_not_inherit_chapter():
md = "# Chapter 1\n\nReal content.\n\n# Index\n\nbinomial, 45\n"
chunks = _mk(md)
assert chunks[0].chapter.startswith("Chapter 1")
assert chunks[-1].chapter == "" # store.retrieval_units drops these
def test_unnumbered_topic_header_refines_section_only():
md = "# Chapter 2\n\nA.\n\n## Standard Deviation\n\nB.\n"
chunks = _mk(md)
assert chunks[-1].chapter.startswith("Chapter 2")
assert chunks[-1].section == "Standard Deviation"
def test_statement_marker_inside_display_math_does_not_split():
md = "Some setup text.\n$$\nExample = 5\n$$\nmore of the same block.\n"
chunks = _mk(md)
assert len(chunks) == 1
def test_small_parent_has_no_children():
p = Chunk(id="d::p1", doc_id="d", source="s", type="text", chapter="", section="",
text="short", n_tokens=est_tokens("short"))
assert split_children(p) == []
def test_large_parent_splits_with_parent_id_and_full_coverage():
paras = [f"Paragraph {i} " + "word " * 60 for i in range(10)]
text = "\n\n".join(paras)
p = Chunk(id="d::p1", doc_id="d", source="s", type="text", chapter="C", section="S",
text=text, n_tokens=est_tokens(text))
assert p.n_tokens > CHILD_MAX_TOK
kids = split_children(p)
assert len(kids) >= 2
assert all(k.parent_id == "d::p1" and not k.is_parent for k in kids)
assert all(k.chapter == "C" and k.section == "S" for k in kids)
joined = "\n".join(k.text for k in kids)
assert all(f"Paragraph {i}" in joined for i in range(10)) # nothing dropped
def test_paragraphs_keep_display_blocks_intact():
text = "before\n\n$$\na = 1\n\nb = 2\n$$\n\nafter"
paras = _paragraphs(text)
assert len(paras) == 3
assert "a = 1" in paras[1] and "b = 2" in paras[1] # blank line inside $$ kept
def test_est_tokens_floor():
assert est_tokens("") == 1
assert est_tokens("x" * (CHARS_PER_TOKEN * 10)) == 10