tests/test_api.py

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