test_mongodb.py

"""
Simple MongoDB verification test script
This script tests if data is being stored in MongoDB after API calls.
"""
import os
import sys
import logging
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_mongodb_connection():
    """Test MongoDB connection and verify collections exist"""
    mongodb_uri = os.getenv("MONGODB_URI")
    if not mongodb_uri:
        logger.error("❌ MONGODB_URI environment variable not set!")
        logger.info("Please set MONGODB_URI environment variable")
        return False
    
    db_name = os.getenv("MONGODB_DB_NAME", "colorization_db")
    
    try:
        logger.info("Connecting to MongoDB...")
        client = MongoClient(mongodb_uri, serverSelectionTimeoutMS=5000)
        
        # Test connection
        client.admin.command('ping')
        logger.info("✅ Connected to MongoDB successfully!")
        
        db = client[db_name]
        logger.info(f"Using database: {db_name}")
        
        # List all collections
        collections = db.list_collection_names()
        logger.info(f"Collections found: {collections}")
        
        # Check each collection
        collections_to_check = ["api_calls", "image_uploads", "colorizations"]
        
        for collection_name in collections_to_check:
            collection = db[collection_name]
            total_count = collection.count_documents({})
            recent_count = collection.count_documents({
                "timestamp" if collection_name == "api_calls" else 
                ("uploaded_at" if collection_name == "image_uploads" else "created_at"): {
                    "$gte": datetime.utcnow() - timedelta(hours=24)
                }
            })
            
            logger.info(f"\n📊 Collection: {collection_name}")
            logger.info(f"   Total documents: {total_count}")
            logger.info(f"   Documents in last 24 hours: {recent_count}")
            
            # Show sample document
            sample = collection.find_one(sort=[("_id", -1)])
            if sample:
                logger.info(f"   Latest document ID: {sample.get('_id')}")
                if collection_name == "api_calls":
                    logger.info(f"   Latest endpoint: {sample.get('endpoint')} at {sample.get('timestamp')}")
                elif collection_name == "image_uploads":
                    logger.info(f"   Latest image: {sample.get('image_id')} at {sample.get('uploaded_at')}")
                elif collection_name == "colorizations":
                    logger.info(f"   Latest result: {sample.get('result_id')} at {sample.get('created_at')}")
        
        # Show recent API calls
        api_calls = db["api_calls"]
        recent_calls = list(api_calls.find({
            "timestamp": {"$gte": datetime.utcnow() - timedelta(minutes=30)}
        }).sort("timestamp", -1).limit(10))
        
        if recent_calls:
            logger.info("\n📝 Recent API calls (last 30 minutes):")
            for call in recent_calls:
                logger.info(f"   {call.get('method')} {call.get('endpoint')} - "
                          f"Status: {call.get('status_code')} - "
                          f"Time: {call.get('timestamp')}")
        else:
            logger.info("\n📝 No API calls in the last 30 minutes")
        
        client.close()
        logger.info("\n✅ MongoDB test completed successfully!")
        return True
        
    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__":
    success = test_mongodb_connection()
    sys.exit(0 if success else 1)