gemini test

tests/test_gemini_router.py

@@ -0,0 +1,598 @@
+"""
+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 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 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 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 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 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 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 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 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 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"])