Function test_chunking_sizes

tests/main/test_md_parser.py:314–381

Test that the chunking logic produces chunks that: - Have token counts between the lower and upper bounds (except possibly the final chunk) - Include the header enrichment in each chunk's text - Include the expected overlap between consecutive chunks

(
    chunk_size: int,
    rollup: bool,
)

@pytest.mark.parametrize("rollup", [False, True])
@pytest.mark.parametrize("chunk_size", [20, 500])
def test_chunking_sizes(
    chunk_size: int,
    rollup: bool,
):
    """
    Test that the chunking logic produces chunks that:
      - Have token counts between the lower and upper bounds
        (except possibly the final chunk)
      - Include the header enrichment in each chunk's text
      - Include the expected overlap between consecutive chunks
    """
    # Create a long text consisting of 200 repeated tokens ("word")
    long_text = " ".join(["word"] * 200)  # 200 tokens
    md_text = f"""# Chapter 1
{long_text}
"""

    # Set chunking configuration.
    # Here chunk_size=50 means that (with variation_percent=0.2)
    # we expect chunks to have between 40 and 60 tokens.
    config = MarkdownChunkConfig(
        chunk_size=chunk_size, rollup=rollup, overlap_tokens=5, variation_percent=0.2
    )

    # Produce the enriched chunks from the tree.
    chunks = chunk_markdown(md_text, config)

    # Compute the allowed bounds.
    lower_bound = config.chunk_size * (1 - config.variation_percent)
    upper_bound = config.chunk_size * (1 + config.variation_percent)

    # Verify each chunk's token count.
    # For all chunks except possibly the final one,
    # we expect at least lower_bound tokens.
    for i, chunk in enumerate(chunks):
        tokens = count_words(chunk)
        if i < len(chunks) - 1:
            assert (
                tokens >= lower_bound
            ), f"Chunk {i} has {tokens} tokens, expected at least {lower_bound}"
        assert (
            tokens <= 2 * upper_bound
        ), (  # relaxed check
            f"Chunk {i} has {tokens} tokens, expected at most {upper_bound}"
        )

    # Check that each chunk is enriched with the header context.
    # Each chunk's text should contain "Chapter 1" since that is the header path.
    for i, chunk in enumerate(chunks):
        assert "Chapter 1" in chunk, f"Chunk {i} is missing header enrichment"

    # Verify that consecutive chunks share the expected overlap.
    # For each consecutive pair of chunks, the last `overlap_tokens`
    # tokens of the previous chunk
    # should appear among the first tokens of the next chunk.
    if len(chunks) > 1:
        for i in range(len(chunks) - 1):
            prev_tokens = chunks[i].split()