rag-from-scratch/tests/chunks/test_segmenter.py

import pytest
from chunks.parser import parse
from chunks.config import ChunkerConfig
from chunks.segmenter import segment


@pytest.fixture
def cfg():
    return ChunkerConfig()


def test_paragraph_block(cfg):
    tokens, lines = parse("# Titolo\n\nParagrafo di testo.\n")
    blocks = segment(tokens, lines, cfg)
    para = next(b for b in blocks if b.kind == "paragraph")
    assert "Paragrafo di testo." in para.content
    assert para.header_path == [{"level": 1, "text": "Titolo"}]


def test_code_block_is_atomic(cfg):
    tokens, lines = parse("# T\n\n```python\nprint('hello')\n```\n")
    blocks = segment(tokens, lines, cfg)
    code = next(b for b in blocks if b.kind == "code")
    assert code.atomic is True


def test_table_block_is_atomic(cfg):
    tokens, lines = parse("# T\n\n| A | B |\n|---|---|\n| 1 | 2 |\n")
    blocks = segment(tokens, lines, cfg)
    table = next(b for b in blocks if b.kind == "table")
    assert table.atomic is True


def test_list_block_is_atomic(cfg):
    tokens, lines = parse("# T\n\n- item 1\n- item 2\n")
    blocks = segment(tokens, lines, cfg)
    lst = next(b for b in blocks if b.kind == "list")
    assert lst.atomic is True


def test_heading_stack_depth(cfg):
    md = "# H1\n\n## H2\n\n### H3\n\nTesto.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    para = next(b for b in blocks if b.kind == "paragraph")
    assert para.header_path == [
        {"level": 1, "text": "H1"},
        {"level": 2, "text": "H2"},
        {"level": 3, "text": "H3"},
    ]


def test_context_depth_limits_header_path():
    cfg = ChunkerConfig(context_depth=2)
    md = "# H1\n\n## H2\n\n### H3\n\nTesto.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    para = next(b for b in blocks if b.kind == "paragraph")
    assert len(para.header_path) == 2


def test_skip_headings(cfg):
    md = "# Titolo\n\n## Indice\n\nContenuto saltato.\n\n## Sezione Reale\n\nContenuto reale.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    contents = [b.content for b in blocks]
    assert not any("saltato" in c for c in contents)
    assert any("reale" in c for c in contents)


def test_skip_pre_heading(cfg):
    md = "Testo prima del primo heading.\n\n# Titolo\n\nTesto dopo.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    contents = [b.content for b in blocks]
    assert not any("prima del primo" in c for c in contents)
    assert any("dopo" in c for c in contents)


def test_skip_pre_heading_disabled():
    cfg = ChunkerConfig(skip_pre_heading=False)
    md = "Testo prima.\n\n# Titolo\n\nTesto dopo.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    contents = [b.content for b in blocks]
    assert any("prima" in c for c in contents)


def test_block_ids_unique(cfg):
    md = "# T\n\nPara 1.\n\nPara 2.\n\nPara 3.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    ids = [b.id for b in blocks]
    assert len(ids) == len(set(ids))


def test_heading_reset_on_same_level(cfg):
    md = "# H1a\n\n## H2a\n\n# H1b\n\nTesto.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    para = next(b for b in blocks if b.kind == "paragraph")
    assert para.header_path == [{"level": 1, "text": "H1b"}]


def test_thematic_break(cfg):
    md = "# T\n\nPara.\n\n---\n\nPara 2.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    kinds = [b.kind for b in blocks]
    assert "thematic_break" in kinds


def test_blockquote(cfg):
    md = "# T\n\n> Una citazione.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    bq = next((b for b in blocks if b.kind == "blockquote"), None)
    assert bq is not None
    assert "citazione" in bq.content


def test_source_line_numbers(cfg):
    md = "# Titolo\n\nParagrafo.\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    para = next(b for b in blocks if b.kind == "paragraph")
    assert para.start_line == 2
    assert para.end_line == 3


def test_math_block_is_atomic(cfg):
    md = "# T\n\n$$\nE=mc^2\n$$\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    math = next((b for b in blocks if b.kind == "math"), None)
    assert math is not None
    assert math.atomic is True
    assert "E=mc^2" in math.content


def test_math_block_plain_text(cfg):
    md = "# T\n\n$$\n\\int_0^1 f(x) dx\n$$\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    math = next(b for b in blocks if b.kind == "math")
    assert math.plain_text == "[formula]"


def test_html_block_with_table(cfg):
    md = "# T\n\n<table><tr><td>A</td></tr></table>\n"
    tokens, lines = parse(md)
    blocks = segment(tokens, lines, cfg)
    html = next((b for b in blocks if b.kind == "html"), None)
    assert html is not None
    assert "<table" in html.content