402 lines
14 KiB
Python
402 lines
14 KiB
Python
# SPDX-FileCopyrightText: 2026 James R. Barlow
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""End-to-end tests for the batch web interface in ``webui/``.
|
|
|
|
These drive a real uvicorn server over real HTTP and run real OCR, so they are
|
|
slow. They are skipped unless the ``webui`` extra is installed::
|
|
|
|
uv sync --extra webui --group test
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import mimetypes
|
|
import socket
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
import zipfile
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("fastapi", reason="webui extra not installed")
|
|
pytest.importorskip("uvicorn", reason="webui extra not installed")
|
|
pytest.importorskip("multipart", reason="python-multipart not installed")
|
|
|
|
# webui/ lives at the repo root and is deliberately not part of the installed
|
|
# ocrmypdf distribution, so put the repo root on the path explicitly.
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
RESOURCES = Path(__file__).parent / "resources"
|
|
|
|
pytestmark = pytest.mark.slow
|
|
|
|
|
|
# --------------------------------------------------------------- HTTP helpers
|
|
|
|
|
|
def encode_multipart(
|
|
files: list[tuple[str, str, bytes]], fields: dict[str, str]
|
|
) -> tuple[bytes, str]:
|
|
"""Build a multipart/form-data body without pulling in a HTTP library."""
|
|
boundary = f"----ocrmypdf{uuid.uuid4().hex}"
|
|
buffer = BytesIO()
|
|
for name, value in fields.items():
|
|
buffer.write(f"--{boundary}\r\n".encode())
|
|
buffer.write(f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode())
|
|
buffer.write(value.encode() + b"\r\n")
|
|
for field_name, filename, content in files:
|
|
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
|
buffer.write(f"--{boundary}\r\n".encode())
|
|
buffer.write(
|
|
f'Content-Disposition: form-data; name="{field_name}"; '
|
|
f'filename="{filename}"\r\n'.encode()
|
|
)
|
|
buffer.write(f"Content-Type: {content_type}\r\n\r\n".encode())
|
|
buffer.write(content + b"\r\n")
|
|
buffer.write(f"--{boundary}--\r\n".encode())
|
|
return buffer.getvalue(), f"multipart/form-data; boundary={boundary}"
|
|
|
|
|
|
class Client:
|
|
"""The smallest HTTP client that can exercise the API."""
|
|
|
|
def __init__(self, base_url: str):
|
|
"""Bind this client to a running server."""
|
|
self.base_url = base_url
|
|
|
|
def request(self, method: str, path: str, data=None, content_type=None):
|
|
request = urllib.request.Request(self.base_url + path, data=data, method=method)
|
|
if content_type:
|
|
request.add_header("Content-Type", content_type)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
return response.status, response.read(), dict(response.headers)
|
|
except urllib.error.HTTPError as error:
|
|
return error.code, error.read(), dict(error.headers)
|
|
|
|
def get_json(self, path: str):
|
|
status, body, _ = self.request("GET", path)
|
|
assert status == 200, body
|
|
return json.loads(body)
|
|
|
|
def post_batch(self, files: list[tuple[str, bytes]], options: dict | None = None):
|
|
body, content_type = encode_multipart(
|
|
[("files", name, content) for name, content in files],
|
|
{"options": json.dumps(options or {})},
|
|
)
|
|
status, raw, _ = self.request("POST", "/api/batches", body, content_type)
|
|
return status, (json.loads(raw) if raw else None)
|
|
|
|
|
|
# ------------------------------------------------------------------- fixtures
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def server(tmp_path_factory):
|
|
"""Run the real ASGI app on a real port for the duration of the module."""
|
|
import uvicorn
|
|
|
|
work_dir = tmp_path_factory.mktemp("webui-work")
|
|
# A module-scoped fixture cannot use the function-scoped monkeypatch.
|
|
monkeypatch = pytest.MonkeyPatch()
|
|
monkeypatch.setenv("OCRMYPDF_WEBUI_WORK_DIR", str(work_dir))
|
|
monkeypatch.setenv("OCRMYPDF_WEBUI_MAX_UPLOAD_MB", "2")
|
|
monkeypatch.setenv("OCRMYPDF_WEBUI_MAX_FILES", "4")
|
|
monkeypatch.setenv("OCRMYPDF_WEBUI_WORKERS", "2")
|
|
monkeypatch.setenv("OCRMYPDF_WEBUI_OCR_JOBS", "1")
|
|
monkeypatch.setenv("OCRMYPDF_WEBUI_BATCH_TTL_SECONDS", "600")
|
|
|
|
from webui.config import get_settings, installed_languages
|
|
|
|
get_settings.cache_clear()
|
|
installed_languages.cache_clear()
|
|
|
|
from webui.app import app
|
|
|
|
with socket.socket() as probe:
|
|
probe.bind(("127.0.0.1", 0))
|
|
port = probe.getsockname()[1]
|
|
|
|
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
|
|
uv_server = uvicorn.Server(config)
|
|
thread = threading.Thread(target=uv_server.run, daemon=True)
|
|
thread.start()
|
|
|
|
deadline = time.monotonic() + 30
|
|
while not uv_server.started:
|
|
if time.monotonic() > deadline:
|
|
raise RuntimeError("uvicorn did not start")
|
|
time.sleep(0.05)
|
|
|
|
yield Client(f"http://127.0.0.1:{port}")
|
|
|
|
uv_server.should_exit = True
|
|
thread.join(timeout=30)
|
|
get_settings.cache_clear()
|
|
monkeypatch.undo()
|
|
|
|
|
|
def wait_for_batch(client: Client, batch_id: str, timeout: float = 600.0) -> dict:
|
|
"""Poll until every file in the batch reaches a terminal state."""
|
|
deadline = time.monotonic() + timeout
|
|
while True:
|
|
batch = client.get_json(f"/api/batches/{batch_id}")
|
|
if batch["finished"]:
|
|
return batch
|
|
if time.monotonic() > deadline:
|
|
pytest.fail(f"batch did not finish in {timeout}s: {batch}")
|
|
time.sleep(0.5)
|
|
|
|
|
|
def pdf_text(data: bytes) -> str:
|
|
"""Extract text from an in-memory PDF."""
|
|
from pdfminer.high_level import extract_text
|
|
|
|
return extract_text(BytesIO(data))
|
|
|
|
|
|
# ---------------------------------------------------------------------- tests
|
|
|
|
|
|
def test_config_advertises_capabilities(server):
|
|
config = server.get_json("/api/config")
|
|
assert "eng" in config["languages"]
|
|
assert ".pdf" in config["accepted_extensions"]
|
|
assert config["max_files"] == 4
|
|
assert config["max_upload_bytes"] == 2 * 1024 * 1024
|
|
assert config["defaults"]["mode"] == "skip-text"
|
|
|
|
|
|
def test_health(server):
|
|
status, body, _ = server.request("GET", "/healthz")
|
|
assert status == 200
|
|
assert body == b"ok"
|
|
|
|
|
|
def test_index_page_is_served(server):
|
|
status, body, headers = server.request("GET", "/")
|
|
assert status == 200
|
|
assert b"OCRmyPDF" in body
|
|
assert "text/html" in headers["content-type"]
|
|
|
|
|
|
def test_batch_of_several_files_ocrs_and_downloads(server):
|
|
"""The core flow: many files in, searchable PDFs out, plus a zip."""
|
|
inputs = [
|
|
("linn.png", (RESOURCES / "linn.png").read_bytes()),
|
|
("ccitt.pdf", (RESOURCES / "ccitt.pdf").read_bytes()),
|
|
("trivial.pdf", (RESOURCES / "trivial.pdf").read_bytes()),
|
|
]
|
|
status, batch = server.post_batch(
|
|
inputs, {"languages": ["eng"], "mode": "skip-text", "output_type": "pdf"}
|
|
)
|
|
assert status == 201, batch
|
|
assert batch["total"] == 3
|
|
assert {file["name"] for file in batch["files"]} == {
|
|
"linn.png",
|
|
"ccitt.pdf",
|
|
"trivial.pdf",
|
|
}
|
|
|
|
batch = wait_for_batch(server, batch["id"])
|
|
by_name = {file["name"]: file for file in batch["files"]}
|
|
|
|
# An image of text must come back as a PDF that actually contains that text.
|
|
linn = by_name["linn.png"]
|
|
assert linn["status"] == "succeeded", linn
|
|
assert linn["output_name"] == "linn.pdf"
|
|
status, pdf, headers = server.request(
|
|
"GET", f"/api/batches/{batch['id']}/files/{linn['id']}"
|
|
)
|
|
assert status == 200
|
|
assert pdf.startswith(b"%PDF")
|
|
assert "linn.pdf" in headers["content-disposition"]
|
|
assert "linnsequencer" in pdf_text(pdf).lower()
|
|
|
|
# A blank page is not an error; it just yields no text.
|
|
assert by_name["trivial.pdf"]["status"] == "succeeded"
|
|
assert by_name["ccitt.pdf"]["status"] == "succeeded"
|
|
assert batch["downloadable"] == 3
|
|
assert batch["completed"] == 3
|
|
|
|
# The zip must contain one entry per successful file.
|
|
status, archive, headers = server.request(
|
|
"GET", f"/api/batches/{batch['id']}/download"
|
|
)
|
|
assert status == 200
|
|
assert headers["content-type"] == "application/zip"
|
|
with zipfile.ZipFile(BytesIO(archive)) as zf:
|
|
assert sorted(zf.namelist()) == ["ccitt.pdf", "linn.pdf", "trivial.pdf"]
|
|
assert zf.read("linn.pdf").startswith(b"%PDF")
|
|
|
|
|
|
def test_duplicate_names_are_disambiguated_in_zip(server):
|
|
content = (RESOURCES / "trivial.pdf").read_bytes()
|
|
status, batch = server.post_batch(
|
|
[("same.pdf", content), ("same.pdf", content)], {"output_type": "pdf"}
|
|
)
|
|
assert status == 201
|
|
batch = wait_for_batch(server, batch["id"])
|
|
assert batch["downloadable"] == 2
|
|
_, archive, _ = server.request("GET", f"/api/batches/{batch['id']}/download")
|
|
with zipfile.ZipFile(BytesIO(archive)) as zf:
|
|
assert sorted(zf.namelist()) == ["same (1).pdf", "same.pdf"]
|
|
|
|
|
|
def test_file_with_existing_text_fails_helpfully_in_normal_mode(server):
|
|
"""A per-file failure must be reported per file, not as a batch error."""
|
|
status, batch = server.post_batch(
|
|
[
|
|
("graph_ocred.pdf", (RESOURCES / "graph_ocred.pdf").read_bytes()),
|
|
("trivial.pdf", (RESOURCES / "trivial.pdf").read_bytes()),
|
|
],
|
|
{"mode": "normal", "output_type": "pdf"},
|
|
)
|
|
assert status == 201
|
|
batch = wait_for_batch(server, batch["id"])
|
|
by_name = {file["name"]: file for file in batch["files"]}
|
|
|
|
failed = by_name["graph_ocred.pdf"]
|
|
assert failed["status"] == "failed"
|
|
assert "already contains text" in failed["error"]
|
|
# The healthy file in the same batch still succeeds and is downloadable.
|
|
assert by_name["trivial.pdf"]["status"] == "succeeded"
|
|
assert batch["downloadable"] == 1
|
|
|
|
status, _, _ = server.request(
|
|
"GET", f"/api/batches/{batch['id']}/files/{failed['id']}"
|
|
)
|
|
assert status == 409
|
|
|
|
status, log, _ = server.request(
|
|
"GET", f"/api/batches/{batch['id']}/files/{failed['id']}/log"
|
|
)
|
|
assert status == 200
|
|
assert len(log) > 0
|
|
|
|
|
|
def test_skip_text_mode_handles_the_same_file(server):
|
|
status, batch = server.post_batch(
|
|
[("graph_ocred.pdf", (RESOURCES / "graph_ocred.pdf").read_bytes())],
|
|
{"mode": "skip-text", "output_type": "pdf"},
|
|
)
|
|
assert status == 201
|
|
batch = wait_for_batch(server, batch["id"])
|
|
assert batch["files"][0]["status"] == "succeeded"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"options,expected_status",
|
|
[
|
|
({"languages": ["eng; rm -rf /"]}, 422), # command injection attempt
|
|
({"languages": ["klingon"]}, 422), # uninstalled language
|
|
({"optimize": 9}, 422), # out of range
|
|
({"mode": "--evil"}, 422), # not a member of the enum
|
|
({"image_dpi": -1}, 422), # out of range
|
|
({"unknown_option": True}, 422), # extra fields forbidden
|
|
],
|
|
)
|
|
def test_invalid_options_are_rejected(server, options, expected_status):
|
|
status, body = server.post_batch(
|
|
[("trivial.pdf", (RESOURCES / "trivial.pdf").read_bytes())], options
|
|
)
|
|
assert status == expected_status, body
|
|
|
|
|
|
def test_unsupported_file_type_is_rejected(server):
|
|
status, body = server.post_batch([("evil.exe", b"MZ" + b"\0" * 100)])
|
|
assert status == 415
|
|
assert "unsupported type" in body["detail"]
|
|
|
|
|
|
def test_oversized_file_is_rejected(server):
|
|
status, body = server.post_batch([("big.pdf", b"%PDF-1.7\n" + b"x" * 3_000_000)])
|
|
assert status == 413
|
|
assert "limit" in body["detail"]
|
|
|
|
|
|
def test_empty_file_is_rejected(server):
|
|
status, body = server.post_batch([("empty.pdf", b"")])
|
|
assert status == 400
|
|
|
|
|
|
def test_too_many_files_is_rejected(server):
|
|
content = (RESOURCES / "trivial.pdf").read_bytes()
|
|
status, body = server.post_batch(
|
|
[(f"file{n}.pdf", content) for n in range(5)] # limit is 4
|
|
)
|
|
assert status == 413
|
|
assert "Too many files" in body["detail"]
|
|
|
|
|
|
def test_path_traversal_in_filename_is_neutralized(server):
|
|
status, batch = server.post_batch(
|
|
[("../../etc/passwd.pdf", (RESOURCES / "trivial.pdf").read_bytes())],
|
|
{"output_type": "pdf"},
|
|
)
|
|
assert status == 201
|
|
assert batch["files"][0]["name"] == "passwd.pdf"
|
|
batch = wait_for_batch(server, batch["id"])
|
|
_, archive, _ = server.request("GET", f"/api/batches/{batch['id']}/download")
|
|
with zipfile.ZipFile(BytesIO(archive)) as zf:
|
|
assert zf.namelist() == ["passwd.pdf"]
|
|
|
|
|
|
def test_delete_removes_the_batch_and_its_files(server):
|
|
status, batch = server.post_batch(
|
|
[("trivial.pdf", (RESOURCES / "trivial.pdf").read_bytes())],
|
|
{"output_type": "pdf"},
|
|
)
|
|
assert status == 201
|
|
batch_id = batch["id"]
|
|
wait_for_batch(server, batch_id)
|
|
|
|
from webui.config import get_settings
|
|
|
|
directory = get_settings().work_dir / batch_id
|
|
assert directory.exists()
|
|
|
|
status, _, _ = server.request("DELETE", f"/api/batches/{batch_id}")
|
|
assert status == 204
|
|
assert not directory.exists()
|
|
|
|
status, _, _ = server.request("GET", f"/api/batches/{batch_id}")
|
|
assert status == 404
|
|
|
|
|
|
def test_unknown_batch_returns_404(server):
|
|
status, _, _ = server.request("GET", "/api/batches/" + "0" * 32)
|
|
assert status == 404
|
|
|
|
|
|
def test_sweep_removes_expired_batches(server):
|
|
"""The TTL reaper must delete batches whether or not they were downloaded."""
|
|
status, batch = server.post_batch(
|
|
[("trivial.pdf", (RESOURCES / "trivial.pdf").read_bytes())],
|
|
{"output_type": "pdf"},
|
|
)
|
|
assert status == 201
|
|
batch_id = batch["id"]
|
|
wait_for_batch(server, batch_id)
|
|
|
|
from webui.app import app
|
|
|
|
manager = app.state.manager
|
|
directory = manager.settings.work_dir / batch_id
|
|
assert directory.exists()
|
|
|
|
# Pretend enough time has passed for the TTL to lapse.
|
|
manager.sweep(now=time.time() + manager.settings.batch_ttl_seconds + 1)
|
|
assert not directory.exists()
|
|
assert manager.get_batch(batch_id) is None
|