Function test_markdown_chunking

tests/main/test_md_parser.py:233–309

Given a Markdown document with sections and sub-sections, this test verifies that: - The tree is built correctly from the document. - The chunking process produces distinct chunks with enriched header context. - A header-only node does not duplicate the header in its own chunk

(
    sample_markdown,
    markdown_sections,
    chunk_size_factor: int,
    rollup: bool,
)

@pytest.mark.parametrize("chunk_size_factor", [1.2, 100])
@pytest.mark.parametrize("rollup", [True, False])
def test_markdown_chunking(
    sample_markdown,
    markdown_sections,
    chunk_size_factor: int,
    rollup: bool,
):
    """
    Given a Markdown document with sections and sub-sections, this test verifies that:
      - The tree is built correctly from the document.
      - The chunking process produces distinct chunks with enriched header context.
      - A header-only node does not duplicate the header in its own chunk.

    The sample document has:
      - Chapter 1 with a preamble.
      - Section 1.1 with content.
      - Section 1.2 with bullet content.
      - Chapter 2 with its own content.
    """

    ch1, sec1_1, sec1_2, ch2 = markdown_sections
    chunk_size = chunk_size_factor * count_words(ch1.content)
    config = MarkdownChunkConfig(
        chunk_size=chunk_size,
        overlap_tokens=5,
        variation_percent=0.2,
        rollup=rollup,
    )

    # Structure-aware chunking of the text into enriched chunks.
    chunks: List[str] = chunk_markdown(sample_markdown, config)

    if rollup and chunk_size > count_words(sample_markdown):
        assert len(chunks) == 1, f"Expected 1 chunk, got {len(chunks)}"
        assert (
            chunks[0].split() == sample_markdown.split()
        ), "Chunk does not match original Markdown"
        # check that line-breaks in each section content are preserved
        for section in markdown_sections:
            assert (
                section.content in chunks[0]
            ), f"Section content not found in the chunk: {section.content}"

    if not rollup or chunk_size < count_words(sample_markdown):
        # Based on our document structure, we expect four chunks:
        # 1. Chapter 1's preamble content (enriched with prefix "# Chapter 1")
        # 2. Section 1.1 content (enriched with prefix "# Chapter 1 \n\n # Section 1.1")
        # 3. Section 1.2 content (enriched with prefix "# Chapter 1 \n\n # Section 1.2")
        # 4. Chapter 2 content (enriched with prefix "# Chapter 2")
        assert len(chunks) == 4, f"Expected 4 chunks, got {len(chunks)}"

        assert (
            chunks[0].split() == ch1.to_markdown().split()
        ), "Chunk 1 does not match Chapter 1 preamble"
        assert (
            ch1.content.strip() in chunks[0]
        ), "Chapter 1 content not preserved in Chunk 1"

        assert chunks[1].split() == (