tests/unit/test_table_chunking.py

"""
Unit tests for table-aware chunking.

The defect these lock down: a Markdown table contains no sentence-ending
punctuation, so the sentence splitter treats an entire table as one "sentence"
and _group_into_chunks emitted it whole regardless of size. Measured before the
fix, a 400-row table produced a single 14,703-character chunk — roughly twice
gemini-embedding-001's input limit, so most of the table was silently truncated
and never indexed, and what survived was one vector averaging 400 unrelated
rows.
"""

import pytest

from app.config import settings
from app.services.chunking import ChunkingService, split_markdown_table
from app.services.extractors import PageContent

pytestmark = pytest.mark.unit


def make_table(n_rows: int, label: str = "[Table 1 on page 5]") -> str:
    rows = [f"| {h} ft | {36 + h} in | 4.0 ft | Resin |" for h in range(n_rows)]
    return (
        f"{label}\n"
        "| Mining Height | Bolt Length | Spacing | Support |\n"
        "| --- | --- | --- | --- |\n" + "\n".join(rows)
    )


def chunk(text: str, page: int = 5):
    pages = [PageContent(page_number=page, text=text, char_start=0, char_end=len(text))]
    return ChunkingService().chunk_document(full_text=text, pages=pages)


# ── The core defect ────────────────────────────────────────────────────────────


def test_large_table_no_longer_becomes_one_giant_chunk():
    chunks = chunk(make_table(400))
    assert len(chunks) > 1
    for c in chunks:
        assert len(c.text) <= settings.MAX_CHUNK_CHARS


def test_no_rows_are_lost_when_splitting():
    """Splitting must not drop data — that would be worse than the bug."""
    n = 400
    chunks = chunk(make_table(n))
    recovered = sum(c.text.count("| Resin |") for c in chunks)
    assert recovered == n


def test_every_part_repeats_the_header():
    """A bare row of numbers is not an answer without its column names."""
    chunks = chunk(make_table(400))
    for c in chunks:
        assert "| Mining Height | Bolt Length | Spacing | Support |" in c.text
        assert "| --- | --- | --- | --- |" in c.text


def test_parts_are_labelled_so_partial_lists_are_not_read_as_complete():
    chunks = chunk(make_table(400))
    labels = [
        ln for c in chunks for ln in c.text.splitlines() if ln.startswith("[Table")
    ]
    assert len(labels) == len(chunks)
    assert "part 1 of" in labels[0]
    # Well-formed brackets, not "[Table 1 on page 5 (part 1 of 2]"
    for label in labels:
        assert label.startswith("[") and label.endswith("]")
        assert label.count("[") == 1 and label.count("]") == 1


def test_row_column_relationship_survives_the_split():
    """The value must still be readable as belonging to its row and column."""
    chunks = chunk(make_table(400))
    holder = next(c for c in chunks if "| 300 ft | 336 in |" in c.text)
    assert "| Mining Height | Bolt Length |" in holder.text


# ── Behaviour that must not change ─────────────────────────────────────────────


def test_small_table_is_left_alone():
    chunks = chunk(make_table(20))
    assert len(chunks) == 1
    assert "part 1 of" not in chunks[0].text  # no spurious part labels


def test_prose_only_document_is_unaffected():
    prose = " ".join(f"Sentence number {i} about mine ventilation." for i in range(40))
    before = ChunkingService().chunk_document(
        full_text=prose,
        pages=[
            PageContent(page_number=1, text=prose, char_start=0, char_end=len(prose))
        ],
    )
    assert before
    for c in before:
        assert len(c.text) <= settings.MAX_CHUNK_CHARS


def test_table_caption_is_kept_with_the_table():
    """
    A one-line caption is below the minimum-words filter, so emitting it
    separately would discard it. It must ride with the first part.
    """
    caption = "Roof bolt patterns shall conform to the schedule below."
    chunks = chunk(f"{caption}\n\n{make_table(400)}")
    assert caption in chunks[0].text


def test_long_prose_before_a_table_is_split_separately():
    long_prose = " ".join(
        f"Ventilation requirement number {i} applies to all working sections."
        for i in range(120)
    )
    chunks = chunk(f"{long_prose}\n\n{make_table(200)}")
    assert len(chunks) > 2
    for c in chunks:
        assert len(c.text) <= settings.MAX_CHUNK_CHARS
    # Both kinds of content survived.
    assert any("Ventilation requirement" in c.text for c in chunks)
    assert sum(c.text.count("| Resin |") for c in chunks) == 200


def test_table_without_a_label_still_splits():
    rows = [f"| {h} | {h * 2} |" for h in range(600)]
    table = "| A | B |\n| --- | --- |\n" + "\n".join(rows)
    chunks = chunk(table)
    assert len(chunks) > 1
    for c in chunks:
        assert "| A | B |" in c.text
        assert len(c.text) <= settings.MAX_CHUNK_CHARS


def test_short_table_chunks_are_not_dropped_by_min_words_filter():
    """Tables are exempt from the tiny-fragment filter — the data matters."""
    tiny = "[Table 1 on page 2]\n| A | B |\n| --- | --- |\n| 1 | 2 |"
    chunks = chunk(tiny)
    assert chunks, "a small table must not be discarded as a fragment"
    assert "| 1 | 2 |" in chunks[0].text


# ── The infinite loop this work uncovered ──────────────────────────────────────


def test_oversized_sentence_after_prose_terminates():
    """
    Regression: _group_into_chunks used to spin forever here.

    A Markdown table has no sentence boundaries, so it merges into the
    preceding sentence and produces one "sentence" larger than chunk_size.
    The old loop flushed, restored a non-empty overlap, and re-tested the same
    oversized sentence without advancing — appending chunks until memory ran
    out. Ingestion of any document with prose followed by a large table hung
    permanently.
    """
    prose = " ".join(
        f"Ventilation requirement number {i} applies to all sections."
        for i in range(120)
    )
    table = "| A | B |\n| --- | --- |\n" + "\n".join(
        f"| {h} | {h * 2} |" for h in range(200)
    )

    chunks = chunk(f"{prose}\n\n{table}")

    assert chunks
    assert len(chunks) < 100, "runaway chunk count suggests the loop is spinning"
    for c in chunks:
        assert len(c.text) <= settings.MAX_CHUNK_CHARS


def test_single_oversized_sentence_alone_terminates():
    huge = "word " * 5000  # one 5000-word "sentence", no boundaries
    chunks = chunk(huge)
    assert chunks
    for c in chunks:
        assert len(c.text) <= settings.MAX_CHUNK_CHARS


def test_repeated_identical_sentences_get_distinct_offsets():
    """
    The old code looked sentences back up with list.index(), which returns the
    first match — so a repeated sentence attributed later chunks to the wrong
    character offset, and therefore the wrong page.
    """
    repeated = "The mine foreman shall record the reading. " * 200
    pages = [
        PageContent(page_number=1, text=repeated[:2000], char_start=0, char_end=2000),
        PageContent(
            page_number=2,
            text=repeated[2000:],
            char_start=2000,
            char_end=len(repeated),
        ),
    ]
    chunks = ChunkingService().chunk_document(full_text=repeated, pages=pages)

    assert len(chunks) > 1
    starts = [c.char_start for c in chunks]
    assert starts == sorted(starts), "chunk offsets must advance monotonically"
    assert len(set(starts)) > 1, "all chunks collapsed onto the same offset"


# ── split_markdown_table directly ──────────────────────────────────────────────


def test_split_returns_table_unchanged_when_it_fits():
    table = make_table(5)
    assert split_markdown_table(table, max_chars=10_000) == [table.strip()]


def test_split_handles_malformed_table_without_raising():
    assert split_markdown_table("", 1000) == []
    assert split_markdown_table("| only a header |", 1000) == ["| only a header |"]


def test_every_part_stays_within_budget():
    parts = split_markdown_table(make_table(500), max_chars=2000)
    assert len(parts) > 1
    for p in parts:
        assert len(p) <= 2000


def test_oversized_single_row_is_emitted_rather_than_dropped():
    """A row wider than the budget cannot be split further — keep it anyway."""
    wide = "| " + "x" * 3000 + " |"
    table = f"| A |\n| --- |\n{wide}"
    parts = split_markdown_table(table, max_chars=500)
    assert any("x" * 3000 in p for p in parts)