brain-of-reese/tests/unit/test_chunker.py

"""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"