tests/ai/test_agent_integration.py

"""
Integration tests for the AI agent functionality.

These tests verify the end-to-end functionality of the AI agent with
realistic scenarios and actual tool execution.
"""
import pytest
from fastapi.testclient import TestClient
from unittest.mock import AsyncMock, MagicMock, patch
from main import app
from models.conversation import Conversation
from models.message import Message
from uuid import UUID, uuid4


@pytest.fixture
def client():
    """Create a test client for the API."""
    return TestClient(app)


@pytest.mark.asyncio
async def test_agent_full_add_task_workflow(client):
    """Test the complete workflow for adding a task via AI agent."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to simulate add_task workflow
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "I've added the task 'Buy groceries' to your list successfully!",
            "conversation_id": str(uuid4()),
            "tool_calls": [
                {
                    "id": "call_123",
                    "function": {
                        "name": "add_task",
                        "arguments": "{\"user_id\":\"test-user-123\",\"title\":\"Buy groceries\",\"description\":\"Need to buy milk, bread, and eggs\",\"priority\":\"medium\"}"
                    }
                }
            ],
            "requires_action": True
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a request to add a task
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Add a task: Buy groceries - need milk, bread, and eggs"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "Buy groceries" in data["response"]
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) > 0
        assert data["tool_calls"][0]["function"]["name"] == "add_task"


@pytest.mark.asyncio
async def test_agent_full_list_tasks_workflow(client):
    """Test the complete workflow for listing tasks via AI agent."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to simulate list_tasks workflow
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "Here are your tasks: 1. Buy groceries (pending), 2. Clean house (pending), 3. Pay bills (completed)",
            "conversation_id": str(uuid4()),
            "tool_calls": [
                {
                    "id": "call_456",
                    "function": {
                        "name": "list_tasks",
                        "arguments": "{\"user_id\":\"test-user-123\",\"status\":\"all\"}"
                    }
                }
            ],
            "requires_action": True
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a request to list tasks
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Show me all my tasks"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "tasks" in data["response"].lower()
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) > 0
        assert data["tool_calls"][0]["function"]["name"] == "list_tasks"


@pytest.mark.asyncio
async def test_agent_full_complete_task_workflow(client):
    """Test the complete workflow for completing a task via AI agent."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to simulate complete_task workflow
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "I've marked task 'Buy groceries' as completed successfully!",
            "conversation_id": str(uuid4()),
            "tool_calls": [
                {
                    "id": "call_789",
                    "function": {
                        "name": "complete_task",
                        "arguments": "{\"user_id\":\"test-user-123\",\"task_id\":\"1\"}"
                    }
                }
            ],
            "requires_action": True
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a request to complete a task
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Complete task 1"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "completed" in data["response"].lower()
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) > 0
        assert data["tool_calls"][0]["function"]["name"] == "complete_task"


@pytest.mark.asyncio
async def test_agent_full_delete_task_workflow(client):
    """Test the complete workflow for deleting a task via AI agent."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to simulate delete_task workflow
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "I've deleted task 'Clean house' from your list successfully!",
            "conversation_id": str(uuid4()),
            "tool_calls": [
                {
                    "id": "call_abc",
                    "function": {
                        "name": "delete_task",
                        "arguments": "{\"user_id\":\"test-user-123\",\"task_id\":\"2\"}"
                    }
                }
            ],
            "requires_action": True
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a request to delete a task
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Delete task 2"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "deleted" in data["response"].lower()
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) > 0
        assert data["tool_calls"][0]["function"]["name"] == "delete_task"


@pytest.mark.asyncio
async def test_agent_full_update_task_workflow(client):
    """Test the complete workflow for updating a task via AI agent."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to simulate update_task workflow
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "I've updated task 'Buy groceries' with the new due date successfully!",
            "conversation_id": str(uuid4()),
            "tool_calls": [
                {
                    "id": "call_def",
                    "function": {
                        "name": "update_task",
                        "arguments": "{\"user_id\":\"test-user-123\",\"task_id\":\"1\",\"due_date\":\"2024-12-31\"}"
                    }
                }
            ],
            "requires_action": True
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a request to update a task
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Update task 1: Change due date to December 31st"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "updated" in data["response"].lower()
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) > 0
        assert data["tool_calls"][0]["function"]["name"] == "update_task"


@pytest.mark.asyncio
async def test_agent_error_handling_workflow(client):
    """Test the agent's error handling workflow."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent to raise an exception
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(side_effect=Exception("API Error"))
        mock_agent_class.return_value = mock_agent_instance

        # Send a request that will cause an error
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Add a task: Test error handling"
            },
            headers={"Content-Type": "application/json"}
        )

        # Should return a server error
        assert response.status_code == 500


@pytest.mark.asyncio
async def test_agent_unrecognized_command_response(client):
    """Test the agent's response to unrecognized commands."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to return a clarification response
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "I'm sorry, I didn't understand that command. Could you please rephrase or be more specific?",
            "conversation_id": str(uuid4()),
            "tool_calls": [],
            "requires_action": False
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send an unrecognized command
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Random command that doesn't make sense for task management"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "sorry" in data["response"].lower() or "understand" in data["response"].lower()
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) == 0
        assert data["requires_action"] is False


@pytest.mark.asyncio
async def test_agent_multiple_tool_calls_in_single_request(client):
    """Test the agent's ability to handle multiple tool calls in a single request."""
    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance to return multiple tool calls
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "I've processed your request by adding a task and listing your tasks.",
            "conversation_id": str(uuid4()),
            "tool_calls": [
                {
                    "id": "call_multi_1",
                    "function": {
                        "name": "add_task",
                        "arguments": "{\"user_id\":\"test-user-123\",\"title\":\"New task\",\"priority\":\"high\"}"
                    }
                },
                {
                    "id": "call_multi_2",
                    "function": {
                        "name": "list_tasks",
                        "arguments": "{\"user_id\":\"test-user-123\",\"status\":\"all\"}"
                    }
                }
            ],
            "requires_action": True
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a complex request that might trigger multiple tools
        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Add a high priority task 'New task' and then show me all my tasks"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert "response" in data
        assert "conversation_id" in data
        assert "tool_calls" in data
        assert len(data["tool_calls"]) == 2
        assert data["tool_calls"][0]["function"]["name"] in ["add_task", "list_tasks"]
        assert data["tool_calls"][1]["function"]["name"] in ["add_task", "list_tasks"]


@pytest.mark.asyncio
async def test_agent_conversation_management(client):
    """Test that the agent properly manages conversation state."""
    test_conversation_id = uuid4()

    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager methods
        mock_conv_instance = MagicMock()
        mock_conv_instance.get_conversation = AsyncMock(return_value=MagicMock(id=test_conversation_id, user_id="test-user-123"))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent instance
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "Processed your message in conversation",
            "conversation_id": str(test_conversation_id),
            "tool_calls": [],
            "requires_action": False
        })
        mock_agent_class.return_value = mock_agent_instance

        # Send a request with an existing conversation ID
        response = client.post(
            f"/api/test-user-123/chat",
            params={"conversation_id": str(test_conversation_id)},
            json={
                "message": "Continue working on my tasks"
            },
            headers={"Content-Type": "application/json"}
        )

        # Verify the response
        assert response.status_code == 200
        data = response.json()

        assert data["conversation_id"] == str(test_conversation_id)
        assert "response" in data


@pytest.mark.asyncio
async def test_agent_response_timing_and_performance(client):
    """Test that the agent responds within reasonable time limits."""
    import time

    with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
         patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

        # Mock conversation manager
        mock_conv_instance = MagicMock()
        mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
        mock_conv_instance.add_message = AsyncMock()
        mock_conv_instance.update_conversation_timestamp = AsyncMock()
        mock_conv_manager.return_value = mock_conv_instance

        # Mock agent with a simulated processing time
        mock_agent_instance = MagicMock()
        mock_agent_instance.process_message = AsyncMock(return_value={
            "response": "Task processed successfully",
            "conversation_id": str(uuid4()),
            "tool_calls": [],
            "requires_action": False
        })
        mock_agent_class.return_value = mock_agent_instance

        # Measure response time
        start_time = time.time()

        response = client.post(
            "/api/test-user-123/chat",
            json={
                "message": "Add a simple task: Test timing"
            },
            headers={"Content-Type": "application/json"}
        )

        end_time = time.time()
        response_time = end_time - start_time

        # Verify response is successful and timely
        assert response.status_code == 200
        assert response_time < 5.0  # Should respond in under 5 seconds


@pytest.mark.asyncio
async def test_agent_user_isolation(client):
    """Test that different users' tasks are properly isolated."""
    test_users = ["user-1", "user-2", "user-3"]

    for user_id in test_users:
        with patch('ai.endpoints.chat.ConversationManager') as mock_conv_manager, \
             patch('ai.endpoints.chat.TodoAgent') as mock_agent_class:

            # Mock conversation manager
            mock_conv_instance = MagicMock()
            mock_conv_instance.create_conversation = AsyncMock(return_value=MagicMock(id=uuid4()))
            mock_conv_instance.add_message = AsyncMock()
            mock_conv_instance.update_conversation_timestamp = AsyncMock()
            mock_conv_manager.return_value = mock_conv_instance

            # Mock agent instance
            mock_agent_instance = MagicMock()
            mock_agent_instance.process_message = AsyncMock(return_value={
                "response": f"Processed for {user_id}",
                "conversation_id": str(uuid4()),
                "tool_calls": [],
                "requires_action": False
            })
            mock_agent_class.return_value = mock_agent_instance

            response = client.post(
                f"/api/{user_id}/chat",
                json={
                    "message": "Show me my tasks"
                },
                headers={"Content-Type": "application/json"}
            )

            # Verify each user gets a proper response
            assert response.status_code == 200
            data = response.json()
            assert "response" in data
            # Verify the response mentions the correct user
            assert user_id in data["response"] or user_id.replace("-", "") in data["response"]


if __name__ == "__main__":
    pytest.main([__file__])