tests/test_pdf.py

import pytest
from PIL import Image

from unlimited_ocr_rdna4.errors import InferenceError
from unlimited_ocr_rdna4.pdf import PDFPageRenderer


def _make_pdf(path) -> None:
    image = Image.new("RGB", (120, 80), "white")
    try:
        image.save(path, format="PDF", resolution=72)
    finally:
        image.close()


def test_pdf_renders_one_page_at_a_time(tmp_path) -> None:
    pdf = tmp_path / "fixture.pdf"
    _make_pdf(pdf)
    with PDFPageRenderer(
        pdf,
        tmp_path,
        dpi=72,
        start_page=1,
        max_pages=1,
        max_page_pixels=20_000,
        max_total_pixels=20_000,
        max_page_rendered_bytes=1_000_000,
        max_rendered_bytes=1_000_000,
    ) as renderer:
        assert tuple(renderer.page_numbers) == (1,)
        rendered = renderer.render(1)
        assert rendered.is_file()


def test_pdf_rejects_page_over_pixel_limit(tmp_path) -> None:
    pdf = tmp_path / "fixture.pdf"
    _make_pdf(pdf)
    with (
        PDFPageRenderer(
            pdf,
            tmp_path,
            dpi=72,
            start_page=1,
            max_pages=1,
            max_page_pixels=1_000,
            max_total_pixels=20_000,
            max_page_rendered_bytes=1_000_000,
            max_rendered_bytes=1_000_000,
        ) as renderer,
        pytest.raises(InferenceError, match="would render"),
    ):
        renderer.render(1)