Spaces:
Sleeping
Sleeping
File size: 5,229 Bytes
235461a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | from unittest.mock import patch
import pytest
from fastapi.testclient import TestClient
# Mock dependencies before importing the app
with patch("langchain_google_genai.ChatGoogleGenerativeAI"):
with patch("langgraph.prebuilt.create_react_agent"):
with patch("langgraph.checkpoint.postgres.PostgresSaver"):
with patch("llm.connection_manager.get_checkpointer"):
with patch("llm.connection_manager._test_connection", return_value=True):
from server.main import app
client = TestClient(app)
class TestAPI:
"""Test cases for the API endpoints."""
def test_health_check(self):
"""Test the health check endpoint."""
response = client.get("/health")
assert response.status_code == 200
data = response.json()
assert data["status"] in ["healthy", "degraded", "unhealthy"]
assert data["service"] == "ai-image-editor-api"
def test_root_endpoint(self):
"""Test the root endpoint."""
response = client.get("/")
assert response.status_code == 200
data = response.json()
assert "message" in data
@patch("server.main.chat_with_agent")
def test_chat_endpoint_basic(self, mock_chat_with_agent):
"""Test basic chat endpoint without image generation."""
mock_chat_with_agent.return_value = ("Hello! I can help you with image editing.", None)
request_data = {"message": "Hello", "selected_images": [], "user_id": "test_user", "client_ip": "127.0.0.1"}
response = client.post("/chat", json=request_data)
assert response.status_code == 200
data = response.json()
assert data["response"] == "Hello! I can help you with image editing."
assert data["status"] == "success"
assert data["generated_image"] is None
@patch("server.main.chat_with_agent")
def test_chat_endpoint_with_image_generation(self, mock_chat_with_agent):
"""Test chat endpoint with image generation."""
generated_image_data = {
"id": "test-uuid-123",
"url": "https://test-bucket.s3.amazonaws.com/test-url",
"title": "Generated Test Image",
"description": "AI-generated image: A beautiful sunset",
"timestamp": "2024-01-01T00:00:00Z",
"type": "generated",
}
mock_chat_with_agent.return_value = ("I've generated an image for you!", generated_image_data)
request_data = {"message": "Generate an image of a sunset", "selected_images": [], "user_id": "test_user", "client_ip": "127.0.0.1"}
response = client.post("/chat", json=request_data)
assert response.status_code == 200
data = response.json()
assert data["response"] == "I've generated an image for you!"
assert data["status"] == "success"
assert data["generated_image"] is not None
assert data["generated_image"]["id"] == "test-uuid-123"
@patch("server.main.chat_with_agent")
def test_chat_endpoint_with_selected_images(self, mock_chat_with_agent):
"""Test chat endpoint with selected images."""
mock_chat_with_agent.return_value = ("I see your selected images!", None)
request_data = {
"message": "Edit these images",
"selected_images": [
{
"id": "img-1",
"url": "https://example.com/img1.jpg",
"title": "Test Image 1",
"description": "A test image",
"timestamp": "2024-01-01T00:00:00Z",
"type": "uploaded",
}
],
"user_id": "test_user",
"client_ip": "127.0.0.1",
}
response = client.post("/chat", json=request_data)
assert response.status_code == 200
# Verify that chat_with_agent was called with the correct data
mock_chat_with_agent.assert_called_once()
call_args = mock_chat_with_agent.call_args
assert call_args[1]["message"] == "Edit these images"
assert call_args[1]["user_id"] == "test_user"
assert len(call_args[1]["selected_images"]) == 1
@patch("server.main.chat_with_agent")
def test_chat_endpoint_error_handling(self, mock_chat_with_agent):
"""Test chat endpoint error handling."""
mock_chat_with_agent.side_effect = Exception("Agent error")
request_data = {"message": "Hello", "selected_images": [], "user_id": "test_user", "client_ip": "127.0.0.1"}
response = client.post("/chat", json=request_data)
assert response.status_code == 500
data = response.json()
assert "Error processing request" in data["detail"]
def test_chat_endpoint_missing_client_ip(self):
"""Test chat endpoint with missing client IP."""
request_data = {"message": "Hello", "selected_images": [], "user_id": "test_user"}
response = client.post("/chat", json=request_data)
assert response.status_code == 200
data = response.json()
assert data["status"] == "error"
assert "Client IP not found" in data["response"]
if __name__ == "__main__":
pytest.main([__file__])
|