Update try-on tests for the API-based service

Mock tryon_service.generate instead of CatVTON/AutoMasker; add tests for the
Gemini provider image parsing and the missing-key error path.

tests/test_tryon.py

@@ -1,5 +1,5 @@
+import base64
 import pytest
-import torch
 from fastapi.testclient import TestClient
 from unittest.mock import patch, MagicMock
 from app.main import app
@@ -8,10 +8,17 @@ 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_idm_vton_tryon
+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
@@ -19,6 +26,7 @@ def override_db(db_session):
     yield
     app.dependency_overrides.clear()
 
+
 @pytest.fixture
 def test_user_and_token(db_session):
     # Create a user
@@ -46,11 +54,12 @@ def test_user_and_token(db_session):
     # 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_idm_vton_tryon.delay") as mock_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)}
@@ -61,9 +70,10 @@ def test_trigger_try_on_success(db_session, test_user_and_token):
         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,
@@ -79,10 +89,11 @@ def test_get_try_on_task(db_session, test_user_and_token):
     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
 
-    # Pre-insert a task
     task = TryOnTask(
         user_id=user.id,
         garment_item_id=item.id,
@@ -97,88 +108,74 @@ def test_celery_task_execution(db_session, test_user_and_token):
     class SafeSessionWrapper:
         def __init__(self, session):
             self._session = session
+
         def __getattr__(self, name):
             if name == "close":
                 return lambda: None
             return getattr(self._session, name)
 
-    # Mock responses for NVIDIA VLM and Gemini Imagen
-    def mock_post_responses(url, *args, **kwargs):
+    # 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
-        if "integrate.api.nvidia.com" in url:
-            mock_resp.json.return_value = {
-                "choices": [
-                    {
-                        "message": {
-                            "content": "A photo of a fair-skinned woman standing with brown hair"
-                        }
-                    }
-                ]
-            }
-        elif "generativelanguage.googleapis.com" in url:
-            import base64
-            # 1x1 pixel base64 image
-            dummy_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
-            mock_resp.json.return_value = {
-                "predictions": [
-                    {
-                        "bytesBase64Encoded": dummy_b64
-                    }
-                ]
-            }
-        elif "api.openai.com" in url:
-            dummy_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
-            mock_resp.json.return_value = {
-                "data": [
-                    {
-                        "b64_json": dummy_b64
-                    }
-                ]
-            }
+        mock_resp.content = DUMMY_PNG_BYTES
+        mock_resp.raise_for_status = lambda: None
         return mock_resp
 
-    # Mock AutoMasker
-    mock_automasker = MagicMock()
-    from PIL import Image
-    dummy_mask = Image.new("L", (768, 1024), 255)
-    mock_automasker.return_value = {"mask": dummy_mask}
-    
-    mock_mask_processor = MagicMock()
-    mock_mask_processor.blur.return_value = dummy_mask
-    
-    # Mock CatVTON Pipeline
-    mock_pipeline = MagicMock()
-    def mock_pipeline_call(image, condition_image, mask, **kwargs):
-        return [image]
-    mock_pipeline.side_effect = mock_pipeline_call
-
-    # Mock requests.get to return a dummy 1x1 PNG image
-    def mock_get_responses(url, *args, **kwargs):
+    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
+
+    def mock_post(url, *args, **kwargs):
+        assert "generativelanguage.googleapis.com" in url
         mock_resp = MagicMock()
         mock_resp.status_code = 200
-        mock_resp.content = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15c4\x00\x00\x00\rIDATx\x9cc`\x00\x00\x00\x02\x00\x01H\xaf\xa4q\x00\x00\x00\x00IEND\xaeB`\x82'
+        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
 
-    # Mock SessionLocal in app.workers.tasks to return our wrapped session
-    with patch("app.workers.tasks.SessionLocal", return_value=SafeSessionWrapper(db_session)):
-        # Mock get_automasker and get_catvton_pipeline to avoid loading models
-        with patch("app.workers.tasks.get_automasker", return_value=(mock_automasker, mock_mask_processor)):
-            with patch("app.workers.tasks.get_catvton_pipeline", return_value=mock_pipeline):
-                # Mock the storage service upload_file to return a valid URL
-                with patch("app.workers.tasks.storage_service.upload_file", return_value="http://localhost:8002/static/tryon_result.jpg") as mock_upload:
-                    # Mock requests.post to avoid real network calls
-                    with patch("app.workers.tasks.requests.post", side_effect=mock_post_responses) as mock_post:
-                        # Mock requests.get to return a dummy PNG image
-                        with patch("app.workers.tasks.requests.get", side_effect=mock_get_responses) as mock_get:
-                            # Patch settings to ensure OpenAI API key is set for test
-                            with patch("app.workers.tasks.settings.OPENAI_API_KEY", "test-openai-key"):
-                                res = run_idm_vton_tryon(str(task_id_val))
-                            
-                            assert res is True
-                            
-                            # Reload task from DB
-                            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()
+    with patch.object(ts_module.settings, "TRYON_PROVIDER", "gemini"), \
+         patch.object(ts_module.settings, "GEMINI_API_KEY", "test-gemini-key"), \
+         patch.object(ts_module.settings, "GEMINI_IMAGE_MODEL", "gemini-2.5-flash-image"), \
+         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_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
+
+    with patch.object(ts_module.settings, "TRYON_PROVIDER", "gemini"), \
+         patch.object(ts_module.settings, "GEMINI_API_KEY", None):
+        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")])