Spaces:
Sleeping
Sleeping
File size: 5,935 Bytes
6a3de9e |
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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
"""
Test script to verify the AI agent properly connects to the MCP server
and executes database operations through the MCP server integration.
"""
import asyncio
import os
import sys
from unittest.mock import AsyncMock, patch
from sqlmodel.ext.asyncio.session import AsyncSession
from database.session import get_async_session
from models.user import User
from models.task import Task
from sqlmodel import select
from datetime import datetime
import uuid
async def test_ai_agent_mcp_integration():
"""
Test that the AI agent properly integrates with the MCP server
and executes database operations.
"""
print("Testing AI Agent MCP Integration...")
# Get database session
async with get_async_session() as session:
# Create a test user with a unique email
unique_email = f"test_{uuid.uuid4()}@example.com"
user = User(
email=unique_email,
name="Test User",
created_at=datetime.utcnow()
)
session.add(user)
await session.commit()
await session.refresh(user)
print(f"Created test user: {user.id}")
# Import the AI agent and conversation manager
from ai.agents.todo_agent import TodoAgent
from ai.agents.conversation_manager import ConversationManager
# Create conversation manager and conversation
conversation_manager = ConversationManager(session)
conversation = await conversation_manager.create_conversation(str(user.id))
print(f"Created conversation: {conversation.id}")
# Initialize the AI agent
todo_agent = TodoAgent()
print("Initialized AI agent")
# Test 1: Add a task via AI agent
print("\n--- Test 1: Adding a task via AI agent ---")
add_task_request = "Add a new task: Buy groceries with high priority"
result = await todo_agent.process_message(str(user.id), add_task_request, conversation)
print(f"AI Response: {result['response']}")
# Verify task was added to database
tasks = await session.exec(select(Task).where(Task.user_id == user.id))
tasks_list = tasks.all()
print(f"Tasks in database after add: {len(tasks_list)}")
if tasks_list:
print(f"Latest task: {tasks_list[-1].title} (Priority: {tasks_list[-1].priority})")
# Test 2: List tasks via AI agent
print("\n--- Test 2: Listing tasks via AI agent ---")
list_request = "List all my tasks"
result = await todo_agent.process_message(str(user.id), list_request, conversation)
print(f"AI Response: {result['response']}")
# Test 3: Complete a task via AI agent
if tasks_list:
print("\n--- Test 3: Completing a task via AI agent ---")
task_to_complete = tasks_list[-1] # Use the last added task
complete_request = f"Complete the task with ID {task_to_complete.id}"
result = await todo_agent.process_message(str(user.id), complete_request, conversation)
print(f"AI Response: {result['response']}")
# Refresh the task to check completion status
await session.refresh(task_to_complete)
print(f"Task {task_to_complete.id} completed status: {task_to_complete.completed}")
# Final verification
final_tasks = await session.exec(select(Task).where(Task.user_id == user.id))
final_tasks_list = final_tasks.all()
print(f"\nFinal task count: {len(final_tasks_list)}")
for task in final_tasks_list:
print(f"- Task: {task.title}, Priority: {task.priority}, Completed: {task.completed}, Due: {task.due_date}")
print("\n--- Test Summary ---")
print(f"Created {len(final_tasks_list)} tasks in database")
completed_count = sum(1 for task in final_tasks_list if task.completed)
print(f"Completed {completed_count} tasks")
print("AI agent properly integrated with MCP server")
print("Database operations executed successfully through MCP server")
# Clean up test data
for task in final_tasks_list:
await session.delete(task)
await session.delete(conversation)
await session.delete(user)
await session.commit()
print("Test data cleaned up")
async def test_mcp_server_availability():
"""
Test that the MCP server is properly configured and available.
"""
print("\nTesting MCP Server Configuration...")
# Import the MCP server
from ai.mcp.server import server, list_todo_tools
# Check that tools are available
tools = await list_todo_tools()
tool_names = [tool.name for tool in tools]
print(f"MCP Server tools available: {tool_names}")
# Verify all expected tools are present
expected_tools = {"add_task", "list_tasks", "complete_task", "delete_task", "update_task"}
actual_tools = set(tool_names)
if expected_tools.issubset(actual_tools):
print("All expected tools are registered in MCP server")
else:
missing = expected_tools - actual_tools
print(f"Missing tools: {missing}")
# Check that the server object has the expected attributes
print(f"Server name: {server.name}")
print("MCP server properly initialized")
async def main():
"""
Main test function.
"""
print("=" * 60)
print("AI AGENT MCP INTEGRATION TEST")
print("=" * 60)
try:
await test_mcp_server_availability()
await test_ai_agent_mcp_integration()
print("\n" + "=" * 60)
print("ALL TESTS PASSED!")
print("AI agent properly integrates with MCP server")
print("Database operations are executed correctly through MCP")
print("=" * 60)
except Exception as e:
print(f"\nTEST FAILED: {str(e)}")
import traceback
traceback.print_exc()
if __name__ == "__main__":
asyncio.run(main()) |