File size: 4,949 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 |
"""
Test script to make API calls and verify MongoDB storage
This script makes a simple health check API call and verifies it's stored in MongoDB
"""
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_and_mongodb(base_url: str = "http://localhost:7860", mongodb_uri: str = None):
"""Make API calls and verify they're stored in MongoDB"""
if not mongodb_uri:
mongodb_uri = os.getenv("MONGODB_URI")
if not mongodb_uri:
logger.error("❌ MONGODB_URI not provided!")
return False
db_name = os.getenv("MONGODB_DB_NAME", "colorization_db")
logger.info("=" * 60)
logger.info("Testing API calls and MongoDB storage")
logger.info("=" * 60)
# Step 1: Make a health check API call
logger.info(f"\n1. Making health check API call to {base_url}/health")
try:
response = requests.get(f"{base_url}/health", timeout=10)
logger.info(f" Status Code: {response.status_code}")
if response.ok:
logger.info(f" Response: {response.json()}")
else:
logger.warning(f" API returned error: {response.text}")
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(" You can start it with: uvicorn app.main_sdxl:app --reload")
return False
# Step 2: Wait a bit for MongoDB write
logger.info("\n2. Waiting 3 seconds for MongoDB write to complete...")
time.sleep(3)
# Step 3: Verify MongoDB storage
logger.info("\n3. 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"]
# Count total calls
total_calls = api_calls.count_documents({})
logger.info(f" Total API calls in database: {total_calls}")
# Check recent calls (last 5 minutes)
recent_calls = list(api_calls.find({
"timestamp": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
}).sort("timestamp", -1).limit(10))
logger.info(f" Recent API calls (last 5 minutes): {len(recent_calls)}")
if recent_calls:
logger.info("\n Recent API call details:")
for call in recent_calls[:5]:
logger.info(f" - {call.get('method')} {call.get('endpoint')}")
logger.info(f" Status: {call.get('status_code')}")
logger.info(f" Timestamp: {call.get('timestamp')}")
logger.info(f" IP: {call.get('ip_address')}")
logger.info("")
logger.info(" ✅ MongoDB storage verification PASSED!")
logger.info(" ✅ API calls are being stored in MongoDB with timestamps!")
client.close()
return True
else:
logger.warning(" ⚠️ No recent API calls found in MongoDB")
logger.info(" This might mean:")
logger.info(" 1. The API server is not logging to MongoDB")
logger.info(" 2. MONGODB_URI is not set in the API server environment")
logger.info(" 3. The write hasn't completed yet (try waiting longer)")
client.close()
return False
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 calls and MongoDB storage")
parser.add_argument("--base-url", type=str, default="http://localhost:7860",
help="API base URL (default: http://localhost:7860)")
parser.add_argument("--mongodb-uri", type=str, default=os.getenv("MONGODB_URI", ""),
help="MongoDB connection string")
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_and_mongodb(args.base_url, args.mongodb_uri)
sys.exit(0 if success else 1)
|