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