ai-resume/corpus/tests/test_chunker.py

"""Chunking rules: size, overlap, determinism, and the hard one — no split inside
a fence, a table or front-matter (docs/phase-4-corpus.md DoD item 3)."""

from __future__ import annotations

import os
import unittest

from corpus import chunker, loader
from corpus.config import Config
from corpus.tests import support

REPO_ROOT = support.REPO_ROOT


def content_documents():
    cfg = Config.from_env(root=REPO_ROOT, provider="local-hashing", dim=32)
    return loader.load_documents(cfg)


class BlockParsingTest(unittest.TestCase):
    def test_fence_is_one_block(self):
        blocks = chunker.parse_blocks(support.FENCE)
        kinds = [block.kind for block in blocks]
        self.assertEqual(kinds, ["para", "fence", "para"])
        self.assertIn("chunk boundary must never land", blocks[1].text)

    def test_table_is_one_block(self):
        blocks = chunker.parse_blocks(support.TABLE)
        table = [block for block in blocks if block.kind == "table"]
        self.assertEqual(len(table), 1)
        self.assertEqual(len(table[0].text.splitlines()), 5)  # header, rule, 3 rows

    def test_front_matter_in_body_is_refused(self):
        with self.assertRaises(chunker.ChunkError):
            chunker.parse_blocks("---\nid: x\n---\n\n# Heading\n")

    def test_unterminated_fence_runs_to_end(self):
        blocks = chunker.parse_blocks("intro\n\n```go\nfunc main() {}\n")
        self.assertEqual(blocks[-1].kind, "fence")
        self.assertTrue(blocks[-1].text.rstrip().endswith("}"))

    def test_headings_start_blocks(self):
        blocks = chunker.parse_blocks("# One\n\ntext\n\n## Two\n\nmore\n")
        self.assertEqual([block.kind for block in blocks], ["heading", "para", "heading", "para"])

    def test_list_items_are_separate_blocks(self):
        blocks = chunker.parse_blocks("- one\n- two\n  continued\n- three\n")
        self.assertEqual([block.kind for block in blocks], ["list", "list", "list"])
        self.assertIn("continued", blocks[1].text)


class ChunkingTest(unittest.TestCase):
    def test_small_body_is_one_chunk(self):
        chunks = chunker.chunk_document("# Title\n\nA short body.\n")
        self.assertEqual(len(chunks), 1)
        self.assertEqual(chunks[0].index, 0)
        self.assertFalse(chunks[0].oversize)

    def test_long_body_respects_maximum_and_overlaps(self):
        chunks = chunker.chunk_document(support.PROSE)
        self.assertGreater(len(chunks), 1)
        for chunk in chunks:
            self.assertLessEqual(chunk.tokens, chunker.TOKEN_MAX, f"chunk {chunk.index} too large")
        # Overlap: consecutive chunks share at least one block.
        for previous, current in zip(chunks, chunks[1:]):
            self.assertTrue(
                set(id(block) for block in previous.blocks) & set(id(block) for block in current.blocks),
                f"chunks {previous.index} and {current.index} do not overlap",
            )

    def test_chunk_indices_are_sequential(self):
        chunks = chunker.chunk_document(support.PROSE)
        self.assertEqual([chunk.index for chunk in chunks], list(range(len(chunks))))

    def test_deterministic(self):
        first = chunker.chunk_document(support.PROSE)
        second = chunker.chunk_document(support.PROSE)
        self.assertEqual([chunk.text for chunk in first], [chunk.text for chunk in second])

    def test_code_fence_is_never_split(self):
        body = support.PROSE + "\n\n" + support.FENCE + "\n\n" + support.PROSE
        chunks = chunker.chunk_document(body)
        fences = [block for block in chunker.parse_blocks(body) if block.kind == "fence"]
        self.assertTrue(fences)
        for fence in fences:
            holders = [chunk for chunk in chunks if fence in chunk.blocks]
            self.assertTrue(holders, "the fenced block disappeared from every chunk")
            for holder in holders:
                self.assertIn(fence.text, holder.text, "a chunk holds part of a fenced block")

    def test_oversize_single_block_is_flagged_not_split(self):
        huge = "```go\n" + "\n".join(f"line{i} := {i}" for i in range(1200)) + "\n```\n"
        chunks = chunker.chunk_document(huge)
        self.assertEqual(len(chunks), 1)
        self.assertTrue(chunks[0].oversize)
        self.assertGreater(chunks[0].tokens, chunker.TOKEN_MAX)


class RealContentTest(unittest.TestCase):
    """The same rules, over every document in content/."""

    def test_no_boundary_splits_a_protected_region(self):
        for doc in content_documents():
            regions = chunker.protected_regions(doc.body)
            chunks = chunker.chunk_document(doc.body)
            chunker.verify_no_split(chunks, regions)  # raises on failure

    def test_every_chunk_is_whole_blocks(self):
        for doc in content_documents():
            expected = [block.text for block in chunker.parse_blocks(doc.body)]
            seen = []
            for chunk in chunker.chunk_document(doc.body):
                for block in chunk.blocks:
                    self.assertIn(block.text, chunk.text)
                    if not seen or seen[-1] != block.text:
                        seen.append(block.text)
            # Every block appears, in order, exactly as parsed (overlap repeats a
            # block at the start of the next chunk, which is expected).
            position = 0
            for text in expected:
                while position < len(seen) and seen[position] != text:
                    position += 1
                self.assertLess(position, len(seen), f"{doc.path}: block was dropped: {text[:40]!r}")
                position += 1

    def test_chunks_are_within_the_size_band(self):
        for doc in content_documents():
            for chunk in chunker.chunk_document(doc.body):
                if chunk.oversize:
                    self.assertEqual(len(chunk.blocks), 1)
                    continue
                self.assertLessEqual(chunk.tokens, chunker.TOKEN_MAX, f"{doc.path} chunk {chunk.index}")

    def test_content_tree_is_chunked_completely(self):
        total = 0
        for doc in content_documents():
            chunks = chunker.chunk_document(doc.body)
            self.assertTrue(chunks, f"{doc.path}: no chunks")
            total += len(chunks)
        self.assertGreater(total, 0)


if __name__ == "__main__":
    unittest.main()