217 lines
7.4 KiB
Python
217 lines
7.4 KiB
Python
"""Tests for FastAPI backend in app.py."""
|
|
|
|
import base64
|
|
import io
|
|
from unittest.mock import patch
|
|
|
|
from fastapi.testclient import TestClient
|
|
from PIL import Image
|
|
|
|
from cleaner import CleanerError, CleanResult
|
|
|
|
|
|
def test_health_check():
|
|
"""GET /api/health returns 200 OK with status ok."""
|
|
from app import app
|
|
|
|
client = TestClient(app)
|
|
response = client.get("/api/health")
|
|
assert response.status_code == 200
|
|
assert response.json() == {"status": "ok"}
|
|
|
|
|
|
def test_root_fallback_when_no_index_html(monkeypatch, tmp_path):
|
|
"""GET / returns a graceful fallback when static/index.html does not exist."""
|
|
import app as app_module
|
|
|
|
# Point STATIC_DIR to a temporary directory without index.html
|
|
monkeypatch.setattr(app_module, "STATIC_DIR", tmp_path)
|
|
client = TestClient(app_module.app)
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert response.json()["status"] == "ok"
|
|
|
|
|
|
def test_root_serves_index_html_when_present(monkeypatch, tmp_path):
|
|
"""GET / serves static/index.html when it exists."""
|
|
import app as app_module
|
|
|
|
index_file = tmp_path / "index.html"
|
|
index_file.write_text("<!DOCTYPE html><html><body>Unwrap Clothes</body></html>")
|
|
monkeypatch.setattr(app_module, "STATIC_DIR", tmp_path)
|
|
|
|
client = TestClient(app_module.app)
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert "Unwrap Clothes" in response.text
|
|
assert "text/html" in response.headers.get("content-type", "")
|
|
|
|
|
|
def _create_test_image_bytes() -> bytes:
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (50, 50), color="blue").save(buf, format="JPEG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def test_clean_missing_api_key_returns_400(monkeypatch):
|
|
"""POST /api/clean without API key when OPENROUTER_API_KEY is unset returns 400."""
|
|
from app import app
|
|
|
|
monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
|
|
client = TestClient(app)
|
|
image_bytes = _create_test_image_bytes()
|
|
files = {"file": ("test.jpg", image_bytes, "image/jpeg")}
|
|
|
|
response = client.post("/api/clean", files=files)
|
|
assert response.status_code == 400
|
|
assert "API key" in response.json()["detail"]
|
|
|
|
|
|
def test_clean_uses_env_api_key_when_form_omitted(monkeypatch):
|
|
"""POST /api/clean uses OPENROUTER_API_KEY from environment if api_key form field is omitted."""
|
|
from app import app
|
|
|
|
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-v1-env-key")
|
|
client = TestClient(app)
|
|
image_bytes = _create_test_image_bytes()
|
|
files = {"file": ("test.jpg", image_bytes, "image/jpeg")}
|
|
|
|
mock_result = CleanResult(
|
|
image_bytes=image_bytes,
|
|
media_type="image/jpeg",
|
|
cost=0.005,
|
|
width=50,
|
|
height=50,
|
|
original_dimensions=(50, 50),
|
|
model_dimensions=(50, 50),
|
|
)
|
|
with patch("app.clean_garment", return_value=mock_result) as mock_clean:
|
|
response = client.post("/api/clean", files=files)
|
|
assert response.status_code == 200
|
|
mock_clean.assert_called_once()
|
|
_, kwargs = mock_clean.call_args
|
|
assert kwargs["api_key"] == "sk-or-v1-env-key"
|
|
|
|
|
|
def test_clean_empty_file_returns_400(monkeypatch):
|
|
"""POST /api/clean with empty file returns 400 Bad Request."""
|
|
from app import app
|
|
|
|
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-v1-env-key")
|
|
client = TestClient(app)
|
|
files = {"file": ("empty.jpg", b"", "image/jpeg")}
|
|
|
|
response = client.post("/api/clean", files=files)
|
|
assert response.status_code == 400
|
|
assert "empty" in response.json()["detail"].lower()
|
|
|
|
|
|
def test_clean_invalid_image_file_returns_400(monkeypatch):
|
|
"""POST /api/clean with invalid image content returns 400 Bad Request."""
|
|
from app import app
|
|
|
|
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-v1-env-key")
|
|
client = TestClient(app)
|
|
files = {"file": ("corrupt.jpg", b"not-a-valid-image-content", "image/jpeg")}
|
|
|
|
response = client.post("/api/clean", files=files)
|
|
assert response.status_code == 400
|
|
assert "invalid" in response.json()["detail"].lower()
|
|
|
|
|
|
def test_clean_cleaner_error_returns_502(monkeypatch):
|
|
"""POST /api/clean when clean_garment raises CleanerError returns 502 Bad Gateway."""
|
|
from app import app
|
|
|
|
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-v1-env-key")
|
|
client = TestClient(app)
|
|
image_bytes = _create_test_image_bytes()
|
|
files = {"file": ("test.jpg", image_bytes, "image/jpeg")}
|
|
|
|
with patch("app.clean_garment", side_effect=CleanerError("OpenRouter upstream failure")):
|
|
response = client.post("/api/clean", files=files)
|
|
assert response.status_code == 502
|
|
assert "OpenRouter upstream failure" in response.json()["detail"]
|
|
|
|
|
|
def test_clean_success_happy_path():
|
|
"""POST /api/clean with valid inputs returns full JSON result with image data URI."""
|
|
from app import app
|
|
|
|
client = TestClient(app)
|
|
image_bytes = _create_test_image_bytes()
|
|
output_bytes = b"cleaned-image-data-bytes"
|
|
mock_result = CleanResult(
|
|
image_bytes=output_bytes,
|
|
media_type="image/png",
|
|
cost=0.0125,
|
|
width=1000,
|
|
height=1500,
|
|
original_dimensions=(1000, 1500),
|
|
model_dimensions=(800, 1200),
|
|
)
|
|
|
|
files = {"file": ("garment.jpg", image_bytes, "image/jpeg")}
|
|
data = {
|
|
"prompt": "Custom studio prompt with white background.",
|
|
"model": "custom/test-model",
|
|
"api_key": "sk-or-v1-custom-form-key",
|
|
"restore_res": "true",
|
|
}
|
|
|
|
with patch("app.clean_garment", return_value=mock_result) as mock_clean:
|
|
response = client.post("/api/clean", files=files, data=data)
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
|
|
# Assert response keys and values
|
|
expected_data_uri = f"data:image/png;base64,{base64.b64encode(output_bytes).decode()}"
|
|
assert payload["image"] == expected_data_uri
|
|
assert payload["cost"] == 0.0125
|
|
assert payload["width"] == 1000
|
|
assert payload["height"] == 1500
|
|
assert payload["original_dimensions"] == [1000, 1500]
|
|
assert payload["model_dimensions"] == [800, 1200]
|
|
assert payload["was_rescaled"] is True
|
|
assert payload["media_type"] == "image/png"
|
|
|
|
# Assert arguments passed to clean_garment
|
|
mock_clean.assert_called_once_with(
|
|
image_data=image_bytes,
|
|
prompt="Custom studio prompt with white background.",
|
|
api_key="sk-or-v1-custom-form-key",
|
|
model="custom/test-model",
|
|
restore_res=True,
|
|
)
|
|
|
|
|
|
def test_clean_default_parameters():
|
|
"""POST /api/clean uses default prompt, model, and restore_res if omitted in form."""
|
|
from app import app, DEFAULT_PROMPT
|
|
from cleaner import DEFAULT_MODEL
|
|
|
|
client = TestClient(app)
|
|
image_bytes = _create_test_image_bytes()
|
|
mock_result = CleanResult(
|
|
image_bytes=image_bytes,
|
|
media_type="image/jpeg",
|
|
cost=None,
|
|
width=50,
|
|
height=50,
|
|
original_dimensions=(50, 50),
|
|
model_dimensions=(50, 50),
|
|
)
|
|
|
|
files = {"file": ("garment.jpg", image_bytes, "image/jpeg")}
|
|
data = {"api_key": "sk-or-v1-key"}
|
|
|
|
with patch("app.clean_garment", return_value=mock_result) as mock_clean:
|
|
response = client.post("/api/clean", files=files, data=data)
|
|
assert response.status_code == 200
|
|
mock_clean.assert_called_once_with(
|
|
image_data=image_bytes,
|
|
prompt=DEFAULT_PROMPT,
|
|
api_key="sk-or-v1-key",
|
|
model=DEFAULT_MODEL,
|
|
restore_res=True,
|
|
)
|