Spaces:
Sleeping
Sleeping
| """ | |
| Rigorous Tests for Gemini Router. | |
| Tests cover: | |
| 1. Request/Response Models | |
| 2. Helper functions (get_queue_position, create_job) | |
| 3. Job creation endpoints (5 types) | |
| 4. Job status polling | |
| 5. Video download | |
| 6. Job cancellation | |
| 7. Models endpoint | |
| Uses mocked auth, database, and worker pool. | |
| """ | |
| import pytest | |
| import os | |
| import tempfile | |
| from datetime import datetime | |
| from unittest.mock import patch, MagicMock, AsyncMock | |
| from fastapi.testclient import TestClient | |
| # ============================================================================= | |
| # 1. Request Model Tests | |
| # ============================================================================= | |
| class TestRequestModels: | |
| """Test request model validation.""" | |
| def test_generate_animation_prompt_request(self): | |
| """GenerateAnimationPromptRequest validates correctly.""" | |
| from routers.gemini import GenerateAnimationPromptRequest | |
| req = GenerateAnimationPromptRequest( | |
| base64_image="base64data", | |
| mime_type="image/png", | |
| custom_prompt="Make it dramatic" | |
| ) | |
| assert req.base64_image == "base64data" | |
| assert req.mime_type == "image/png" | |
| assert req.custom_prompt == "Make it dramatic" | |
| def test_edit_image_request(self): | |
| """EditImageRequest validates correctly.""" | |
| from routers.gemini import EditImageRequest | |
| req = EditImageRequest( | |
| base64_image="base64data", | |
| mime_type="image/png", | |
| prompt="Add colors" | |
| ) | |
| assert req.prompt == "Add colors" | |
| def test_generate_video_request_defaults(self): | |
| """GenerateVideoRequest has correct defaults.""" | |
| from routers.gemini import GenerateVideoRequest | |
| req = GenerateVideoRequest( | |
| base64_image="base64data", | |
| mime_type="image/png", | |
| prompt="Animate" | |
| ) | |
| assert req.aspect_ratio == "16:9" | |
| assert req.resolution == "720p" | |
| assert req.number_of_videos == 1 | |
| def test_generate_video_request_custom_values(self): | |
| """GenerateVideoRequest accepts custom values.""" | |
| from routers.gemini import GenerateVideoRequest | |
| req = GenerateVideoRequest( | |
| base64_image="base64data", | |
| mime_type="image/png", | |
| prompt="Animate", | |
| aspect_ratio="9:16", | |
| resolution="1080p", | |
| number_of_videos=2 | |
| ) | |
| assert req.aspect_ratio == "9:16" | |
| assert req.resolution == "1080p" | |
| assert req.number_of_videos == 2 | |
| def test_generate_text_request(self): | |
| """GenerateTextRequest validates correctly.""" | |
| from routers.gemini import GenerateTextRequest | |
| req = GenerateTextRequest(prompt="Hello world") | |
| assert req.prompt == "Hello world" | |
| assert req.model is None | |
| def test_analyze_image_request(self): | |
| """AnalyzeImageRequest validates correctly.""" | |
| from routers.gemini import AnalyzeImageRequest | |
| req = AnalyzeImageRequest( | |
| base64_image="base64data", | |
| mime_type="image/jpeg", | |
| prompt="Describe this" | |
| ) | |
| assert req.prompt == "Describe this" | |
| # ============================================================================= | |
| # 2. Job Creation Endpoints Tests | |
| # ============================================================================= | |
| class TestJobCreationEndpoints: | |
| """Test job creation endpoints.""" | |
| def test_generate_animation_prompt_requires_auth(self): | |
| """generate-animation-prompt requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post( | |
| "/gemini/generate-animation-prompt", | |
| json={"base64_image": "abc", "mime_type": "image/png"} | |
| ) | |
| assert response.status_code in [401, 403, 422] | |
| def test_edit_image_requires_auth(self): | |
| """edit-image requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post( | |
| "/gemini/edit-image", | |
| json={"base64_image": "abc", "mime_type": "image/png", "prompt": "edit"} | |
| ) | |
| assert response.status_code in [401, 403, 422] | |
| def test_generate_video_requires_auth(self): | |
| """generate-video requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post( | |
| "/gemini/generate-video", | |
| json={"base64_image": "abc", "mime_type": "image/png", "prompt": "animate"} | |
| ) | |
| assert response.status_code in [401, 403, 422] | |
| def test_generate_text_requires_auth(self): | |
| """generate-text requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post( | |
| "/gemini/generate-text", | |
| json={"prompt": "Hello"} | |
| ) | |
| assert response.status_code in [401, 403, 422] | |
| def test_analyze_image_requires_auth(self): | |
| """analyze-image requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post( | |
| "/gemini/analyze-image", | |
| json={"base64_image": "abc", "mime_type": "image/png", "prompt": "analyze"} | |
| ) | |
| assert response.status_code in [401, 403, 422] | |
| # ============================================================================= | |
| # 3. Job Status Endpoint Tests | |
| # ============================================================================= | |
| class TestJobStatusEndpoint: | |
| """Test GET /job/{job_id} endpoint.""" | |
| def test_job_status_requires_auth(self): | |
| """Job status requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/job/job_123") | |
| assert response.status_code in [401, 403, 422] | |
| def test_job_status_not_found(self): | |
| """Return 404 for non-existent job.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_user.credits = 100 | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = None | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/job/job_nonexistent") | |
| assert response.status_code == 404 | |
| assert "not found" in response.json()["detail"].lower() | |
| def test_job_status_queued(self): | |
| """Return queued status with position.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_user.credits = 100 | |
| mock_job = MagicMock() | |
| mock_job.job_id = "job_123" | |
| mock_job.job_type = "text" | |
| mock_job.status = "queued" | |
| mock_job.created_at = datetime.utcnow() | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = mock_job | |
| # For queue position query | |
| mock_position_result = MagicMock() | |
| mock_position_result.scalar.return_value = 2 | |
| mock_db.execute.side_effect = [mock_result, mock_position_result] | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/job/job_123") | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert data["status"] == "queued" | |
| assert "position" in data | |
| def test_job_status_completed(self): | |
| """Return completed status with output.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_user.credits = 100 | |
| mock_job = MagicMock() | |
| mock_job.job_id = "job_123" | |
| mock_job.job_type = "text" | |
| mock_job.status = "completed" | |
| mock_job.created_at = datetime.utcnow() | |
| mock_job.completed_at = datetime.utcnow() | |
| mock_job.output_data = {"result": "Generated text"} | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = mock_job | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/job/job_123") | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert data["status"] == "completed" | |
| assert "output" in data | |
| def test_job_status_failed(self): | |
| """Return failed status with error.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_user.credits = 100 | |
| mock_job = MagicMock() | |
| mock_job.job_id = "job_123" | |
| mock_job.job_type = "text" | |
| mock_job.status = "failed" | |
| mock_job.created_at = datetime.utcnow() | |
| mock_job.completed_at = datetime.utcnow() | |
| mock_job.error_message = "API rate limited" | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = mock_job | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/job/job_123") | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert data["status"] == "failed" | |
| assert "error" in data | |
| # ============================================================================= | |
| # 4. Video Download Endpoint Tests | |
| # ============================================================================= | |
| class TestDownloadEndpoint: | |
| """Test GET /download/{job_id} endpoint.""" | |
| def test_download_requires_auth(self): | |
| """Download requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/download/job_123") | |
| assert response.status_code in [401, 403, 422] | |
| def test_download_job_not_found(self): | |
| """Return 404 for non-existent job.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = None | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/download/job_nonexistent") | |
| assert response.status_code == 404 | |
| def test_download_video_not_ready(self): | |
| """Return 400 if video not ready.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_job = MagicMock() | |
| mock_job.job_id = "job_123" | |
| mock_job.job_type = "video" | |
| mock_job.status = "processing" # Not completed | |
| mock_job.output_data = None | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = mock_job | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/download/job_123") | |
| assert response.status_code == 400 | |
| assert "not ready" in response.json()["detail"].lower() | |
| # ============================================================================= | |
| # 5. Job Cancellation Endpoint Tests | |
| # ============================================================================= | |
| class TestCancelEndpoint: | |
| """Test POST /job/{job_id}/cancel endpoint.""" | |
| def test_cancel_requires_auth(self): | |
| """Cancel requires authentication.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post("/gemini/job/job_123/cancel") | |
| assert response.status_code in [401, 403, 422] | |
| def test_cancel_job_not_found(self): | |
| """Return 404 for non-existent job.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = None | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post("/gemini/job/job_nonexistent/cancel") | |
| assert response.status_code == 404 | |
| def test_cancel_only_queued_jobs(self): | |
| """Only queued jobs can be cancelled.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_job = MagicMock() | |
| mock_job.job_id = "job_123" | |
| mock_job.status = "processing" # Cannot cancel | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = mock_job | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post("/gemini/job/job_123/cancel") | |
| assert response.status_code == 400 | |
| assert "cannot cancel" in response.json()["detail"].lower() | |
| def test_cancel_queued_job_success(self): | |
| """Successfully cancel a queued job.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| from core.dependencies import get_current_user | |
| from core.database import get_db | |
| app = FastAPI() | |
| mock_user = MagicMock() | |
| mock_user.user_id = "test-user" | |
| mock_job = MagicMock() | |
| mock_job.job_id = "job_123" | |
| mock_job.status = "queued" | |
| async def mock_get_db(): | |
| mock_db = AsyncMock() | |
| mock_result = MagicMock() | |
| mock_result.scalar_one_or_none.return_value = mock_job | |
| mock_db.execute.return_value = mock_result | |
| yield mock_db | |
| app.dependency_overrides[get_current_user] = lambda: mock_user | |
| app.dependency_overrides[get_db] = mock_get_db | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.post("/gemini/job/job_123/cancel") | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert data["success"] == True | |
| assert data["status"] == "cancelled" | |
| # ============================================================================= | |
| # 6. Models Endpoint Tests | |
| # ============================================================================= | |
| class TestModelsEndpoint: | |
| """Test GET /models endpoint.""" | |
| def test_get_models(self): | |
| """Get available models.""" | |
| from routers.gemini import router | |
| from fastapi import FastAPI | |
| app = FastAPI() | |
| app.include_router(router) | |
| client = TestClient(app) | |
| response = client.get("/gemini/models") | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert "models" in data | |
| assert "text_generation" in data["models"] | |
| assert "video_generation" in data["models"] | |
| if __name__ == "__main__": | |
| pytest.main([__file__, "-v"]) | |