477 lines
18 KiB
Python
Executable File
477 lines
18 KiB
Python
Executable File
#!/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:]]))
|