111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
"""Unit tests for haiku.rag.converters.pdf_split.
|
|
|
|
The integration test that pins split-and-merge against a real-PDF baseline
|
|
lives in tests/test_converters.py — that path requires docling installed and
|
|
is gated accordingly. These tests cover the byte-level split mechanism and
|
|
the docling-core concatenate contract in isolation.
|
|
"""
|
|
|
|
import io
|
|
from pathlib import Path
|
|
|
|
import pypdfium2 as pdfium
|
|
import pytest
|
|
|
|
from haiku.rag.converters.pdf_split import iter_pdf_slices
|
|
|
|
|
|
def _make_pdf(page_count: int, tmp_path: Path) -> Path:
|
|
"""Synthesize a minimal valid multi-page PDF with pypdfium2. Each page is
|
|
an A4-sized blank. Returns the path."""
|
|
doc = pdfium.PdfDocument.new()
|
|
try:
|
|
for _ in range(page_count):
|
|
doc.new_page(width=595.0, height=842.0) # A4 in points
|
|
out = tmp_path / "synth.pdf"
|
|
with open(out, "wb") as f:
|
|
doc.save(f)
|
|
return out
|
|
finally:
|
|
doc.close()
|
|
|
|
|
|
def test_iter_pdf_slices_partitions_pages(tmp_path):
|
|
src = _make_pdf(7, tmp_path)
|
|
slices = list(iter_pdf_slices(src, slice_size=3))
|
|
|
|
assert [(s, e) for s, e, _ in slices] == [(1, 3), (4, 6), (7, 7)]
|
|
|
|
page_counts = []
|
|
for _, _, pdf_bytes in slices:
|
|
d = pdfium.PdfDocument(io.BytesIO(pdf_bytes))
|
|
try:
|
|
page_counts.append(len(d))
|
|
finally:
|
|
d.close()
|
|
assert page_counts == [3, 3, 1]
|
|
assert sum(page_counts) == 7
|
|
|
|
|
|
def test_iter_pdf_slices_single_slice_when_doc_fits(tmp_path):
|
|
src = _make_pdf(4, tmp_path)
|
|
slices = list(iter_pdf_slices(src, slice_size=10))
|
|
assert len(slices) == 1
|
|
start, end, pdf_bytes = slices[0]
|
|
assert (start, end) == (1, 4)
|
|
d = pdfium.PdfDocument(io.BytesIO(pdf_bytes))
|
|
try:
|
|
assert len(d) == 4
|
|
finally:
|
|
d.close()
|
|
|
|
|
|
def test_iter_pdf_slices_rejects_zero_slice_size(tmp_path):
|
|
src = _make_pdf(2, tmp_path)
|
|
with pytest.raises(ValueError, match="slice_size must be >= 1"):
|
|
list(iter_pdf_slices(src, slice_size=0))
|
|
|
|
|
|
def test_concatenate_shifts_page_nos_and_unique_self_refs():
|
|
"""Pins the docling-core contract we rely on: when two docs (each with
|
|
items on page 1) are concatenated, the second doc's items move to page 2
|
|
and self_refs across both stay unique."""
|
|
pytest.importorskip("docling_core")
|
|
from docling_core.types.doc.base import BoundingBox, CoordOrigin
|
|
from docling_core.types.doc.document import (
|
|
DoclingDocument,
|
|
PageItem,
|
|
ProvenanceItem,
|
|
Size,
|
|
)
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
def _make_one_page_doc(name: str, text: str) -> DoclingDocument:
|
|
d = DoclingDocument(name=name)
|
|
d.pages[1] = PageItem(page_no=1, size=Size(width=595.0, height=842.0))
|
|
d.add_text(
|
|
label=DocItemLabel.TEXT,
|
|
text=text,
|
|
prov=ProvenanceItem(
|
|
page_no=1,
|
|
bbox=BoundingBox(
|
|
l=0.0, t=0.0, r=100.0, b=20.0, coord_origin=CoordOrigin.TOPLEFT
|
|
),
|
|
charspan=(0, len(text)),
|
|
),
|
|
)
|
|
return d
|
|
|
|
a = _make_one_page_doc("a", "alpha")
|
|
b = _make_one_page_doc("b", "beta")
|
|
|
|
merged = DoclingDocument.concatenate([a, b])
|
|
|
|
assert len(merged.texts) == 2
|
|
refs = [t.self_ref for t in merged.texts]
|
|
assert len(set(refs)) == 2, f"self_refs collided: {refs}"
|
|
|
|
page_nos = sorted({p.page_no for t in merged.texts for p in t.prov})
|
|
assert page_nos == [1, 2], f"expected b's page 1 to shift to page 2, got {page_nos}"
|
|
|
|
assert sorted(merged.pages.keys()) == [1, 2]
|