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