#!/usr/bin/env -S uv run --script # /// script # requires-python = ">=3.11" # dependencies = ["pytest", "httpx", "pillow", "click"] # /// import shutil import subprocess import sys import threading import urllib.parse import zipfile from concurrent.futures import ThreadPoolExecutor 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]) @pytest.mark.parametrize("extension", [".zip", ".mp4"]) @pytest.mark.parametrize("unfinished", [False, True]) def test_filehost_resumes_unfinished_downloads( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, extension: str, unfinished: bool, ) -> None: directory = tmp_path / "album" directory.mkdir() destination = ( directory / "download.zip" if extension == ".zip" else tmp_path / "album.mp4" ) with destination.open(mode="wb") as output: output.truncate(6 * 1024 * 1024) control_file = destination.with_name(f"{destination.name}.aria2") if unfinished: control_file.touch() url = "https://example.com/download" monkeypatch.setattr( download_imageset, "resolve_download_url_with_alldebrid", Mock(return_value=url) ) monkeypatch.setattr( download_imageset, "_resolve_url", Mock(return_value=download_imageset.UrlMeta( suggested_filename=f"download{extension}", suggested_extension=extension, size_bytes=destination.stat().st_size, )), ) run_command = Mock() monkeypatch.setattr(download_imageset, "_TRICKLEBAR_APP", tmp_path / "missing.app") monkeypatch.setattr(download_imageset, "run_command", run_command) monkeypatch.setattr(download_imageset, "process_imageset", Mock()) result = CliRunner().invoke( cli=download_imageset.cli, args=["filehost", url, "--name", "album", "--cwd", str(tmp_path), "--token", "test"], ) assert result.exit_code == 0, result.output if unfinished: assert f"unfinished download; resuming {destination}" in caplog.text run_command.assert_called_once_with(args=[ "aria2c", "--continue=true", f"--dir={destination.parent}", f"--out={destination.name}", url, ]) else: assert "unfinished download; resuming" not in caplog.text run_command.assert_not_called() def test_aria2_download_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: failure = subprocess.CalledProcessError(returncode=1, cmd="aria2c") monkeypatch.setattr(download_imageset, "run_command", Mock(side_effect=failure)) with pytest.raises(subprocess.CalledProcessError): download_imageset.download_url_with_aria2( url="https://example.com/archive.zip", destination=Path("archive.zip") ) @pytest.mark.parametrize( ("status", "file_state", "error"), [ ("complete", "valid", None), ("error", "valid", "download failed"), ("complete", "missing", "without a downloaded file"), ("complete", "wrong_path", "instead of"), ("complete", "unfinished", "download is unfinished"), ], ) def test_tricklebar_waits_for_callback( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, status: str, file_state: str, error: str | None, ) -> None: destination = tmp_path / "download & photos.zip" downloaded_path = tmp_path / "wrong.zip" if file_state == "wrong_path" else destination if file_state != "missing": downloaded_path.write_bytes(data=b"downloaded archive") if file_state == "unfinished": destination.with_name(f"{destination.name}.aria2").touch() url = "https://example.com/photos.zip?token=a+b&other=value" opened = threading.Event() callback_paths: list[Path] = [] def open_download(args: list[str]) -> None: assert args[0] == "open" parsed = urllib.parse.urlparse(url=args[1]) assert "+" not in parsed.query assert parsed.scheme == "tricklebar" assert parsed.netloc == "add-download" query = urllib.parse.parse_qs(qs=parsed.query) assert query["url"] == [url] assert query["dir"] == [str(destination.parent)] assert query["out"] == [destination.name] assert query["continue"] == ["true"] callback_paths.append(Path(query["on_done"][0])) opened.set() monkeypatch.setattr(download_imageset, "run_command", open_download) with ThreadPoolExecutor(max_workers=1) as executor: download = executor.submit( download_imageset.download_url_with_tricklebar, url=url, destination=destination ) assert opened.wait(timeout=5) assert not download.done() callback = callback_paths[0] subprocess.run(args=[str(callback), status, str(downloaded_path), "gid", url], check=True) if error is None: assert download.result(timeout=5) is None else: with pytest.raises(download_imageset.click.ClickException, match=error): download.result(timeout=5) assert not callback.parent.exists() @pytest.mark.parametrize("installed", [False, True]) def test_download_selects_backend( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, installed: bool ) -> None: application = tmp_path / "TrickleBar.app" if installed: application.mkdir() monkeypatch.setattr(download_imageset, "_TRICKLEBAR_APP", application) aria2 = Mock() tricklebar = Mock() monkeypatch.setattr(download_imageset, "download_url_with_aria2", aria2) monkeypatch.setattr(download_imageset, "download_url_with_tricklebar", tricklebar) destination = tmp_path / "photos" / "album.zip" url = "https://example.com/album.zip" download_imageset.download_url(url=url, destination=destination) assert destination.parent.is_dir() selected = tricklebar if installed else aria2 unused = aria2 if installed else tricklebar selected.assert_called_once_with(url=url, destination=destination.resolve()) unused.assert_not_called() @pytest.mark.parametrize("process_directory", [False, True]) def test_process_skips_unfinished_archive( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, process_directory: bool, ) -> 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"image data") archive_path.with_name(f"{archive_path.name}.aria2").touch() trash = Mock() monkeypatch.setattr(download_imageset, "move_to_trash", trash) result = CliRunner().invoke( cli=download_imageset.cli, args=["process", str(tmp_path if process_directory else archive_path)], ) assert result.exit_code == 0, result.output assert "download is unfinished" in caplog.text assert not (tmp_path / "photo.png").exists() assert not (tmp_path / "gallery").exists() assert archive_path.is_file() trash.assert_not_called() if __name__ == "__main__": sys.exit(pytest.main(args=[str(Path(__file__).resolve()), *sys.argv[1:]]))