validate_setup.py

"""
Validation script to test MongoDB & Redis integration.
Run this after setup to verify everything is working.
"""

import os
import sys
import time
from datetime import datetime

# Load environment variables from .env file
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    print("Warning: python-dotenv not installed. Environment variables should be set manually.")
    print("Install with: pip install python-dotenv")

# Color codes for terminal output
GREEN = '\033[92m'
RED = '\033[91m'
YELLOW = '\033[93m'
BLUE = '\033[94m'
RESET = '\033[0m'


def print_header(text):
    """Print section header."""
    print(f"\n{BLUE}{'='*60}{RESET}")
    print(f"{BLUE}{text:^60}{RESET}")
    print(f"{BLUE}{'='*60}{RESET}\n")


def print_success(text):
    """Print success message."""
    print(f"{GREEN}✓{RESET} {text}")


def print_error(text):
    """Print error message."""
    print(f"{RED}✗{RESET} {text}")


def print_warning(text):
    """Print warning message."""
    print(f"{YELLOW}⚠{RESET} {text}")


def print_info(text):
    """Print info message."""
    print(f"  {text}")


def test_environment_variables():
    """Test that required environment variables are set."""
    print_header("Environment Variables")
    
    required_vars = {
        'MONGO_URI': 'MongoDB connection string',
        'REDIS_HOST': 'Redis host',
        'REDIS_PORT': 'Redis port',
        'ADMIN_TOKEN': 'Admin authentication token',
        'IPINFO_TOKEN': 'IPinfo API token'
    }
    
    optional_vars = {
        'REDIS_PASSWORD': 'Redis password',
        'HF_TOKEN': 'HuggingFace API token',
        'MAX_WORKERS': 'Worker thread count'
    }
    
    all_set = True
    
    for var, description in required_vars.items():
        value = os.getenv(var)
        if value:
            print_success(f"{var} is set ({description})")
        else:
            print_error(f"{var} is NOT set ({description})")
            all_set = False
    
    print()
    for var, description in optional_vars.items():
        value = os.getenv(var)
        if value:
            print_info(f"{var} is set ({description})")
        else:
            print_warning(f"{var} is not set ({description}) - optional")
    
    return all_set


def test_mongodb_connection():
    """Test MongoDB connection."""
    print_header("MongoDB Connection")
    
    try:
        # Add src to path
        sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
        
        from src.utils.mongodb_service import get_mongodb_service
        
        mongodb = get_mongodb_service()
        
        if mongodb._client is None:
            print_error("MongoDB client is None")
            return False
        
        # Try to ping
        mongodb._client.admin.command('ping')
        print_success("MongoDB connection successful")
        
        # Test event logging
        success = mongodb.log_event(
            event_type="DASHBOARD_VIEW",
            device_id="test-device-validation",
            user_id="test-user",
            metadata={"test": True, "timestamp": datetime.now().isoformat()}
        )
        
        if success:
            print_success("Test event logged successfully")
        else:
            print_warning("Could not log test event (non-critical)")
        
        return True
        
    except Exception as e:
        print_error(f"MongoDB connection failed: {str(e)}")
        return False


def test_redis_connection():
    """Test Redis connection."""
    print_header("Redis Connection")
    
    try:
        from src.utils.redis_service import get_redis_service
        
        redis = get_redis_service()
        
        if not redis.is_connected():
            print_error("Redis is not connected")
            return False
        
        print_success("Redis connection successful")
        
        # Test rate limiting
        is_allowed, count = redis.check_rate_limit("test-validation-user", max_requests=5)
        if is_allowed:
            print_success(f"Rate limiting functional (count: {count})")
        else:
            print_warning("Rate limit check returned not allowed (might be cached)")
        
        # Test IP gate
        should_log = redis.should_log_ip("test-validation-device")
        print_success(f"IP logging gate functional (should_log: {should_log})")
        
        return True
        
    except Exception as e:
        print_error(f"Redis connection failed: {str(e)}")
        return False


def test_services():
    """Test all service modules."""
    print_header("Service Modules")
    
    try:
        from src.utils.mongodb_service import get_mongodb_service
        print_success("MongoDB service module imported")
    except Exception as e:
        print_error(f"Failed to import MongoDB service: {str(e)}")
        return False
    
    try:
        from src.utils.redis_service import get_redis_service
        print_success("Redis service module imported")
    except Exception as e:
        print_error(f"Failed to import Redis service: {str(e)}")
        return False
    
    try:
        from src.utils.rate_limit_middleware import RateLimitMiddleware
        print_success("Rate limit middleware imported")
    except Exception as e:
        print_error(f"Failed to import rate limit middleware: {str(e)}")
        return False
    
    try:
        from src.utils.task_queue import get_task_queue
        print_success("Task queue module imported")
    except Exception as e:
        print_error(f"Failed to import task queue: {str(e)}")
        return False
    
    try:
        from src.utils.ip_location_service import get_ip_location_service
        print_success("IP location service imported")
    except Exception as e:
        print_error(f"Failed to import IP location service: {str(e)}")
        return False
    
    try:
        from src.utils.admin_endpoints import router
        print_success("Admin endpoints imported")
    except Exception as e:
        print_error(f"Failed to import admin endpoints: {str(e)}")
        return False
    
    return True


def test_api_health():
    """Test API health endpoint."""
    print_header("API Health Check")
    
    try:
        import requests
        
        api_url = os.getenv("API_URL", "http://localhost:7860")
        
        print_info(f"Testing API at: {api_url}")
        
        response = requests.get(f"{api_url}/health", timeout=5)
        
        if response.status_code == 200:
            print_success("API is responding")
            
            data = response.json()
            
            # Check services
            if data.get("mongodb") == "connected":
                print_success("MongoDB status: connected")
            else:
                print_warning(f"MongoDB status: {data.get('mongodb', 'unknown')}")
            
            if data.get("redis") == "connected":
                print_success("Redis status: connected")
            else:
                print_warning(f"Redis status: {data.get('redis', 'unknown')}")
            
            return True
        else:
            print_error(f"API returned status code: {response.status_code}")
            return False
        
    except requests.exceptions.ConnectionError:
        print_warning("API is not running (this is OK if you haven't started it yet)")
        print_info("Start the API with: python api_server.py")
        return None
    except Exception as e:
        print_error(f"Failed to check API health: {str(e)}")
        return False


def main():
    """Run all validation tests."""
    print(f"{BLUE}")
    print("╔═══════════════════════════════════════════════════════════╗")
    print("║  MongoDB & Redis Integration - Validation Script          ║")
    print("╚═══════════════════════════════════════════════════════════╝")
    print(f"{RESET}")
    
    results = {
        "Environment Variables": test_environment_variables(),
        "Service Modules": test_services(),
        "MongoDB Connection": test_mongodb_connection(),
        "Redis Connection": test_redis_connection(),
        "API Health": test_api_health()
    }
    
    # Summary
    print_header("Validation Summary")
    
    passed = sum(1 for v in results.values() if v is True)
    failed = sum(1 for v in results.values() if v is False)
    skipped = sum(1 for v in results.values() if v is None)
    total = len(results)
    
    for test, result in results.items():
        if result is True:
            print_success(f"{test}: PASSED")
        elif result is False:
            print_error(f"{test}: FAILED")
        else:
            print_warning(f"{test}: SKIPPED")
    
    print()
    print(f"Total: {total} tests")
    print(f"{GREEN}Passed: {passed}{RESET}")
    if failed > 0:
        print(f"{RED}Failed: {failed}{RESET}")
    if skipped > 0:
        print(f"{YELLOW}Skipped: {skipped}{RESET}")
    
    print()
    
    if failed == 0 and passed > 0:
        print(f"{GREEN}{'='*60}{RESET}")
        print(f"{GREEN}✓ All critical tests passed! System is ready.{RESET}")
        print(f"{GREEN}{'='*60}{RESET}")
        return 0
    elif failed > 0:
        print(f"{RED}{'='*60}{RESET}")
        print(f"{RED}✗ Some tests failed. Please check the errors above.{RESET}")
        print(f"{RED}{'='*60}{RESET}")
        return 1
    else:
        print(f"{YELLOW}{'='*60}{RESET}")
        print(f"{YELLOW}⚠ No tests could be completed. Check configuration.{RESET}")
        print(f"{YELLOW}{'='*60}{RESET}")
        return 2


if __name__ == "__main__":
    sys.exit(main())