287 lines
8.9 KiB
Python
287 lines
8.9 KiB
Python
import pytest
|
|
from cleaner import (
|
|
API_URL,
|
|
DEFAULT_MODEL,
|
|
BACKGROUND_PRESETS,
|
|
EXTENSIONS,
|
|
build_prompt,
|
|
encode_image,
|
|
encode_image_bytes,
|
|
restore_resolution,
|
|
restore_resolution_bytes,
|
|
call_openrouter_images,
|
|
clean_garment,
|
|
CleanResult,
|
|
CleanerError,
|
|
)
|
|
import io
|
|
from PIL import Image
|
|
|
|
|
|
def test_constants():
|
|
assert API_URL == "https://openrouter.ai/api/v1/images"
|
|
assert "beige" in BACKGROUND_PRESETS
|
|
assert "white" in BACKGROUND_PRESETS
|
|
assert EXTENSIONS["image/jpeg"] == ".jpg"
|
|
assert EXTENSIONS["image/png"] == ".png"
|
|
assert EXTENSIONS["image/webp"] == ".webp"
|
|
|
|
|
|
def test_build_prompt_presets():
|
|
prompt_beige = build_prompt("beige")
|
|
assert BACKGROUND_PRESETS["beige"] in prompt_beige
|
|
assert "Remove ALL wrinkles" in prompt_beige
|
|
|
|
prompt_white = build_prompt("white")
|
|
assert BACKGROUND_PRESETS["white"] in prompt_white
|
|
assert "Remove ALL wrinkles" in prompt_white
|
|
|
|
# Case insensitivity
|
|
assert build_prompt("BEIGE") == prompt_beige
|
|
|
|
|
|
def test_build_prompt_custom():
|
|
custom = "Vintage wooden floor with warm sidelight."
|
|
prompt_custom = build_prompt(custom)
|
|
assert custom in prompt_custom
|
|
assert "Remove ALL wrinkles" in prompt_custom
|
|
|
|
|
|
def test_encode_image_bytes():
|
|
raw = b"fake-image-bytes"
|
|
encoded = encode_image_bytes(raw, "image/jpeg")
|
|
assert encoded == "data:image/jpeg;base64,ZmFrZS1pbWFnZS1ieXRlcw=="
|
|
|
|
# Default mime
|
|
encoded_default = encode_image_bytes(raw)
|
|
assert encoded_default.startswith("data:image/jpeg;base64,")
|
|
|
|
|
|
def test_encode_image_file(tmp_path):
|
|
img_file = tmp_path / "test.png"
|
|
img_file.write_bytes(b"png-data")
|
|
encoded = encode_image(str(img_file))
|
|
assert encoded == "data:image/png;base64,cG5nLWRhdGE="
|
|
|
|
jpg_file = tmp_path / "test.jpg"
|
|
jpg_file.write_bytes(b"jpg-data")
|
|
encoded_jpg = encode_image(str(jpg_file))
|
|
assert encoded_jpg == "data:image/jpeg;base64,anBnLWRhdGE="
|
|
|
|
unknown_file = tmp_path / "test.customext"
|
|
unknown_file.write_bytes(b"custom-data")
|
|
encoded_unknown = encode_image(str(unknown_file))
|
|
assert encoded_unknown.startswith("data:image/png;base64,")
|
|
|
|
|
|
def create_test_image(size=(100, 100), fmt="JPEG"):
|
|
buf = io.BytesIO()
|
|
img = Image.new("RGB", size, color="blue")
|
|
img.save(buf, format=fmt)
|
|
return buf.getvalue()
|
|
|
|
|
|
def test_restore_resolution_bytes_different_dimensions():
|
|
gen_data = create_test_image(size=(100, 100), fmt="JPEG")
|
|
orig_width, orig_height = 200, 300
|
|
|
|
restored = restore_resolution_bytes(gen_data, orig_width, orig_height, fmt="JPEG")
|
|
assert restored != gen_data
|
|
|
|
with Image.open(io.BytesIO(restored)) as im:
|
|
assert im.size == (200, 300)
|
|
|
|
|
|
def test_restore_resolution_bytes_same_dimensions_noop():
|
|
gen_data = create_test_image(size=(100, 100), fmt="JPEG")
|
|
restored = restore_resolution_bytes(gen_data, 100, 100, fmt="JPEG")
|
|
assert restored is gen_data # Exact same object/bytes, no re-encoding
|
|
|
|
|
|
def test_restore_resolution_files(tmp_path):
|
|
orig_file = tmp_path / "orig.jpg"
|
|
out_file = tmp_path / "out.jpg"
|
|
|
|
orig_img = Image.new("RGB", (300, 400), color="red")
|
|
orig_img.save(str(orig_file), format="JPEG")
|
|
|
|
gen_img = Image.new("RGB", (150, 200), color="green")
|
|
gen_img.save(str(out_file), format="JPEG")
|
|
|
|
restore_resolution(str(out_file), str(orig_file))
|
|
|
|
with Image.open(str(out_file)) as im:
|
|
assert im.size == (300, 400)
|
|
|
|
|
|
def test_restore_resolution_files_same_dimensions(tmp_path):
|
|
orig_file = tmp_path / "orig.jpg"
|
|
out_file = tmp_path / "out.jpg"
|
|
|
|
orig_img = Image.new("RGB", (200, 200), color="red")
|
|
orig_img.save(str(orig_file), format="JPEG")
|
|
orig_img.save(str(out_file), format="JPEG")
|
|
|
|
mtime_before = out_file.stat().st_mtime_ns
|
|
restore_resolution(str(out_file), str(orig_file))
|
|
mtime_after = out_file.stat().st_mtime_ns
|
|
|
|
# File should not have been overwritten
|
|
assert mtime_before == mtime_after
|
|
with Image.open(str(out_file)) as im:
|
|
assert im.size == (200, 200)
|
|
|
|
|
|
def test_call_openrouter_images_success(monkeypatch):
|
|
def mock_post(url, headers, json, timeout):
|
|
assert url == API_URL
|
|
assert headers == {"Authorization": "Bearer test-key"}
|
|
assert json["model"] == DEFAULT_MODEL
|
|
assert json["prompt"] == "test prompt"
|
|
assert json["background"] == "opaque"
|
|
assert json["output_format"] == "jpeg"
|
|
assert json["input_references"][0]["image_url"]["url"] == "data:image/jpeg;base64,abc"
|
|
assert timeout == 300
|
|
|
|
class MockResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return {
|
|
"data": [{"b64_json": "ZGF0YQ==", "media_type": "image/jpeg"}],
|
|
"usage": {"cost": 0.05},
|
|
}
|
|
|
|
return MockResponse()
|
|
|
|
import requests
|
|
monkeypatch.setattr(requests, "post", mock_post)
|
|
|
|
result = call_openrouter_images(
|
|
image_url_or_data_uri="data:image/jpeg;base64,abc",
|
|
prompt="test prompt",
|
|
api_key="test-key",
|
|
)
|
|
assert result["data"][0]["b64_json"] == "ZGF0YQ=="
|
|
assert result["usage"]["cost"] == 0.05
|
|
|
|
|
|
def test_call_openrouter_images_error_json(monkeypatch):
|
|
def mock_post(url, headers, json, timeout):
|
|
class MockResponse:
|
|
status_code = 401
|
|
text = '{"error": {"message": "Invalid API key"}}'
|
|
|
|
def json(self):
|
|
return {"error": {"message": "Invalid API key"}}
|
|
|
|
return MockResponse()
|
|
|
|
import requests
|
|
monkeypatch.setattr(requests, "post", mock_post)
|
|
|
|
with pytest.raises(CleanerError) as exc_info:
|
|
call_openrouter_images(
|
|
image_url_or_data_uri="data:image/jpeg;base64,abc",
|
|
prompt="test prompt",
|
|
api_key="invalid-key",
|
|
)
|
|
assert "API error 401: Invalid API key" in str(exc_info.value)
|
|
|
|
|
|
def test_call_openrouter_images_error_text(monkeypatch):
|
|
def mock_post(url, headers, json, timeout):
|
|
class MockResponse:
|
|
status_code = 502
|
|
text = "Bad Gateway"
|
|
|
|
def json(self):
|
|
raise ValueError("Not JSON")
|
|
|
|
return MockResponse()
|
|
|
|
import requests
|
|
monkeypatch.setattr(requests, "post", mock_post)
|
|
|
|
with pytest.raises(CleanerError) as exc_info:
|
|
call_openrouter_images(
|
|
image_url_or_data_uri="data:image/jpeg;base64,abc",
|
|
prompt="test prompt",
|
|
api_key="test-key",
|
|
)
|
|
assert "API error 502: Bad Gateway" in str(exc_info.value)
|
|
|
|
|
|
def test_clean_garment_bytes(monkeypatch):
|
|
orig_bytes = create_test_image(size=(300, 400), fmt="JPEG")
|
|
generated_bytes = create_test_image(size=(150, 200), fmt="JPEG")
|
|
import base64
|
|
b64_gen = base64.b64encode(generated_bytes).decode()
|
|
|
|
def mock_call_openrouter(image_url_or_data_uri, prompt, api_key, model=DEFAULT_MODEL, api_url=API_URL, timeout=300):
|
|
assert image_url_or_data_uri.startswith("data:image/jpeg;base64,")
|
|
assert prompt == "test prompt"
|
|
assert api_key == "test-key"
|
|
return {
|
|
"data": [{"b64_json": b64_gen, "media_type": "image/jpeg"}],
|
|
"usage": {"cost": 0.02},
|
|
}
|
|
|
|
import cleaner
|
|
monkeypatch.setattr(cleaner, "call_openrouter_images", mock_call_openrouter)
|
|
|
|
# With restore_res=True
|
|
res = clean_garment(
|
|
image_data=orig_bytes,
|
|
prompt="test prompt",
|
|
api_key="test-key",
|
|
restore_res=True,
|
|
)
|
|
assert isinstance(res, CleanResult)
|
|
assert res.media_type == "image/jpeg"
|
|
assert res.cost == 0.02
|
|
assert res.width == 300
|
|
assert res.height == 400
|
|
with Image.open(io.BytesIO(res.image_bytes)) as im:
|
|
assert im.size == (300, 400)
|
|
|
|
# With restore_res=False
|
|
res_unscaled = clean_garment(
|
|
image_data=orig_bytes,
|
|
prompt="test prompt",
|
|
api_key="test-key",
|
|
restore_res=False,
|
|
)
|
|
assert res_unscaled.width == 150
|
|
assert res_unscaled.height == 200
|
|
with Image.open(io.BytesIO(res_unscaled.image_bytes)) as im:
|
|
assert im.size == (150, 200)
|
|
|
|
|
|
def test_clean_garment_file(tmp_path, monkeypatch):
|
|
orig_file = tmp_path / "garment.jpg"
|
|
orig_img = Image.new("RGB", (250, 350), color="yellow")
|
|
orig_img.save(str(orig_file), format="JPEG")
|
|
|
|
gen_bytes = create_test_image(size=(100, 100), fmt="JPEG")
|
|
import base64
|
|
b64_gen = base64.b64encode(gen_bytes).decode()
|
|
|
|
def mock_call_openrouter(image_url_or_data_uri, prompt, api_key, model=DEFAULT_MODEL, api_url=API_URL, timeout=300):
|
|
return {
|
|
"data": [{"b64_json": b64_gen, "media_type": "image/jpeg"}],
|
|
"usage": {"cost": 0.03},
|
|
}
|
|
|
|
import cleaner
|
|
monkeypatch.setattr(cleaner, "call_openrouter_images", mock_call_openrouter)
|
|
|
|
res = clean_garment(
|
|
image_data=str(orig_file),
|
|
prompt="test prompt",
|
|
api_key="test-key",
|
|
restore_res=True,
|
|
)
|
|
assert res.width == 250
|
|
assert res.height == 350
|
|
assert res.cost == 0.03
|