test_with_image.py

"""
Test script to test API with Descratch.png image and verify MongoDB storage
"""
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_with_image_and_mongodb(image_path: str, base_url: str = "http://localhost:7860", 
                                    mongodb_uri: str = None, app_check_token: str = None):
    """Test API with image and verify MongoDB storage"""
    
    if not mongodb_uri:
        mongodb_uri = os.getenv("MONGODB_URI")
        if not mongodb_uri:
            logger.error("❌ MONGODB_URI not provided!")
            return False
    
    if not os.path.exists(image_path):
        logger.error(f"❌ Image not found: {image_path}")
        return False
    
    db_name = os.getenv("MONGODB_DB_NAME", "colorization_db")
    
    logger.info("=" * 80)
    logger.info("Testing API with Image and MongoDB Storage")
    logger.info("=" * 80)
    logger.info(f"Image: {image_path}")
    logger.info(f"API URL: {base_url}")
    logger.info(f"Database: {db_name}\n")
    
    # Step 1: Health check
    logger.info("1. Testing /health endpoint...")
    try:
        response = requests.get(f"{base_url}/health", timeout=10)
        if response.ok:
            logger.info(f"   ✅ Health check passed: {response.json()}")
        else:
            logger.warning(f"   ⚠️ Health check returned: {response.status_code}")
    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("   Start with: uvicorn app.main_sdxl:app --reload")
        return False
    
    # Step 2: Upload image
    logger.info("\n2. Uploading image...")
    upload_url = f"{base_url}/upload"
    headers = {}
    if app_check_token:
        headers["X-Firebase-AppCheck"] = app_check_token
    
    try:
        with open(image_path, "rb") as f:
            files = {"file": (os.path.basename(image_path), f, "image/png")}
            response = requests.post(upload_url, files=files, headers=headers, timeout=120)
        
        if response.ok:
            upload_data = response.json()
            logger.info(f"   ✅ Upload successful!")
            logger.info(f"   Image ID: {upload_data.get('image_id')}")
            image_id = upload_data.get('image_id')
        else:
            logger.error(f"   ❌ Upload failed: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        logger.error(f"   ❌ Upload error: {e}")
        return False
    
    # Step 3: Colorize image
    logger.info("\n3. Colorizing image...")
    colorize_url = f"{base_url}/colorize"
    
    try:
        with open(image_path, "rb") as f:
            files = {"file": (os.path.basename(image_path), f, "image/png")}
            response = requests.post(colorize_url, files=files, headers=headers, timeout=900)
        
        if response.ok:
            colorize_data = response.json()
            logger.info(f"   ✅ Colorization successful!")
            logger.info(f"   Result ID: {colorize_data.get('result_id')}")
            result_id = colorize_data.get('result_id')
        else:
            logger.error(f"   ❌ Colorization failed: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        logger.error(f"   ❌ Colorization error: {e}")
        return False
    
    # Step 4: Wait for MongoDB writes
    logger.info("\n4. Waiting 5 seconds for MongoDB writes to complete...")
    time.sleep(5)
    
    # Step 5: Verify MongoDB storage
    logger.info("\n5. 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"]
        recent_calls = list(api_calls.find({
            "timestamp": {"$gte": datetime.utcnow() - timedelta(minutes=10)}
        }).sort("timestamp", -1))
        
        logger.info(f"\n   📊 Found {len(recent_calls)} API calls in last 10 minutes:")
        for call in recent_calls[:5]:
            logger.info(f"      - {call.get('method')} {call.get('endpoint')} "
                      f"(Status: {call.get('status_code')}) at {call.get('timestamp')}")
        
        # Check image_uploads collection
        image_uploads = db["image_uploads"]
        recent_uploads = list(image_uploads.find({
            "uploaded_at": {"$gte": datetime.utcnow() - timedelta(minutes=10)}
        }).sort("uploaded_at", -1))
        
        logger.info(f"\n   📊 Found {len(recent_uploads)} image uploads in last 10 minutes:")
        for upload in recent_uploads[:3]:
            logger.info(f"      - Image ID: {upload.get('image_id')}")
            logger.info(f"        Filename: {upload.get('filename')}")
            logger.info(f"        Size: {upload.get('file_size')} bytes")
            logger.info(f"        Uploaded at: {upload.get('uploaded_at')}")
        
        # Check colorizations collection
        colorizations = db["colorizations"]
        recent_colorizations = list(colorizations.find({
            "created_at": {"$gte": datetime.utcnow() - timedelta(minutes=10)}
        }).sort("created_at", -1))
        
        logger.info(f"\n   📊 Found {len(recent_colorizations)} colorizations in last 10 minutes:")
        for colorization in recent_colorizations[:3]:
            logger.info(f"      - Result ID: {colorization.get('result_id')}")
            logger.info(f"        Model: {colorization.get('model_type')}")
            logger.info(f"        Processing time: {colorization.get('processing_time')}s")
            logger.info(f"        Created at: {colorization.get('created_at')}")
        
        # Summary
        logger.info("\n" + "=" * 80)
        logger.info("TEST SUMMARY")
        logger.info("=" * 80)
        
        if len(recent_calls) > 0 and len(recent_uploads) > 0 and len(recent_colorizations) > 0:
            logger.info("✅ SUCCESS: All data is being stored in MongoDB!")
            logger.info("✅ API calls are logged with timestamps")
            logger.info("✅ Image uploads are logged with metadata")
            logger.info("✅ Colorizations are logged with processing details")
        else:
            logger.warning("⚠️ Some data might be missing:")
            logger.warning(f"   API calls: {len(recent_calls)}")
            logger.warning(f"   Image uploads: {len(recent_uploads)}")
            logger.warning(f"   Colorizations: {len(recent_colorizations)}")
        
        client.close()
        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__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Test API with image and verify MongoDB")
    parser.add_argument("--image", type=str, default="../Descratch.png",
                       help="Path to image file")
    parser.add_argument("--base-url", type=str, default="http://localhost:7860",
                       help="API base URL")
    parser.add_argument("--mongodb-uri", type=str, default=os.getenv("MONGODB_URI", ""),
                       help="MongoDB connection string")
    parser.add_argument("--app-check", type=str, default=os.getenv("APP_CHECK_TOKEN", ""),
                       help="Firebase App Check token (optional)")
    
    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_with_image_and_mongodb(
        args.image, 
        args.base_url, 
        args.mongodb_uri,
        args.app_check if args.app_check else None
    )
    sys.exit(0 if success else 1)