Add MongoDB verification to test.py

test.py

@@ -6,8 +6,11 @@ import os
 import sys
 import time
 from typing import Optional
+from datetime import datetime, timedelta
 
 import requests
+from pymongo import MongoClient
+from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
 
 
 def configure_logging(verbose: bool) -> None:
@@ -103,6 +106,111 @@ def download_result(base_url: str, result_id: str, output_path: str, auth_bearer
 	logging.info("Saved colorized image to: %s", output_path)
 
 
+def verify_mongodb_storage(mongodb_uri: str, db_name: str = "colorization_db", 
+                           test_endpoint: str = None, wait_seconds: int = 5) -> bool:
+	"""
+	Verify that API calls were stored in MongoDB.
+	
+	Args:
+		mongodb_uri: MongoDB connection string
+		db_name: Database name
+		test_endpoint: Specific endpoint to check (optional)
+		wait_seconds: Wait time before checking (to allow async writes)
+	
+	Returns:
+		True if data is found, False otherwise
+	"""
+	logging.info("Waiting %d seconds for MongoDB writes to complete...", wait_seconds)
+	time.sleep(wait_seconds)
+	
+	try:
+		client = MongoClient(mongodb_uri, serverSelectionTimeoutMS=5000)
+		# Test connection
+		client.admin.command('ping')
+		logging.info("✅ Connected to MongoDB successfully")
+		
+		db = client[db_name]
+		
+		# Check api_calls collection
+		api_calls_collection = db["api_calls"]
+		recent_calls = api_calls_collection.find({
+			"timestamp": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
+		}).sort("timestamp", -1).limit(10)
+		
+		api_calls_count = api_calls_collection.count_documents({
+			"timestamp": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
+		})
+		
+		logging.info("Found %d API calls in the last 5 minutes", api_calls_count)
+		
+		if api_calls_count > 0:
+			logging.info("Recent API calls:")
+			for call in list(recent_calls)[:5]:  # Show last 5
+				logging.info("  - %s %s at %s (status: %d)", 
+				           call.get("method", "N/A"),
+				           call.get("endpoint", "N/A"),
+				           call.get("timestamp", "N/A"),
+				           call.get("status_code", 0))
+		
+		# Check image_uploads collection
+		uploads_collection = db["image_uploads"]
+		recent_uploads = uploads_collection.find({
+			"uploaded_at": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
+		}).sort("uploaded_at", -1).limit(5)
+		
+		uploads_count = uploads_collection.count_documents({
+			"uploaded_at": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
+		})
+		
+		logging.info("Found %d image uploads in the last 5 minutes", uploads_count)
+		
+		if uploads_count > 0:
+			logging.info("Recent uploads:")
+			for upload in list(recent_uploads)[:3]:  # Show last 3
+				logging.info("  - Image ID: %s, Size: %d bytes, Uploaded at: %s",
+				           upload.get("image_id", "N/A"),
+				           upload.get("file_size", 0),
+				           upload.get("uploaded_at", "N/A"))
+		
+		# Check colorizations collection
+		colorizations_collection = db["colorizations"]
+		recent_colorizations = colorizations_collection.find({
+			"created_at": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
+		}).sort("created_at", -1).limit(5)
+		
+		colorizations_count = colorizations_collection.count_documents({
+			"created_at": {"$gte": datetime.utcnow() - timedelta(minutes=5)}
+		})
+		
+		logging.info("Found %d colorizations in the last 5 minutes", colorizations_count)
+		
+		if colorizations_count > 0:
+			logging.info("Recent colorizations:")
+			for colorization in list(recent_colorizations)[:3]:  # Show last 3
+				logging.info("  - Result ID: %s, Model: %s, Time: %.2fs, Created at: %s",
+				           colorization.get("result_id", "N/A"),
+				           colorization.get("model_type", "N/A"),
+				           colorization.get("processing_time", 0),
+				           colorization.get("created_at", "N/A"))
+		
+		client.close()
+		
+		# Return True if we found any recent data
+		if api_calls_count > 0 or uploads_count > 0 or colorizations_count > 0:
+			logging.info("✅ MongoDB storage verification PASSED")
+			return True
+		else:
+			logging.warning("⚠️ No recent data found in MongoDB (this might be normal if no API calls were made)")
+			return False
+			
+	except (ConnectionFailure, ServerSelectionTimeoutError) as e:
+		logging.error("❌ Failed to connect to MongoDB: %s", str(e))
+		return False
+	except Exception as e:
+		logging.error("❌ MongoDB verification error: %s", str(e))
+		return False
+
+
 def main() -> int:
 	parser = argparse.ArgumentParser(description="End-to-end test for Colorize API")
 	parser.add_argument("--base-url", type=str, help="API base URL, e.g. https://<space>.hf.space")
@@ -112,6 +220,9 @@ def main() -> int:
 	parser.add_argument("--app-check", type=str, default=os.getenv("APP_CHECK_TOKEN", ""), help="Optional App Check token")
 	parser.add_argument("--skip-wait", action="store_true", help="Skip waiting for model to load")
 	parser.add_argument("--verbose", action="store_true", help="Verbose logging")
+	parser.add_argument("--verify-mongodb", action="store_true", help="Verify MongoDB storage after API calls")
+	parser.add_argument("--mongodb-uri", type=str, default=os.getenv("MONGODB_URI", ""), help="MongoDB connection string for verification")
+	parser.add_argument("--mongodb-db", type=str, default=os.getenv("MONGODB_DB_NAME", "colorization_db"), help="MongoDB database name")
 	args = parser.parse_args()
 
 	configure_logging(args.verbose)
@@ -155,6 +266,18 @@ def main() -> int:
 		return 1
 
 	logging.info("Test workflow completed successfully.")
+	
+	# Verify MongoDB storage if requested
+	if args.verify_mongodb:
+		if not args.mongodb_uri:
+			logging.warning("⚠️ MongoDB URI not provided. Skipping MongoDB verification.")
+			logging.info("Set MONGODB_URI environment variable or use --mongodb-uri flag")
+		else:
+			logging.info("=" * 60)
+			logging.info("Verifying MongoDB storage...")
+			logging.info("=" * 60)
+			verify_mongodb_storage(args.mongodb_uri, args.mongodb_db)
+	
 	return 0