"""Unit tests: markdown-aware chunker (PLAN §5 policy).""" from __future__ import annotations from itertools import pairwise import pytest from app.rag.chunker import HARD_MAX_CHARS, chunk_markdown, extract_title ANCHOR = "## Big" ANCHOR_PREFIX = f"{ANCHOR}\n\n" def _paras(n: int, char: str = "l", width: int = 300) -> list[str]: return [f"paragraph {i} " + char * (width - 12) for i in range(n)] def test_short_document_is_single_chunk() -> None: doc = "# Title\n\nJust some intro, no section headings at all." chunks = chunk_markdown(doc) assert chunks == [doc.strip()] def test_empty_and_whitespace_only_content() -> None: assert chunk_markdown("") == [] assert chunk_markdown(" \n\n \n") == [] def test_invalid_params_raise() -> None: with pytest.raises(ValueError): chunk_markdown("# x", target_chars=0) with pytest.raises(ValueError): chunk_markdown("# x", overlap_chars=-1) def test_splits_on_headings_and_keeps_nearest_heading() -> None: doc = ( "# Title\n" "intro line\n" "## Alpha\n" "alpha body\n" "### Beta\n" "beta body\n" "## Gamma\n" "gamma body\n" ) chunks = chunk_markdown(doc) assert chunks[0] == "# Title\nintro line" assert chunks[1] == "## Alpha\nalpha body" assert chunks[2] == "### Beta\nbeta body" assert chunks[3] == "## Gamma\ngamma body" def test_document_without_h1_starts_at_first_section() -> None: chunks = chunk_markdown("## Only\n\nbody") assert chunks == ["## Only\n\nbody"] def test_long_section_splits_with_overlap_and_anchor_on_every_chunk() -> None: body = "\n\n".join(_paras(10, width=138)) doc = f"{ANCHOR}\n\n{body}" chunks = chunk_markdown(doc, target_chars=800, overlap_chars=100) assert len(chunks) == 3 # Every chunk keeps its nearest preceding heading (the section anchor). assert all(c.startswith(ANCHOR) for c in chunks) # All chunks respect the target budget (anchor + packed body). assert all(len(c) <= 800 for c in chunks) # Overlap: the tail of each chunk is at the start of the next one. for prev, nxt in pairwise(chunks): assert nxt[len(ANCHOR_PREFIX) :].startswith(prev[-100:]) def test_overlap_zero_disables_tail_carryover() -> None: body = "\n\n".join(_paras(8, width=200)) chunks = chunk_markdown(f"{ANCHOR}\n\n{body}", target_chars=800, overlap_chars=0) assert len(chunks) >= 2 for prev, nxt in pairwise(chunks): assert not nxt[len(ANCHOR_PREFIX) :].startswith(prev[-50:]) def test_code_fences_stay_intact() -> None: doc = ( "## Section\n" "before fence\n" "```\n" "## fake heading inside fence\n" "\n" "still in fence\n" "```\n" "after fence\n" "## Other\n" "other body\n" ) chunks = chunk_markdown(doc) assert any(c.startswith("## Other") for c in chunks) # The fake heading inside the fence never opens a section… assert not any(c.startswith("## fake heading") for c in chunks) # …and the fence itself is whole in the chunk that contains it. fenced = [c for c in chunks if "still in fence" in c] assert len(fenced) == 1 assert "## fake heading inside fence" in fenced[0] assert fenced[0].count("```") == 2 # Blank lines inside the fence did not create extra paragraph chunks. assert not any(c.startswith("before fence\n\n") for c in chunks) def test_fence_block_is_atomic_across_forced_split() -> None: fence = "```\n" + "\n".join(f"code line {i}" for i in range(60)) + "\n```" doc = ( f"{ANCHOR}\n\npara A " + "a" * 300 + f"\n\n{fence}\n\npara B " + "b" * 300 + "\n\npara C " + "c" * 300 ) chunks = chunk_markdown(doc, target_chars=1000, overlap_chars=100) assert len(chunks) >= 2 # The whole fence (first and last code line) lives in one chunk — a # chunk boundary never falls inside a code block. assert any("code line 0" in c and "code line 59" in c for c in chunks) def test_oversized_fence_block_is_split_to_stay_under_hard_cap() -> None: """aipi's embedding endpoint caps requests at ~1024 input tokens — a multi-KB fenced code block must not survive chunking as one piece.""" code = "\n".join(f"int value_{i:03d} = {i}; // padding to grow the line" for i in range(160)) doc = ( "# Big Doc\n\n" "## Usage Example\n\n" f"```cpp\n{code}\n```\n\n" "## After\n\nDone.\n" ) chunks = chunk_markdown(doc) assert len(chunks) >= 3 # No chunk exceeds the hard cap (heading anchor adds a little). assert all(len(c) <= HARD_MAX_CHARS + 60 for c in chunks) # Content survives the split, and later sections are untouched. joined = "\n".join(chunks) assert "value_000" in joined assert "value_159" in joined assert any(c.startswith("## After") for c in chunks) def test_unclosed_fence_does_not_break_sections() -> None: doc = "## A\n\n```\nunterminated fence\n\n## B\n\nbody\n" chunks = chunk_markdown(doc) # "## B" is inside the unterminated fence → not a real heading. assert len(chunks) == 1 assert "## B" in chunks[0] def test_extract_title_prefers_h1() -> None: assert extract_title("# My Title\n\nbody") == "My Title" assert extract_title(" # Indented H1\nbody") == "" # ATX must be at col 0 assert extract_title("## not a title\n\nbody") == "" assert extract_title("## sub only", fallback="stem") == "stem" assert extract_title("", fallback="fallback") == "fallback"