test_api_mongodb.py

"""
Test script to make API calls and verify MongoDB storage
This script makes a simple health check API call and verifies it's stored in MongoDB
"""
import os
import sys
import time
import logging
import requests
from datetime import datetime, timedelta
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def test_api_and_mongodb(base_url: str = "http://localhost:7860", mongodb_uri: str = None):
    """Make API calls and verify they're stored in MongoDB"""
    
    if not mongodb_uri:
        mongodb_uri = os.getenv("MONGODB_URI")
        if not mongodb_uri:
            logger.error("❌ MONGODB_URI not provided!")
            return False
    
    db_name = os.getenv("MONGODB_DB_NAME", "colorization_db")
    
    logger.info("=" * 60)
    logger.info("Testing API calls and MongoDB storage")
    logger.info("=" * 60)
    
    # Step 1: Make a health check API call
    logger.info(f"\n1. Making health check API call to {base_url}/health")
    try:
        response = requests.get(f"{base_url}/health", timeout=10)
        logger.info(f"   Status Code: {response.status_code}")
        if response.ok:
            logger.info(f"   Response: {response.json()}")
        else:
            logger.warning(f"   API returned error: {response.text}")
    except requests.exceptions.RequestException as e:
        logger.error(f"   ❌ Failed to connect to API: {e}")
        logger.info("   Note: Make sure the API server is running")
        logger.info("   You can start it with: uvicorn app.main_sdxl:app --reload")
        return False
    
    # Step 2: Wait a bit for MongoDB write
    logger.info("\n2. Waiting 3 seconds for MongoDB write to complete...")
    time.sleep(3)
    
    # Step 3: Verify MongoDB storage
    logger.info("\n3. Verifying MongoDB storage...")
    try:
        client = MongoClient(mongodb_uri, serverSelectionTimeoutMS=5000)
        client.admin.command('ping')
        logger.info("   ✅ Connected to MongoDB")
        
        db = client[db_name]
        
        # Check api_calls collection
        api_calls = db["api_calls"]
        
        # Count total calls
        total_calls = api_calls.count_documents({})
        logger.info(f"   Total API calls in database: {total_calls}")
        
        # Check recent calls (last 5 minutes)
        recent_calls = list(api_calls.find({
            "timestamp": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
        }).sort("timestamp", -1).limit(10))
        
        logger.info(f"   Recent API calls (last 5 minutes): {len(recent_calls)}")
        
        if recent_calls:
            logger.info("\n   Recent API call details:")
            for call in recent_calls[:5]:
                logger.info(f"      - {call.get('method')} {call.get('endpoint')}")
                logger.info(f"        Status: {call.get('status_code')}")
                logger.info(f"        Timestamp: {call.get('timestamp')}")
                logger.info(f"        IP: {call.get('ip_address')}")
                logger.info("")
            
            logger.info("   ✅ MongoDB storage verification PASSED!")
            logger.info("   ✅ API calls are being stored in MongoDB with timestamps!")
            client.close()
            return True
        else:
            logger.warning("   ⚠️ No recent API calls found in MongoDB")
            logger.info("   This might mean:")
            logger.info("     1. The API server is not logging to MongoDB")
            logger.info("     2. MONGODB_URI is not set in the API server environment")
            logger.info("     3. The write hasn't completed yet (try waiting longer)")
            client.close()
            return False
            
    except (ConnectionFailure, ServerSelectionTimeoutError) as e:
        logger.error(f"   ❌ Failed to connect to MongoDB: {e}")
        return False
    except Exception as e:
        logger.error(f"   ❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Test API calls and MongoDB storage")
    parser.add_argument("--base-url", type=str, default="http://localhost:7860",
                       help="API base URL (default: http://localhost:7860)")
    parser.add_argument("--mongodb-uri", type=str, default=os.getenv("MONGODB_URI", ""),
                       help="MongoDB connection string")
    
    args = parser.parse_args()
    
    if not args.mongodb_uri:
        logger.error("MongoDB URI required! Set MONGODB_URI environment variable or use --mongodb-uri")
        sys.exit(1)
    
    success = test_api_and_mongodb(args.base_url, args.mongodb_uri)
    sys.exit(0 if success else 1)