tests/test_document_processing.py

from __future__ import annotations

from pathlib import Path

import pypdfium2 as pdfium
import pytest
from PIL import Image

from document_processing import (
    MAX_DOCUMENT_BYTES,
    MAX_IMAGE_EDGE,
    MAX_PAGES,
    DocumentError,
    detect_document_kind,
    prepare_document,
)


def make_image(path: Path, *, size: tuple[int, int] = (320, 200)) -> None:
    image = Image.new("RGB", size, "#ffffff")
    image.save(path, format="PNG")


def make_pdf(path: Path, pages: int) -> None:
    document = pdfium.PdfDocument.new()
    try:
        for _ in range(pages):
            page = document.new_page(width=612, height=792)
            page.close()
        document.save(path)
    finally:
        document.close()


def test_detects_image_content_instead_of_extension(tmp_path: Path) -> None:
    path = tmp_path / "not-really-a-pdf.pdf"
    make_image(path)

    assert detect_document_kind(path) == "png"


def test_reencodes_and_downsizes_image(tmp_path: Path) -> None:
    path = tmp_path / "large.png"
    make_image(path, size=(3000, 1000))

    document = prepare_document(
        path,
        delete_source=False,
        enforce_temp_location=False,
    )

    assert document.kind == "png"
    assert document.total_pages == 1
    assert document.pages[0].width == MAX_IMAGE_EDGE
    assert document.pages[0].height < MAX_IMAGE_EDGE
    assert document.pages[0].png.startswith(b"\x89PNG\r\n\x1a\n")


def test_caps_pdf_at_ten_pages(tmp_path: Path) -> None:
    path = tmp_path / "eleven-pages.pdf"
    make_pdf(path, MAX_PAGES + 1)

    document = prepare_document(
        path,
        delete_source=False,
        enforce_temp_location=False,
    )

    assert document.total_pages == MAX_PAGES + 1
    assert document.processed_pages == MAX_PAGES
    assert [page.number for page in document.pages] == list(
        range(1, MAX_PAGES + 1)
    )


def test_deletes_source_after_preparation(tmp_path: Path) -> None:
    path = tmp_path / "upload.png"
    make_image(path)

    prepare_document(
        path,
        delete_source=True,
        enforce_temp_location=False,
    )

    assert not path.exists()


def test_rejects_invalid_and_oversized_documents(tmp_path: Path) -> None:
    invalid = tmp_path / "invalid.pdf"
    invalid.write_bytes(b"not a document")
    with pytest.raises(DocumentError, match="valid PDF"):
        prepare_document(
            invalid,
            delete_source=False,
            enforce_temp_location=False,
        )

    oversized = tmp_path / "oversized.pdf"
    with oversized.open("wb") as handle:
        handle.truncate(MAX_DOCUMENT_BYTES + 1)
    with pytest.raises(DocumentError, match="20 MB"):
        prepare_document(
            oversized,
            delete_source=False,
            enforce_temp_location=False,
        )


def test_rejects_non_temporary_path_when_enforced(tmp_path: Path) -> None:
    path = tmp_path / "upload.png"
    make_image(path)

    # pytest's tmp_path is already below the system temp directory and is valid.
    document = prepare_document(
        path,
        delete_source=False,
        enforce_temp_location=True,
    )
    assert document.processed_pages == 1