unwrap-clothes/tests/test_cleaner.py

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

def test_resolve_api_key(monkeypatch):
    import cleaner
    monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
    assert cleaner.resolve_api_key(None) is None
    assert cleaner.resolve_api_key("   ") is None
    assert cleaner.resolve_api_key("custom-key") == "custom-key"
    monkeypatch.setenv("OPENROUTER_API_KEY", "env-key")
    assert cleaner.resolve_api_key(None) == "env-key"
    assert cleaner.resolve_api_key("custom-key") == "custom-key"


def test_clean_result_to_payload():
    raw = b"img-data"
    res = CleanResult(
        image_bytes=raw,
        media_type="image/jpeg",
        cost=0.015,
        width=100,
        height=200,
        original_dimensions=(100, 200),
        model_dimensions=(50, 100),
    )
    payload = res.to_payload()
    assert payload["cost"] == 0.015
    assert payload["width"] == 100
    assert payload["height"] == 200
    assert payload["original_dimensions"] == [100, 200]
    assert payload["model_dimensions"] == [50, 100]
    assert payload["was_rescaled"] is True
    assert payload["media_type"] == "image/jpeg"
    assert payload["image"].startswith("data:image/jpeg;base64,")