"""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 ( convert_pdf_with_splitting, 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)) @pytest.mark.asyncio async def test_convert_unlinks_slice_tempfile_on_write_failure(tmp_path, monkeypatch): """If writing the slice bytes to the tempfile raises (e.g. ENOSPC), the tempfile is created on disk but never reaches the converter. The original error must surface and the orphaned file must be removed.""" src = _make_pdf(2, tmp_path) monkeypatch.setattr("tempfile.tempdir", str(tmp_path)) import tempfile as _tempfile real_factory = _tempfile.NamedTemporaryFile created: list[Path] = [] def _make_failing_tempfile(*args, **kwargs): handle = real_factory(*args, **kwargs) created.append(Path(handle.name)) original_write = handle.write def _raising_write(_data): # Touch the underlying file once so we know the path exists on # disk and the cleanup actually has something to remove. original_write(b"\0") raise OSError("No space left on device") handle.write = _raising_write return handle monkeypatch.setattr(_tempfile, "NamedTemporaryFile", _make_failing_tempfile) class _UnusedConverter: async def convert_file(self, path: Path, *, source_uri): raise AssertionError("converter must not be reached on write failure") with pytest.raises(OSError, match="No space left on device"): await convert_pdf_with_splitting( _UnusedConverter(), # ty: ignore[invalid-argument-type] src, source_uri=None, slice_size=1, ) assert created, "expected NamedTemporaryFile to be called at least once" leftover = [p for p in created if p.exists()] assert leftover == [], f"tempfiles leaked after write failure: {leftover}" @pytest.mark.asyncio async def test_convert_aborts_and_cleans_up_on_mid_stream_slice_failure( tmp_path, monkeypatch ): """When converting slice 2 of 3 fails, the whole call must raise ValueError naming the failed slice's page range, every tempfile created along the way must be deleted, and the source PDF handle must be closed. """ src = _make_pdf(7, tmp_path) # Pin tempfiles to a per-test dir so we can list leaks deterministically. monkeypatch.setattr("tempfile.tempdir", str(tmp_path)) calls: list[Path] = [] class _FlakyConverter: async def convert_file(self, path: Path, *, source_uri): calls.append(path) if len(calls) == 2: raise RuntimeError("docling exploded") from docling_core.types.doc.document import DoclingDocument return DoclingDocument(name="slice") with pytest.raises(ValueError, match="pages 4-6"): await convert_pdf_with_splitting( _FlakyConverter(), # ty: ignore[invalid-argument-type] src, source_uri=None, slice_size=3, ) # All converter inputs were under tmp_path (the pinned tempdir)… assert all(str(p).startswith(str(tmp_path)) for p in calls) # …and none of them are still on disk. leftover = [p for p in calls if p.exists()] assert leftover == [], f"tempfiles leaked: {leftover}" # We aborted after slice 2; slice 3 was never attempted. assert len(calls) == 2 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]