unwrap-clothes/tests/test_e2e.py

"""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


import pytest

@pytest.mark.parametrize(
    "orig_dims,model_dims,prompt,cost,color",
    [
        ((1200, 1600), (600, 800), "Soft warm beige studio background, evenly lit, no shadows, no props.", 0.0125, "navy"),
        ((1500, 2000), (750, 1000), "Minimalist bright studio setting with directional softbox lighting.", 0.015, "darkgreen"),
    ],
)
def test_e2e_api_clean_end_to_end(orig_dims, model_dims, prompt, cost, color):
    """Test FastAPI POST /api/clean end-to-end with high-res images, Lanczos restoration, and custom prompts."""
    client = TestClient(app)
    orig_width, orig_height = orig_dims
    model_width, model_height = model_dims

    high_res_bytes = _create_synthetic_image(orig_width, orig_height, color=color)
    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": cost},
    }

    class MockResponse:
        status_code = 200

        def json(self):
            return mock_openrouter_response

    files = {"file": ("garment.jpg", high_res_bytes, "image/jpeg")}
    form_data = {
        "prompt": prompt,
        "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()
        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"] == cost

        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}"
        assert "Restored resolution" in proc.stdout
        assert "cost $0.02" in proc.stdout
        assert "Saved:" in proc.stdout

        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()