225 lines
7.8 KiB
Python
225 lines
7.8 KiB
Python
"""End-to-end verification tests for unwrap-clothes."""
|
|
|
|
import http.server
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
|
|
import base64
|
|
import io
|
|
from unittest.mock import patch
|
|
|
|
from fastapi.testclient import TestClient
|
|
from PIL import Image
|
|
|
|
from app import app
|
|
|
|
|
|
def _create_synthetic_image(width: int, height: int, color: str = "red") -> bytes:
|
|
"""Helper to generate a JPEG image with exact dimensions."""
|
|
buf = io.BytesIO()
|
|
img = Image.new("RGB", (width, height), color=color)
|
|
img.save(buf, format="JPEG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def test_e2e_root_serves_html_with_title_and_dropzone():
|
|
"""GET / serves HTML containing web application title and dropzone."""
|
|
client = TestClient(app)
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert "text/html" in response.headers.get("content-type", "")
|
|
|
|
html_lower = response.text.lower()
|
|
# Verify web application title
|
|
assert "unwrap-clothes" in html_lower
|
|
# Verify dropzone element exists in HTML
|
|
assert "drop-zone" in html_lower or "dropzone" in html_lower
|
|
|
|
|
|
def test_e2e_api_clean_end_to_end():
|
|
"""Test FastAPI POST /api/clean end-to-end with high-res image and Lanczos restoration."""
|
|
client = TestClient(app)
|
|
|
|
orig_width, orig_height = 1200, 1600
|
|
model_width, model_height = 600, 800
|
|
|
|
# Create high-resolution synthetic image (1200x1600)
|
|
high_res_bytes = _create_synthetic_image(orig_width, orig_height, color="navy")
|
|
|
|
# Create lower-resolution synthetic image representing model output (600x800)
|
|
low_res_bytes = _create_synthetic_image(model_width, model_height, color="beige")
|
|
b64_low_res = base64.b64encode(low_res_bytes).decode()
|
|
|
|
mock_openrouter_response = {
|
|
"data": [{"b64_json": b64_low_res, "media_type": "image/jpeg"}],
|
|
"usage": {"cost": 0.0125},
|
|
}
|
|
|
|
class MockResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return mock_openrouter_response
|
|
|
|
files = {"file": ("high_res_garment.jpg", high_res_bytes, "image/jpeg")}
|
|
form_data = {
|
|
"prompt": "Soft warm beige studio background, evenly lit, no shadows, no props.",
|
|
"api_key": "sk-or-v1-fake-e2e-key",
|
|
"restore_res": "true",
|
|
}
|
|
|
|
with patch("cleaner.requests.post", return_value=MockResponse()) as mock_post:
|
|
response = client.post("/api/clean", files=files, data=form_data)
|
|
|
|
assert response.status_code == 200
|
|
mock_post.assert_called_once()
|
|
|
|
payload = response.json()
|
|
|
|
# Verify was_rescaled is True
|
|
assert payload["was_rescaled"] is True
|
|
# Verify original_dimensions and model_dimensions match expectations
|
|
assert payload["original_dimensions"] == [orig_width, orig_height]
|
|
assert payload["model_dimensions"] == [model_width, model_height]
|
|
assert payload["width"] == orig_width
|
|
assert payload["height"] == orig_height
|
|
# Verify cost is reported
|
|
assert payload["cost"] == 0.0125
|
|
|
|
# Verify decoded output image matches exact source dimensions thanks to Lanczos restoration
|
|
image_data_uri = payload["image"]
|
|
assert image_data_uri.startswith("data:image/jpeg;base64,")
|
|
encoded_data = image_data_uri.split(",", 1)[1]
|
|
decoded_bytes = base64.b64decode(encoded_data)
|
|
|
|
with Image.open(io.BytesIO(decoded_bytes)) as result_img:
|
|
assert result_img.size == (orig_width, orig_height)
|
|
|
|
|
|
def test_e2e_api_clean_with_custom_prompt_and_1500x2000():
|
|
"""Test FastAPI POST /api/clean with custom prompt and 1500x2000 dimensions."""
|
|
client = TestClient(app)
|
|
|
|
orig_width, orig_height = 1500, 2000
|
|
model_width, model_height = 750, 1000
|
|
|
|
high_res_bytes = _create_synthetic_image(orig_width, orig_height, color="darkgreen")
|
|
low_res_bytes = _create_synthetic_image(model_width, model_height, color="white")
|
|
b64_low_res = base64.b64encode(low_res_bytes).decode()
|
|
|
|
mock_openrouter_response = {
|
|
"data": [{"b64_json": b64_low_res, "media_type": "image/jpeg"}],
|
|
"usage": {"cost": 0.015},
|
|
}
|
|
|
|
class MockResponse:
|
|
status_code = 200
|
|
|
|
def json(self):
|
|
return mock_openrouter_response
|
|
|
|
files = {"file": ("catalog_garment.jpg", high_res_bytes, "image/jpeg")}
|
|
custom_prompt = "Minimalist bright studio setting with directional softbox lighting."
|
|
form_data = {
|
|
"prompt": custom_prompt,
|
|
"api_key": "sk-or-v1-custom-prompt-key",
|
|
"restore_res": "true",
|
|
}
|
|
|
|
with patch("cleaner.requests.post", return_value=MockResponse()) as mock_post:
|
|
response = client.post("/api/clean", files=files, data=form_data)
|
|
|
|
assert response.status_code == 200
|
|
mock_post.assert_called_once()
|
|
|
|
payload = response.json()
|
|
assert payload["was_rescaled"] is True
|
|
assert payload["original_dimensions"] == [orig_width, orig_height]
|
|
assert payload["model_dimensions"] == [model_width, model_height]
|
|
assert payload["width"] == orig_width
|
|
assert payload["height"] == orig_height
|
|
assert payload["cost"] == 0.015
|
|
|
|
encoded_data = payload["image"].split(",", 1)[1]
|
|
decoded_bytes = base64.b64decode(encoded_data)
|
|
with Image.open(io.BytesIO(decoded_bytes)) as result_img:
|
|
assert result_img.size == (orig_width, orig_height)
|
|
|
|
|
|
def test_e2e_cli_subprocess_end_to_end(tmp_path):
|
|
"""Test CLI unwrap_clothes.py end-to-end via subprocess with mocked OpenRouter."""
|
|
orig_width, orig_height = 1200, 1600
|
|
model_width, model_height = 600, 800
|
|
|
|
# Create source image in temporary directory
|
|
src_image_path = tmp_path / "garment.jpg"
|
|
src_bytes = _create_synthetic_image(orig_width, orig_height, color="purple")
|
|
src_image_path.write_bytes(src_bytes)
|
|
|
|
# Expected output path: <image_path>_clean.jpg
|
|
expected_clean_path = tmp_path / "garment_clean.jpg"
|
|
|
|
# Create mock lower-resolution image
|
|
low_res_bytes = _create_synthetic_image(model_width, model_height, color="white")
|
|
b64_low_res = base64.b64encode(low_res_bytes).decode()
|
|
|
|
mock_payload = {
|
|
"data": [{"b64_json": b64_low_res, "media_type": "image/jpeg"}],
|
|
"usage": {"cost": 0.02},
|
|
}
|
|
|
|
class MockOpenRouterHandler(http.server.BaseHTTPRequestHandler):
|
|
def do_POST(self):
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
_ = self.rfile.read(length)
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.end_headers()
|
|
self.wfile.write(json.dumps(mock_payload).encode())
|
|
|
|
def log_message(self, format, *args):
|
|
pass
|
|
|
|
server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), MockOpenRouterHandler)
|
|
server_thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
server_thread.start()
|
|
|
|
repo_root = Path(__file__).resolve().parent.parent
|
|
env = {
|
|
**os.environ,
|
|
"OPENROUTER_API_URL": f"http://127.0.0.1:{server.server_port}/api/v1/images",
|
|
}
|
|
|
|
try:
|
|
cmd = [
|
|
sys.executable,
|
|
"unwrap_clothes.py",
|
|
str(src_image_path),
|
|
"--bg",
|
|
"white",
|
|
"--api-key",
|
|
"fake-key",
|
|
]
|
|
proc = subprocess.run(
|
|
cmd,
|
|
cwd=str(repo_root),
|
|
env=env,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=15,
|
|
)
|
|
|
|
# Verify exit code 0
|
|
assert proc.returncode == 0, f"CLI failed with stderr: {proc.stderr}\nstdout: {proc.stdout}"
|
|
|
|
# Verify output file <image_path>_clean.jpg is created with exact source dimensions
|
|
assert expected_clean_path.is_file(), f"Output file {expected_clean_path} was not created"
|
|
with Image.open(expected_clean_path) as out_img:
|
|
assert out_img.size == (orig_width, orig_height)
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|