mcp-vultr/test_improvements.py

#!/usr/bin/env python3
"""Test our improvements to the Vultr DNS MCP package."""

import ipaddress
from mcp_vultr.server import (
    VultrAPIError, 
    VultrAuthError, 
    VultrRateLimitError, 
    VultrResourceNotFoundError, 
    VultrValidationError
)

def test_exceptions():
    """Test our enhanced exception classes."""
    print("🧪 Testing Enhanced Exception Classes:")
    
    # Test basic exception
    error = VultrAPIError(500, 'Test error')
    print(f"✅ VultrAPIError: {error}")
    
    # Test specific exceptions
    auth_error = VultrAuthError(401, 'Auth failed')
    print(f"✅ VultrAuthError: {auth_error}")
    
    rate_error = VultrRateLimitError(429, 'Too many requests')
    print(f"✅ VultrRateLimitError: {rate_error}")
    
    not_found_error = VultrResourceNotFoundError(404, 'Not found')
    print(f"✅ VultrResourceNotFoundError: {not_found_error}")
    
    validation_error = VultrValidationError(400, 'Bad request')
    print(f"✅ VultrValidationError: {validation_error}")
    
    # Test inheritance
    print(f"✅ VultrAuthError is VultrAPIError: {isinstance(auth_error, VultrAPIError)}")
    print(f"✅ VultrAuthError status_code: {auth_error.status_code}")
    print(f"✅ VultrAuthError message: {auth_error.message}")
    print()

def test_ipv6_validation():
    """Test enhanced IPv6 validation."""
    print("🧪 Testing Enhanced IPv6 Validation:")
    
    test_addresses = [
        ('2001:db8::1', 'Standard format'),
        ('::1', 'Loopback'),
        ('::ffff:192.0.2.1', 'IPv4-mapped'),
        ('2001:0db8:0000:0000:0000:0000:0000:0001', 'Full format'),
        ('fe80::1', 'Link-local'),
        ('invalid::address::bad', 'Invalid (multiple ::)'),
        ('gggg::1', 'Invalid hex'),
        ('192.168.1.1', 'IPv4 (should fail for IPv6)'),
    ]
    
    for addr, description in test_addresses:
        try:
            ipv6_addr = ipaddress.IPv6Address(addr)
            print(f"✅ {addr:<35} -> {ipv6_addr.compressed} ({description})")
            
            # Test our enhanced features
            if ipv6_addr.ipv4_mapped:
                print(f"   📝 IPv4-mapped: {ipv6_addr.ipv4_mapped}")
            if ipv6_addr.is_loopback:
                print(f"   🔄 Loopback address")
            if ipv6_addr.is_link_local:
                print(f"   🔗 Link-local address")
            if ipv6_addr.is_private:
                print(f"   🔒 Private address")
            if ipv6_addr.compressed != addr:
                print(f"   💡 Could compress to: {ipv6_addr.compressed}")
                
        except ipaddress.AddressValueError as e:
            print(f"❌ {addr:<35} -> Invalid: {e} ({description})")
    
    print()

def test_mcp_server_creation():
    """Test that we can create an MCP server (without API key)."""
    print("🧪 Testing MCP Server Creation:")
    
    try:
        from mcp_vultr.server import create_mcp_server
        
        # This should fail without API key (expected)
        try:
            server = create_mcp_server(None)
            print("❌ Should have failed without API key")
        except ValueError as e:
            print(f"✅ Correctly failed without API key: {e}")
        
        # Test with mock API key
        server = create_mcp_server("mock-api-key")
        print(f"✅ MCP server created successfully: {type(server)}")
        print(f"✅ Server has tools: {hasattr(server, '_tools')}")
        print(f"✅ Server has resources: {hasattr(server, '_resources')}")
        
    except Exception as e:
        print(f"❌ Failed to create MCP server: {e}")
    
    print()

def simulate_domain_query():
    """Simulate what a domain query would look like."""
    print("🧪 Simulating Domain Query (Mock Response):")
    
    # This is what the response would look like with a real API key
    mock_domains = [
        {
            "domain": "example.com",
            "date_created": "2024-01-01T00:00:00Z",
            "dns_sec": "disabled"
        },
        {
            "domain": "test-site.net", 
            "date_created": "2024-06-15T12:30:00Z",
            "dns_sec": "enabled"
        }
    ]
    
    print("📋 Available domains (mock data):")
    for domain in mock_domains:
        name = domain['domain']
        created = domain['date_created']
        dnssec = domain['dns_sec']
        print(f"  • {name:<20} (created: {created}, DNSSEC: {dnssec})")
    
    print()
    print("💡 To query real domains, set VULTR_API_KEY and run:")
    print("   export VULTR_API_KEY='your-api-key'")
    print("   uv run vultr-dns-mcp domains list")
    print()

if __name__ == "__main__":
    print("🚀 Testing Vultr DNS MCP Improvements\n")
    print("=" * 50)
    
    test_exceptions()
    test_ipv6_validation()
    test_mcp_server_creation()
    simulate_domain_query()
    
    print("✅ All improvement tests completed successfully!")
    print("\n🎉 Ready to use with: uv run vultr-dns-mcp --help")