test_mongodb_direct.py

"""
Direct MongoDB test - Insert test data and verify storage
This test directly writes to MongoDB to verify the connection and storage works
"""
import os
import sys
import logging
from datetime import datetime
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_direct():
    """Test MongoDB by directly inserting test data"""
    mongodb_uri = os.getenv("MONGODB_URI")
    if not mongodb_uri:
        logger.error("❌ MONGODB_URI environment variable not set!")
        logger.info("Set it with: $env:MONGODB_URI='your_connection_string'")
        return False
    
    db_name = os.getenv("MONGODB_DB_NAME", "colorization_db")
    
    logger.info("=" * 60)
    logger.info("Direct MongoDB Storage Test")
    logger.info("=" * 60)
    
    try:
        # Connect to MongoDB
        logger.info("\n1. Connecting to MongoDB...")
        client = MongoClient(mongodb_uri, serverSelectionTimeoutMS=10000)
        client.admin.command('ping')
        logger.info("   ✅ Connected successfully!")
        
        db = client[db_name]
        logger.info(f"   Using database: {db_name}")
        
        # Test 1: Insert into api_calls collection
        logger.info("\n2. Testing api_calls collection...")
        api_calls = db["api_calls"]
        
        test_api_call = {
            "endpoint": "/health",
            "method": "GET",
            "status_code": 200,
            "timestamp": datetime.utcnow(),
            "request_data": {},
            "response_data": {"status": "healthy", "model_loaded": True},
            "error": None,
            "user_id": None,
            "ip_address": "127.0.0.1",
            "test": True  # Mark as test data
        }
        
        result = api_calls.insert_one(test_api_call)
        logger.info(f"   ✅ Inserted test API call with ID: {result.inserted_id}")
        
        # Verify it was stored
        stored = api_calls.find_one({"_id": result.inserted_id})
        if stored:
            logger.info(f"   ✅ Verified: API call stored successfully")
            logger.info(f"      Endpoint: {stored.get('endpoint')}")
            logger.info(f"      Timestamp: {stored.get('timestamp')}")
        else:
            logger.error("   ❌ Failed to retrieve stored document")
            return False
        
        # Test 2: Insert into image_uploads collection
        logger.info("\n3. Testing image_uploads collection...")
        image_uploads = db["image_uploads"]
        
        test_upload = {
            "image_id": "test-image-123",
            "filename": "test_image.jpg",
            "file_size": 102400,
            "content_type": "image/jpeg",
            "uploaded_at": datetime.utcnow(),
            "user_id": None,
            "ip_address": "127.0.0.1",
            "test": True
        }
        
        result = image_uploads.insert_one(test_upload)
        logger.info(f"   ✅ Inserted test upload with ID: {result.inserted_id}")
        
        stored = image_uploads.find_one({"_id": result.inserted_id})
        if stored:
            logger.info(f"   ✅ Verified: Image upload stored successfully")
            logger.info(f"      Image ID: {stored.get('image_id')}")
            logger.info(f"      Uploaded at: {stored.get('uploaded_at')}")
        
        # Test 3: Insert into colorizations collection
        logger.info("\n4. Testing colorizations collection...")
        colorizations = db["colorizations"]
        
        test_colorization = {
            "result_id": "test-result-456",
            "image_id": "test-image-123",
            "prompt": "Test colorization",
            "model_type": "test",
            "processing_time": 1.23,
            "created_at": datetime.utcnow(),
            "user_id": None,
            "ip_address": "127.0.0.1",
            "test": True
        }
        
        result = colorizations.insert_one(test_colorization)
        logger.info(f"   ✅ Inserted test colorization with ID: {result.inserted_id}")
        
        stored = colorizations.find_one({"_id": result.inserted_id})
        if stored:
            logger.info(f"   ✅ Verified: Colorization stored successfully")
            logger.info(f"      Result ID: {stored.get('result_id')}")
            logger.info(f"      Created at: {stored.get('created_at')}")
        
        # Summary
        logger.info("\n" + "=" * 60)
        logger.info("Test Summary")
        logger.info("=" * 60)
        
        total_api_calls = api_calls.count_documents({})
        total_uploads = image_uploads.count_documents({})
        total_colorizations = colorizations.count_documents({})
        
        logger.info(f"Total documents in api_calls: {total_api_calls}")
        logger.info(f"Total documents in image_uploads: {total_uploads}")
        logger.info(f"Total documents in colorizations: {total_colorizations}")
        
        logger.info("\n✅ All MongoDB storage tests PASSED!")
        logger.info("✅ Collections are created automatically")
        logger.info("✅ Data is being stored with timestamps")
        logger.info("✅ MongoDB integration is working correctly!")
        
        # Optional: Clean up test data
        logger.info("\n5. Cleaning up test data...")
        api_calls.delete_many({"test": True})
        image_uploads.delete_many({"test": True})
        colorizations.delete_many({"test": True})
        logger.info("   ✅ Test data cleaned up")
        
        client.close()
        return True
        
    except (ConnectionFailure, ServerSelectionTimeoutError) as e:
        logger.error(f"❌ Failed to connect to MongoDB: {e}")
        logger.info("\nTroubleshooting:")
        logger.info("1. Check your MongoDB connection string")
        logger.info("2. Verify network connectivity")
        logger.info("3. Check if MongoDB allows connections from your IP")
        return False
    except Exception as e:
        logger.error(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = test_mongodb_direct()
    sys.exit(0 if success else 1)