131 lines
4 KiB
Python
131 lines
4 KiB
Python
|
import io
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
import unwrap_clothes
|
||
|
|
|
||
|
|
|
||
|
|
def test_reexported_symbols():
|
||
|
|
# Ensure backwards compatibility for any caller importing from unwrap_clothes
|
||
|
|
assert hasattr(unwrap_clothes, "API_URL")
|
||
|
|
assert hasattr(unwrap_clothes, "DEFAULT_MODEL")
|
||
|
|
assert hasattr(unwrap_clothes, "BACKGROUND_PRESETS")
|
||
|
|
assert hasattr(unwrap_clothes, "EXTENSIONS")
|
||
|
|
assert hasattr(unwrap_clothes, "build_prompt")
|
||
|
|
assert hasattr(unwrap_clothes, "encode_image")
|
||
|
|
assert hasattr(unwrap_clothes, "restore_resolution")
|
||
|
|
|
||
|
|
|
||
|
|
def test_cli_help():
|
||
|
|
result = subprocess.run(
|
||
|
|
[sys.executable, "unwrap_clothes.py", "--help"],
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
)
|
||
|
|
assert result.returncode == 0
|
||
|
|
assert "Unwrinkle clothing photos and unify backgrounds via OpenRouter." in result.stdout
|
||
|
|
assert "--bg" in result.stdout
|
||
|
|
assert "--model" in result.stdout
|
||
|
|
assert "--api-key" in result.stdout
|
||
|
|
assert "--no-restore-res" in result.stdout
|
||
|
|
|
||
|
|
|
||
|
|
def test_main_file_not_found(monkeypatch, capsys):
|
||
|
|
monkeypatch.setattr(sys, "argv", ["unwrap_clothes.py", "non_existent_file.jpg", "--api-key", "key"])
|
||
|
|
with pytest.raises(SystemExit) as exc:
|
||
|
|
unwrap_clothes.main()
|
||
|
|
assert "Error: file not found: non_existent_file.jpg" in str(exc.value)
|
||
|
|
|
||
|
|
|
||
|
|
def test_main_missing_api_key(tmp_path, monkeypatch, capsys):
|
||
|
|
img = tmp_path / "img.jpg"
|
||
|
|
img.write_bytes(b"data")
|
||
|
|
monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
|
||
|
|
monkeypatch.setattr(sys, "argv", ["unwrap_clothes.py", str(img)])
|
||
|
|
with pytest.raises(SystemExit) as exc:
|
||
|
|
unwrap_clothes.main()
|
||
|
|
assert "missing API key" in str(exc.value)
|
||
|
|
|
||
|
|
|
||
|
|
def test_main_success_flow(tmp_path, monkeypatch, capsys):
|
||
|
|
# Setup test image
|
||
|
|
src_file = tmp_path / "source.jpg"
|
||
|
|
im = Image.new("RGB", (200, 300), color="blue")
|
||
|
|
im.save(str(src_file), format="JPEG")
|
||
|
|
|
||
|
|
out_file = tmp_path / "custom_out.jpg"
|
||
|
|
|
||
|
|
from cleaner import CleanResult
|
||
|
|
|
||
|
|
cleaned_img = Image.new("RGB", (200, 300), color="white")
|
||
|
|
buf = io.BytesIO()
|
||
|
|
cleaned_img.save(buf, format="JPEG")
|
||
|
|
fake_result = CleanResult(
|
||
|
|
image_bytes=buf.getvalue(),
|
||
|
|
media_type="image/jpeg",
|
||
|
|
cost=0.04,
|
||
|
|
width=200,
|
||
|
|
height=300,
|
||
|
|
)
|
||
|
|
|
||
|
|
clean_garment_called = {}
|
||
|
|
|
||
|
|
def mock_clean_garment(image_data, prompt, api_key, model, restore_res):
|
||
|
|
clean_garment_called["image_data"] = image_data
|
||
|
|
clean_garment_called["prompt"] = prompt
|
||
|
|
clean_garment_called["api_key"] = api_key
|
||
|
|
clean_garment_called["model"] = model
|
||
|
|
clean_garment_called["restore_res"] = restore_res
|
||
|
|
return fake_result
|
||
|
|
|
||
|
|
monkeypatch.setattr("unwrap_clothes.clean_garment", mock_clean_garment)
|
||
|
|
monkeypatch.setattr(
|
||
|
|
sys,
|
||
|
|
"argv",
|
||
|
|
[
|
||
|
|
"unwrap_clothes.py",
|
||
|
|
str(src_file),
|
||
|
|
"-o",
|
||
|
|
str(out_file),
|
||
|
|
"--bg",
|
||
|
|
"beige",
|
||
|
|
"--api-key",
|
||
|
|
"test-token",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
|
||
|
|
unwrap_clothes.main()
|
||
|
|
|
||
|
|
assert clean_garment_called["api_key"] == "test-token"
|
||
|
|
assert clean_garment_called["restore_res"] is True
|
||
|
|
assert out_file.exists()
|
||
|
|
assert out_file.read_bytes() == fake_result.image_bytes
|
||
|
|
|
||
|
|
captured = capsys.readouterr()
|
||
|
|
assert f"Submitting to {unwrap_clothes.DEFAULT_MODEL}..." in captured.out
|
||
|
|
assert f"Saved: {out_file} (cost $0.04)" in captured.out
|
||
|
|
|
||
|
|
|
||
|
|
def test_main_api_error_handling(tmp_path, monkeypatch):
|
||
|
|
src_file = tmp_path / "source.jpg"
|
||
|
|
src_file.write_bytes(b"dummy")
|
||
|
|
|
||
|
|
from cleaner import CleanerError
|
||
|
|
|
||
|
|
def mock_clean_garment(*args, **kwargs):
|
||
|
|
raise CleanerError("API error 400: Invalid image format")
|
||
|
|
|
||
|
|
monkeypatch.setattr("unwrap_clothes.clean_garment", mock_clean_garment)
|
||
|
|
monkeypatch.setattr(
|
||
|
|
sys,
|
||
|
|
"argv",
|
||
|
|
["unwrap_clothes.py", str(src_file), "--api-key", "token"],
|
||
|
|
)
|
||
|
|
|
||
|
|
with pytest.raises(SystemExit) as exc:
|
||
|
|
unwrap_clothes.main()
|
||
|
|
assert "API error 400: Invalid image format" in str(exc.value)
|