Files
playground/download_imageset_test.py
T

302 lines
11 KiB
Python

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["pytest", "httpx", "pillow", "click"]
# ///
import shutil
import subprocess
import sys
import zipfile
from pathlib import Path
from functools import partial
from unittest.mock import Mock
import pytest
from click.testing import CliRunner
from PIL import Image
import download_imageset
@pytest.mark.parametrize("extension", [".zip", ".ZIP"])
def test_process_zip_and_directory(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, extension: str
) -> None:
directory = tmp_path / "existing"
nested = directory / "nested"
nested.mkdir(parents=True)
image_path = nested / "photo.png"
Image.new(mode="RGB", size=(10, 10)).save(fp=image_path)
archive_path = tmp_path / f"gallery{extension}"
with zipfile.ZipFile(file=archive_path, mode="w") as archive:
archive.write(filename=image_path, arcname="album/photo.png")
trash = Mock()
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
result = CliRunner().invoke(
cli=download_imageset.cli,
args=["process", str(archive_path), str(directory)],
)
assert result.exit_code == 0, result.output
assert (tmp_path / "gallery" / "gallery__0000.png").is_file()
assert (directory / "existing__0000.png").is_file()
assert not (tmp_path / "gallery" / "album").exists()
assert not nested.exists()
trash.assert_called_once_with(file_paths=[archive_path])
@pytest.mark.parametrize("command", ["process", "gallery"])
@pytest.mark.parametrize("processing_fails", [False, True])
def test_process_archives_in_directory(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
command: str,
processing_fails: bool,
) -> None:
directory = tmp_path / "album"
directory.mkdir()
archive_path = directory / "download.ZIP"
image_path = tmp_path / "photo.png"
Image.new(mode="RGB", size=(10, 10)).save(fp=image_path)
with zipfile.ZipFile(file=archive_path, mode="w") as archive:
if processing_fails:
archive.writestr(zinfo_or_arcname="nested/photo.png", data=b"invalid image")
else:
archive.write(filename=image_path, arcname="nested/photo.png")
trash = Mock()
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
monkeypatch.setattr(download_imageset, "run_command", Mock())
args = (
["process", str(directory)]
if command == "process"
else ["gallery", "https://example.com/album", "--cwd", str(tmp_path)]
)
result = CliRunner().invoke(cli=download_imageset.cli, args=args)
if processing_fails:
assert result.exit_code == 0, result.output
assert archive_path.is_file()
assert (directory / "nested" / "photo.png").read_bytes() == b"invalid image"
trash.assert_not_called()
else:
assert result.exit_code == 0, result.output
assert (directory / "album__0000.png").is_file()
assert not (directory / "nested").exists()
trash.assert_called_once_with(file_paths=[archive_path])
@pytest.mark.parametrize("filename", ["missing", "unsupported.txt", "broken.zip"])
def test_process_skips_invalid_inputs(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, filename: str
) -> None:
source = tmp_path / filename
if filename != "missing":
source.write_text(data="not an archive")
trash = Mock()
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
result = CliRunner().invoke(
cli=download_imageset.cli, args=["process", str(source)]
)
assert result.exit_code == 0, result.output
trash.assert_not_called()
if filename != "missing":
assert source.read_text() == "not an archive"
assert not (tmp_path / "broken").exists()
def test_process_keeps_zip_when_processing_fails(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
archive_path = tmp_path / "gallery.zip"
with zipfile.ZipFile(file=archive_path, mode="w") as archive:
archive.writestr(zinfo_or_arcname="photo.png", data=b"invalid image")
trash = Mock()
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
result = CliRunner().invoke(
cli=download_imageset.cli, args=["process", str(archive_path)]
)
assert result.exit_code == 0, result.output
assert archive_path.is_file()
trash.assert_not_called()
assert "resizing failed" in caplog.text
def test_process_continues_after_resize_failure(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
directory = tmp_path / "album"
directory.mkdir()
failed_image = directory / "broken.png"
failed_image.write_bytes(data=b"broken image")
good_image = directory / "good.png"
Image.new(mode="RGB", size=(6000, 60)).save(fp=good_image)
def processor(image_path: Path, output_path: Path) -> None:
if image_path == failed_image:
output_path.write_bytes(data=b"partial JPEG")
raise OSError("cannot decode image")
download_imageset.resize_with_pillow(image_path=image_path, output_path=output_path)
trash = Mock()
monkeypatch.setattr(download_imageset, "detect_resizer", lambda: processor)
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
result = CliRunner().invoke(cli=download_imageset.cli, args=["process", str(directory)])
assert result.exit_code == 0, result.output
assert failed_image.read_bytes() == b"broken image"
assert not (directory / "resized" / "broken.resized.jpg").exists()
with Image.open(fp=directory / "album__0001.jpg") as resized:
assert resized.size == (5000, 50)
trash.assert_called_once_with(file_paths=[good_image])
assert f"skipping {failed_image}: resizing failed: cannot decode image" in caplog.text
@pytest.mark.parametrize("use_vips", [False, True])
@pytest.mark.parametrize(
("size", "mode", "orientation", "expected_size"),
[
((6000, 60), "RGB", 1, (5000, 50)),
((60, 6000), "RGB", 1, (50, 5000)),
((40, 20), "RGBA", 1, (40, 20)),
((40, 20), "RGB", 6, (20, 40)),
],
)
def test_resize_image(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
use_vips: bool,
size: tuple[int, int],
mode: str,
orientation: int,
expected_size: tuple[int, int],
) -> None:
vips_binary = shutil.which(cmd="vips") if use_vips else None
if use_vips and vips_binary is None:
pytest.skip("vips is not installed")
processor = (
partial(download_imageset.resize_with_vips, vips_binary=vips_binary)
if vips_binary
else download_imageset.resize_with_pillow
)
image_path = tmp_path / "source.png"
output_path = tmp_path / "resized.jpg"
with Image.new(mode=mode, size=size) as image:
exif = Image.Exif()
exif[274] = orientation
image.save(fp=image_path, exif=exif)
processor(image_path=image_path, output_path=output_path)
with Image.open(fp=output_path) as resized:
resized.load()
assert resized.format == "JPEG"
assert resized.mode == "RGB"
assert resized.size == expected_size
assert image_path.is_file()
@pytest.mark.parametrize(
"error", [OSError("vips unavailable"), subprocess.CalledProcessError(1, "vips")]
)
def test_resize_image_falls_back_when_vips_fails(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, error: Exception
) -> None:
image_path = tmp_path / "source.png"
output_path = tmp_path / "resized.jpg"
Image.new(mode="RGB", size=(6000, 60)).save(fp=image_path)
monkeypatch.setattr(download_imageset.subprocess, "run", Mock(side_effect=error))
download_imageset.resize_with_vips(
image_path=image_path, output_path=output_path, vips_binary="vips"
)
with Image.open(fp=output_path) as resized:
resized.load()
assert resized.size == (5000, 50)
@pytest.mark.parametrize("vips_binary", [None, "/usr/bin/vips"])
def test_process_selects_processor_once(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, vips_binary: str | None
) -> None:
directories = [tmp_path / "first", tmp_path / "second"]
for directory in directories:
directory.mkdir()
Image.new(mode="RGB", size=(10, 10)).save(fp=directory / "photo.png")
lookup = Mock(return_value=vips_binary)
vips_processor = Mock()
pillow_processor = Mock()
monkeypatch.setattr(download_imageset.shutil, "which", lookup)
monkeypatch.setattr(download_imageset, "resize_with_vips", vips_processor)
monkeypatch.setattr(download_imageset, "resize_with_pillow", pillow_processor)
monkeypatch.setattr(download_imageset, "_needs_resize", Mock(return_value=True))
monkeypatch.setattr(download_imageset, "move_to_trash", Mock())
result = CliRunner().invoke(
cli=download_imageset.cli, args=["process", *map(str, directories)]
)
assert result.exit_code == 0, result.output
lookup.assert_called_once_with(cmd="vips")
selected_processor = vips_processor if vips_binary else pillow_processor
unused_processor = pillow_processor if vips_binary else vips_processor
assert selected_processor.call_count == 2
unused_processor.assert_not_called()
for directory, call in zip(directories, selected_processor.call_args_list):
expected = {
"image_path": directory / "photo.png",
"output_path": directory / "resized" / "photo.resized.jpg",
}
if vips_binary:
expected["vips_binary"] = vips_binary
assert call.kwargs == expected
@pytest.mark.parametrize("command", ["process", "gallery"])
@pytest.mark.parametrize("existing_output", [None, b"", b"already resized"])
def test_resume_resizing(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
command: str,
existing_output: bytes | None,
) -> None:
directory = tmp_path / "album"
resized_dir = directory / "resized"
resized_dir.mkdir(parents=True)
image_path = directory / "photo.png"
Image.new(mode="RGB", size=(6000, 60)).save(fp=image_path)
output_path = resized_dir / "photo.resized.jpg"
if existing_output is not None:
output_path.write_bytes(data=existing_output)
processor = Mock(wraps=download_imageset.resize_with_pillow)
trash = Mock()
monkeypatch.setattr(download_imageset, "detect_resizer", lambda: processor)
monkeypatch.setattr(download_imageset, "move_to_trash", trash)
if command == "process":
args = ["process", str(directory)]
else:
args = ["gallery", "https://example.com/album", "--cwd", str(tmp_path)]
result = CliRunner().invoke(cli=download_imageset.cli, args=args)
assert result.exit_code == 0, result.output
if existing_output:
processor.assert_not_called()
assert (directory / "album__0001.jpg").read_bytes() == existing_output
else:
processor.assert_called_once_with(image_path=image_path, output_path=output_path)
with Image.open(fp=directory / "album__0001.jpg") as resized:
assert resized.size == (5000, 50)
trash.assert_called_once_with(file_paths=[image_path])
if __name__ == "__main__":
sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))