tests/test_tryon.py

import base64
import pytest
from fastapi.testclient import TestClient
from unittest.mock import patch, MagicMock
from app.main import app
from app.models.wardrobe import WardrobeItem
from app.models.tryon import TryOnTask
from app.models.user import User
from app.core.security import get_password_hash
from app.core.database import get_db
from app.workers.tasks import run_tryon

client = TestClient(app)

# A valid 1x1 PNG, base64-encoded — small enough to pass through PIL in tests.
DUMMY_PNG_B64 = (
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
)
DUMMY_PNG_BYTES = base64.b64decode(DUMMY_PNG_B64)


@pytest.fixture(autouse=True)
def override_db(db_session):
    # Override get_db dependency in FastAPI app
    app.dependency_overrides[get_db] = lambda: db_session
    yield
    app.dependency_overrides.clear()


@pytest.fixture
def test_user_and_token(db_session):
    # Create a user
    user = User(email="test_tryon@example.com", password_hash=get_password_hash("password123"))
    db_session.add(user)
    db_session.commit()
    db_session.refresh(user)

    # Create a wardrobe item owned by user
    item = WardrobeItem(
        user_id=user.id,
        category="shirt",
        image_url="http://localhost:8002/static/sample.jpg"
    )
    db_session.add(item)
    db_session.commit()
    db_session.refresh(item)

    # Mock the authentication dependency to return this user
    from app.api.deps import get_current_user
    app.dependency_overrides[get_current_user] = lambda: user

    yield user, item

    # Clean up overrides
    app.dependency_overrides.clear()


def test_trigger_try_on_success(db_session, test_user_and_token):
    user, item = test_user_and_token

    # We mock the celery task delay
    with patch("app.api.v1.tryon.run_tryon.delay") as mock_delay:
        response = client.post(
            "/api/v1/try-on/run",
            json={"garment_item_id": str(item.id)}
        )
        assert response.status_code == 201
        data = response.json()
        assert data["status"] == "PENDING"
        assert data["garment_item_id"] == str(item.id)
        mock_delay.assert_called_once_with(data["id"])


def test_get_try_on_task(db_session, test_user_and_token):
    user, item = test_user_and_token

    # Pre-insert a task
    task = TryOnTask(
        user_id=user.id,
        garment_item_id=item.id,
        status="PROCESSING"
    )
    db_session.add(task)
    db_session.commit()
    db_session.refresh(task)

    response = client.get(f"/api/v1/try-on/tasks/{task.id}")
    assert response.status_code == 200
    data = response.json()
    assert data["status"] == "PROCESSING"


def test_celery_task_execution(db_session, test_user_and_token):
    """The try-on task should load the garment, call the image-gen service and upload."""
    user, item = test_user_and_token

    task = TryOnTask(
        user_id=user.id,
        garment_item_id=item.id,
        status="PENDING"
    )
    db_session.add(task)
    db_session.commit()
    db_session.refresh(task)

    task_id_val = task.id

    class SafeSessionWrapper:
        def __init__(self, session):
            self._session = session

        def __getattr__(self, name):
            if name == "close":
                return lambda: None
            return getattr(self._session, name)

    # Garment image is fetched over HTTP (no local /static file in tests).
    def mock_get_responses(url, *args, **kwargs):
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.content = DUMMY_PNG_BYTES
        mock_resp.raise_for_status = lambda: None
        return mock_resp

    with patch("app.workers.tasks.SessionLocal", return_value=SafeSessionWrapper(db_session)):
        # Image generation is provider-specific; mock the service so the task is provider-agnostic.
        with patch("app.workers.tasks.tryon_service.generate", return_value=DUMMY_PNG_BYTES) as mock_generate:
            with patch(
                "app.workers.tasks.storage_service.upload_file",
                return_value="http://localhost:8002/static/tryon_result.jpg",
            ) as mock_upload:
                with patch("app.workers.tasks.requests.get", side_effect=mock_get_responses):
                    res = run_tryon(str(task_id_val))

                assert res is True
                mock_generate.assert_called_once()
                # garments are passed as (bytes, category) tuples
                _, kwargs = mock_generate.call_args
                assert kwargs["garments"][0][1] == "shirt"

                updated_task = db_session.query(TryOnTask).filter(TryOnTask.id == task_id_val).first()
                assert updated_task.status == "SUCCESS"
                assert updated_task.result_url == "http://localhost:8002/static/tryon_result.jpg"
                mock_upload.assert_called_once()


def test_tryon_service_gemini_parses_image():
    """The Gemini provider should decode the inlineData image from a generateContent response."""
    from app.services import tryon_service as ts_module
    from app.services import gemini_keys as gk_module

    def mock_post(url, *args, **kwargs):
        assert "generativelanguage.googleapis.com" in url
        mock_resp = MagicMock()
        mock_resp.status_code = 200
        mock_resp.raise_for_status = lambda: None
        mock_resp.json.return_value = {
            "candidates": [
                {"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": DUMMY_PNG_B64}}]}}
            ]
        }
        return mock_resp

    with patch.object(ts_module.settings, "TRYON_PROVIDER", "gemini"), \
         patch.object(ts_module.settings, "GEMINI_IMAGE_MODEL", "gemini-2.5-flash-image"), \
         patch.object(gk_module.gemini_keys, "_keys", ["test-gemini-key"]), \
         patch.object(gk_module.gemini_keys, "_idx", 0), \
         patch.object(ts_module.requests, "post", side_effect=mock_post):
        svc = ts_module.TryOnService()
        out = svc.generate(person_bytes=DUMMY_PNG_BYTES, garments=[(DUMMY_PNG_BYTES, "shirt")])
        assert out == DUMMY_PNG_BYTES


def test_gemini_key_rotation_on_failure():
    """A failing key (429) should rotate to the next key, which then succeeds."""
    from app.services import gemini_keys as gk_module

    calls = []

    def mock_post(url, *args, **kwargs):
        calls.append(url)
        mock_resp = MagicMock()
        # First key -> 429 (quota), second key -> 200
        if len(calls) == 1:
            mock_resp.status_code = 429
            mock_resp.text = "rate limited"
        else:
            mock_resp.status_code = 200
            mock_resp.text = "ok"
        return mock_resp

    with patch.object(gk_module.gemini_keys, "_keys", ["bad-key", "good-key"]), \
         patch.object(gk_module.gemini_keys, "_idx", 0), \
         patch.object(gk_module.requests, "post", side_effect=lambda *a, **k: None):
        resp = gk_module.gemini_keys.request(
            lambda key: mock_post(f"https://x?key={key}")
        )
        assert resp.status_code == 200
        assert "key=bad-key" in calls[0]
        assert "key=good-key" in calls[1]


def test_gemini_no_rotation_on_hard_quota():
    """A 429 with 'limit: 0' (model not on plan) must fail fast — exactly one call, no rotation."""
    from app.services import gemini_keys as gk_module

    calls = []

    def do_request(key):
        calls.append(key)
        mock_resp = MagicMock()
        mock_resp.status_code = 429
        mock_resp.text = '{"error":{"code":429,"message":"...limit: 0, model: gemini-2.5-flash-preview-image"}}'
        return mock_resp

    with patch.object(gk_module.gemini_keys, "_keys", ["k1", "k2", "k3", "k4", "k5"]), \
         patch.object(gk_module.gemini_keys, "_idx", 0):
        resp = gk_module.gemini_keys.request(do_request)
        # request() returns the 429 (caller raise_for_status), and only ONE key was tried.
        assert resp.status_code == 429
        assert len(calls) == 1


def test_tryon_service_disabled_raises():
    """With no API key configured, generate() should raise a clear error (task -> FAILED)."""
    from app.services import tryon_service as ts_module
    from app.services import gemini_keys as gk_module

    with patch.object(ts_module.settings, "TRYON_PROVIDER", "gemini"), \
         patch.object(gk_module.gemini_keys, "_keys", []):
        svc = ts_module.TryOnService()
        assert svc.enabled is False
        with pytest.raises(RuntimeError):
            svc.generate(person_bytes=DUMMY_PNG_BYTES, garments=[(DUMMY_PNG_BYTES, "shirt")])