""" 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)