test_agentai_integration.py

"""
Test script for InklyAI AgentAI integration.
"""

import requests
import json
import time
import os
from typing import Dict, Any


class AgentAIClient:
    """Client for testing AgentAI integration."""
    
    def __init__(self, base_url: str = "http://localhost:5000"):
        """
        Initialize the client.
        
        Args:
            base_url: Base URL of the API server
        """
        self.base_url = base_url
        self.session = requests.Session()
    
    def health_check(self) -> Dict[str, Any]:
        """Check API health."""
        response = self.session.get(f"{self.base_url}/health")
        return response.json()
    
    def register_agent(self, agent_id: str, signature_template: str) -> Dict[str, Any]:
        """Register an agent."""
        data = {
            "agent_id": agent_id,
            "signature_template": signature_template
        }
        response = self.session.post(f"{self.base_url}/register-agent", json=data)
        return response.json()
    
    def verify_signature(self, agent_id: str, signature_image: str, context: Dict[str, Any] = None) -> Dict[str, Any]:
        """Verify agent signature."""
        data = {
            "agent_id": agent_id,
            "signature_image": signature_image
        }
        if context:
            data["context"] = context
        
        response = self.session.post(f"{self.base_url}/verify-signature", json=data)
        return response.json()
    
    def get_agent_stats(self, agent_id: str) -> Dict[str, Any]:
        """Get agent statistics."""
        response = self.session.get(f"{self.base_url}/agent-stats/{agent_id}")
        return response.json()
    
    def list_agents(self) -> Dict[str, Any]:
        """List all agents."""
        response = self.session.get(f"{self.base_url}/list-agents")
        return response.json()
    
    def deactivate_agent(self, agent_id: str) -> Dict[str, Any]:
        """Deactivate an agent."""
        response = self.session.post(f"{self.base_url}/deactivate-agent/{agent_id}")
        return response.json()
    
    def reactivate_agent(self, agent_id: str) -> Dict[str, Any]:
        """Reactivate an agent."""
        response = self.session.post(f"{self.base_url}/reactivate-agent/{agent_id}")
        return response.json()


def test_agentai_integration():
    """Test the AgentAI integration."""
    print("🧪 Testing InklyAI AgentAI Integration")
    print("=" * 50)
    
    # Initialize client
    client = AgentAIClient()
    
    # Test 1: Health check
    print("\n1. Testing health check...")
    health = client.health_check()
    print(f"   Health status: {health}")
    assert health['status'] == 'healthy', "Health check failed"
    print("   ✅ Health check passed")
    
    # Test 2: Register agents
    print("\n2. Testing agent registration...")
    
    # Check if sample data exists
    if not os.path.exists('data/samples/john_doe_1.png'):
        print("   ⚠️  Sample data not found. Run demo.py first to create sample signatures.")
        return
    
    # Register test agents
    agents = [
        ("test_agent_001", "data/samples/john_doe_1.png"),
        ("test_agent_002", "data/samples/jane_smith_1.png"),
        ("test_agent_003", "data/samples/bob_wilson_1.png")
    ]
    
    for agent_id, signature_template in agents:
        result = client.register_agent(agent_id, signature_template)
        if result['success']:
            print(f"   ✅ Registered {agent_id}")
        else:
            print(f"   ❌ Failed to register {agent_id}: {result.get('error', 'Unknown error')}")
    
    # Test 3: List agents
    print("\n3. Testing agent listing...")
    agents_list = client.list_agents()
    print(f"   Total agents: {agents_list.get('total_agents', 0)}")
    for agent in agents_list.get('agents', []):
        print(f"   - {agent['agent_id']} (active: {agent['is_active']})")
    print("   ✅ Agent listing passed")
    
    # Test 4: Signature verification
    print("\n4. Testing signature verification...")
    
    # Test genuine signatures
    test_cases = [
        ("test_agent_001", "data/samples/john_doe_2.png", "Genuine signature"),
        ("test_agent_002", "data/samples/jane_smith_2.png", "Genuine signature"),
        ("test_agent_001", "data/samples/jane_smith_1.png", "Forged signature"),
        ("test_agent_002", "data/samples/bob_wilson_1.png", "Forged signature")
    ]
    
    for agent_id, signature_image, expected in test_cases:
        result = client.verify_signature(agent_id, signature_image)
        if result['success']:
            is_verified = result['is_verified']
            similarity = result['similarity_score']
            confidence = result['confidence']
            print(f"   {agent_id} vs {signature_image.split('/')[-1]}: "
                  f"Verified={is_verified}, Similarity={similarity:.3f}, Confidence={confidence:.3f}")
        else:
            print(f"   ❌ Verification failed for {agent_id}: {result.get('error', 'Unknown error')}")
    
    print("   ✅ Signature verification passed")
    
    # Test 5: Agent statistics
    print("\n5. Testing agent statistics...")
    for agent_id, _ in agents:
        stats = client.get_agent_stats(agent_id)
        if stats['success']:
            agent_stats = stats['stats']
            print(f"   {agent_id}: {agent_stats['total_verifications']} verifications, "
                  f"success rate: {agent_stats['success_rate']:.2%}")
        else:
            print(f"   ❌ Failed to get stats for {agent_id}")
    
    print("   ✅ Agent statistics passed")
    
    # Test 6: Agent deactivation/reactivation
    print("\n6. Testing agent deactivation/reactivation...")
    
    # Deactivate an agent
    deactivate_result = client.deactivate_agent("test_agent_001")
    if deactivate_result['success']:
        print("   ✅ Deactivated test_agent_001")
    else:
        print(f"   ❌ Failed to deactivate test_agent_001: {deactivate_result.get('error')}")
    
    # Try to verify with deactivated agent
    deactivated_result = client.verify_signature("test_agent_001", "data/samples/john_doe_2.png")
    if not deactivated_result['is_verified']:
        print("   ✅ Deactivated agent correctly rejected verification")
    else:
        print("   ⚠️  Deactivated agent still accepted verification")
    
    # Reactivate the agent
    reactivate_result = client.reactivate_agent("test_agent_001")
    if reactivate_result['success']:
        print("   ✅ Reactivated test_agent_001")
    else:
        print(f"   ❌ Failed to reactivate test_agent_001: {reactivate_result.get('error')}")
    
    print("   ✅ Agent deactivation/reactivation passed")
    
    # Test 7: Batch verification
    print("\n7. Testing batch verification...")
    
    batch_requests = [
        {
            "agent_id": "test_agent_001",
            "signature_image": "data/samples/john_doe_2.png",
            "context": {"test": True}
        },
        {
            "agent_id": "test_agent_002", 
            "signature_image": "data/samples/jane_smith_2.png",
            "context": {"test": True}
        }
    ]
    
    batch_data = {"verification_requests": batch_requests}
    response = client.session.post(f"{client.base_url}/batch-verify", json=batch_data)
    batch_result = response.json()
    
    if batch_result['success']:
        print(f"   ✅ Batch verification processed {batch_result['total_processed']} requests")
        for result in batch_result['results']:
            print(f"   - {result['agent_id']}: Verified={result['is_verified']}, "
                  f"Similarity={result['similarity_score']:.3f}")
    else:
        print(f"   ❌ Batch verification failed: {batch_result.get('error')}")
    
    print("\n🎉 All tests completed successfully!")
    print("\nAgentAI Integration Features Demonstrated:")
    print("✅ Health monitoring")
    print("✅ Agent registration and management")
    print("✅ Signature verification")
    print("✅ Agent statistics and monitoring")
    print("✅ Agent activation/deactivation")
    print("✅ Batch processing")
    print("✅ Error handling and logging")


def test_api_endpoints():
    """Test individual API endpoints."""
    print("\n🔍 Testing Individual API Endpoints")
    print("=" * 40)
    
    client = AgentAIClient()
    
    # Test configuration
    config_response = client.session.get(f"{client.base_url}/config")
    config = config_response.json()
    print(f"Current configuration: {json.dumps(config, indent=2)}")
    
    # Test error handling
    print("\nTesting error handling...")
    
    # Test with invalid agent ID
    invalid_result = client.verify_signature("invalid_agent", "data/samples/john_doe_1.png")
    print(f"Invalid agent test: {invalid_result}")
    
    # Test with missing fields
    response = client.session.post(f"{client.base_url}/verify-signature", json={})
    error_result = response.json()
    print(f"Missing fields test: {error_result}")


if __name__ == "__main__":
    print("InklyAI AgentAI Integration Test Suite")
    print("=" * 50)
    
    # Check if API server is running
    try:
        client = AgentAIClient()
        health = client.health_check()
        print(f"✅ API server is running: {health['status']}")
    except requests.exceptions.ConnectionError:
        print("❌ API server is not running. Please start it with:")
        print("   python flask_api.py")
        print("\nThen run this test script again.")
        exit(1)
    
    # Run tests
    test_agentai_integration()
    test_api_endpoints()
    
    print("\n" + "=" * 50)
    print("Test suite completed!")
    print("\nTo start the API server manually:")
    print("   python flask_api.py")
    print("\nAPI will be available at: http://localhost:5000")
    print("API documentation: http://localhost:5000/health")