homelab-codex-ws/packages/kb-mail/tests/test_chunking.py

"""Unit tests for the shared chunker — moved 1:1 from
jobs/documents-ingest/tests/test_chunk_embed.py (module 5, faza mailowa, plan §3, Krok 0)."""
from __future__ import annotations

import pytest

from kb_mail.chunking import OVERLAP_CHARS, TARGET_CHARS, chunk_text, hard_split, split_paragraphs


def _paragraph(char: str, length: int) -> str:
    return char * length


class TestSplitParagraphs:
    def test_splits_on_blank_lines(self):
        text = "first para\n\nsecond para\n\nthird para"
        assert split_paragraphs(text) == ["first para", "second para", "third para"]

    def test_no_blank_lines_returns_single_paragraph(self):
        text = "line one\nline two\nline three"
        assert split_paragraphs(text) == [text]

    def test_drops_empty_fragments(self):
        text = "a\n\n\n\n\n\nb"
        assert split_paragraphs(text) == ["a", "b"]


class TestHardSplit:
    def test_overlap_equal_to_size_raises(self):
        # step = size - overlap would be 0 -> `start` never advances -> infinite loop.
        with pytest.raises(ValueError):
            hard_split("x" * 1000, size=100, overlap=100)

    def test_overlap_greater_than_size_raises(self):
        with pytest.raises(ValueError):
            hard_split("x" * 1000, size=100, overlap=150)

    def test_short_text_returns_single_chunk(self):
        assert hard_split("short", size=100, overlap=10) == ["short"]

    def test_splits_with_overlap(self):
        text = "A" * 5000
        chunks = hard_split(text, size=2400, overlap=600)
        assert len(chunks) == 3
        # consecutive windows overlap by exactly `overlap` characters
        assert chunks[0][-600:] == chunks[1][:600]
        assert chunks[1][-600:] == chunks[2][:600]
        # covers the full text, in order, no gaps
        assert chunks[0] + chunks[1][600:] + chunks[2][600:] == text

    def test_last_chunk_reaches_end_of_text(self):
        text = "B" * 5000
        chunks = hard_split(text, size=2400, overlap=600)
        assert chunks[-1] == text[-len(chunks[-1]):]
        assert text.endswith(chunks[-1])


class TestChunkText:
    def test_empty_content_returns_no_chunks(self):
        assert chunk_text("") == []
        assert chunk_text(None) == []
        assert chunk_text("   \n\n  ") == []

    def test_document_shorter_than_one_chunk_returns_single_chunk(self):
        text = "A short document."
        assert chunk_text(text, size=TARGET_CHARS, overlap=OVERLAP_CHARS) == [text]

    def test_multi_paragraph_document_splits_on_boundaries(self):
        para1 = _paragraph("a", 1000)
        para2 = _paragraph("b", 1000)
        para3 = _paragraph("c", 1000)
        text = f"{para1}\n\n{para2}\n\n{para3}"

        chunks = chunk_text(text, size=2400, overlap=600)

        assert len(chunks) == 2
        assert para1 in chunks[0]
        assert para2 in chunks[0]
        assert para3 in chunks[-1]

    def test_overlap_between_consecutive_chunks(self):
        para1 = _paragraph("a", 1000)
        para2 = _paragraph("b", 1000)
        para3 = _paragraph("c", 1000)
        text = f"{para1}\n\n{para2}\n\n{para3}"

        chunks = chunk_text(text, size=2400, overlap=600)

        # the tail of chunk[0] (the overlap window) reappears at the start of chunk[1]
        assert chunks[0][-600:] == chunks[1][: len(chunks[0][-600:])]

    def test_single_oversized_paragraph_falls_back_to_hard_split(self):
        text = "X" * 6000  # no blank lines at all
        chunks = chunk_text(text, size=2400, overlap=600)
        assert len(chunks) > 1
        assert all(len(c) <= 2400 for c in chunks)

    def test_paragraph_larger_than_target_is_hard_split_within_mixed_document(self):
        small = _paragraph("s", 100)
        huge = _paragraph("h", 6000)
        text = f"{small}\n\n{huge}"

        chunks = chunk_text(text, size=2400, overlap=600)

        assert chunks[0] == small
        assert len(chunks) > 2
        assert all(len(c) <= 2400 for c in chunks[1:])