File size: 8,453 Bytes
96a6d41 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
"""
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)
|