Aammyy/mirage
0
1"""2Testing and Validation Suite for Mirage AI Avatar System3Tests end-to-end functionality, latency, and performance4"""5import asyncio6import time7import aiohttp8import json9import numpy as np10import cv211import logging12from pathlib import Path13import subprocess14import psutil15from typing import Dict, Any, List16 17logging.basicConfig(level=logging.INFO)18logger = logging.getLogger(__name__)19 20class MirageSystemTester:21 """Comprehensive testing suite for the AI avatar system"""22 23 def __init__(self, base_url: str = "http://localhost:7860"):24 self.base_url = base_url25 self.session = None26 self.test_results = {}27 28 async def __aenter__(self):29 self.session = aiohttp.ClientSession()30 return self31 32 async def __aexit__(self, exc_type, exc_val, exc_tb):33 if self.session:34 await self.session.close()35 36 async def test_health_endpoint(self) -> bool:37 """Test basic health endpoint"""38 try:39 async with self.session.get(f"{self.base_url}/health") as response:40 data = await response.json()41 42 success = (43 response.status == 200 and44 data.get("status") == "ok" and45 data.get("system") == "real-time-ai-avatar"46 )47 48 self.test_results["health"] = {49 "success": success,50 "status": response.status,51 "data": data52 }53 54 logger.info(f"Health check: {'✅ PASS' if success else '❌ FAIL'}")55 return success56 57 except Exception as e:58 logger.error(f"Health check failed: {e}")59 self.test_results["health"] = {"success": False, "error": str(e)}60 return False61 62 async def test_pipeline_initialization(self) -> bool:63 """Test AI pipeline initialization"""64 try:65 start_time = time.time()66 async with self.session.post(f"{self.base_url}/initialize") as response:67 data = await response.json()68 init_time = time.time() - start_time69 70 success = (71 response.status == 200 and72 data.get("status") in ["success", "already_initialized"]73 )74 75 self.test_results["initialization"] = {76 "success": success,77 "status": response.status,78 "data": data,79 "init_time_seconds": init_time80 }81 82 logger.info(f"Pipeline init: {'✅ PASS' if success else '❌ FAIL'} ({init_time:.1f}s)")83 return success84 85 except Exception as e:86 logger.error(f"Pipeline initialization failed: {e}")87 self.test_results["initialization"] = {"success": False, "error": str(e)}88 return False89 90 async def test_reference_image_upload(self) -> bool:91 """Test reference image upload functionality"""92 try:93 # Create a test image94 test_image = np.zeros((512, 512, 3), dtype=np.uint8)95 cv2.circle(test_image, (256, 200), 50, (255, 255, 255), -1) # Face-like circle96 cv2.circle(test_image, (230, 180), 10, (0, 0, 0), -1) # Eye97 cv2.circle(test_image, (280, 180), 10, (0, 0, 0), -1) # Eye98 cv2.ellipse(test_image, (256, 220), (20, 10), 0, 0, 180, (0, 0, 0), 2) # Mouth99 100 # Encode as JPEG101 _, encoded = cv2.imencode('.jpg', test_image)102 image_data = encoded.tobytes()103 104 # Upload test image105 form_data = aiohttp.FormData()106 form_data.add_field('file', image_data, filename='test_face.jpg', content_type='image/jpeg')107 108 async with self.session.post(f"{self.base_url}/set_reference", data=form_data) as response:109 data = await response.json()110 111 success = (112 response.status == 200 and113 data.get("status") == "success"114 )115 116 self.test_results["reference_upload"] = {117 "success": success,118 "status": response.status,119 "data": data120 }121 122 logger.info(f"Reference upload: {'✅ PASS' if success else '❌ FAIL'}")123 return success124 125 except Exception as e:126 logger.error(f"Reference image upload failed: {e}")127 self.test_results["reference_upload"] = {"success": False, "error": str(e)}128 return False129 130 async def test_websocket_connections(self) -> bool:131 """Test WebSocket connections for audio and video"""132 try:133 import websockets134 135 # Test audio WebSocket136 audio_success = await self._test_websocket_endpoint("/audio")137 138 # Test video WebSocket139 video_success = await self._test_websocket_endpoint("/video")140 141 success = audio_success and video_success142 143 self.test_results["websockets"] = {144 "success": success,145 "audio_success": audio_success,146 "video_success": video_success147 }148 149 logger.info(f"WebSocket connections: {'✅ PASS' if success else '❌ FAIL'}")150 return success151 152 except Exception as e:153 logger.error(f"WebSocket test failed: {e}")154 self.test_results["websockets"] = {"success": False, "error": str(e)}155 return False156 157 async def _test_websocket_endpoint(self, endpoint: str) -> bool:158 """Test a specific WebSocket endpoint"""159 try:160 import websockets161 162 ws_url = self.base_url.replace("http://", "ws://") + endpoint163 164 async with websockets.connect(ws_url) as websocket:165 # Send test data166 if endpoint == "/audio":167 # Send 160ms of silence (16kHz, 16-bit)168 test_audio = np.zeros(int(16000 * 0.160), dtype=np.int16)169 await websocket.send(test_audio.tobytes())170 else: # video171 # Send a small test JPEG172 test_frame = np.zeros((256, 256, 3), dtype=np.uint8)173 _, encoded = cv2.imencode('.jpg', test_frame, [cv2.IMWRITE_JPEG_QUALITY, 50])174 await websocket.send(encoded.tobytes())175 176 # Wait for response177 response = await asyncio.wait_for(websocket.recv(), timeout=5.0)178 return len(response) > 0179 180 except Exception as e:181 logger.error(f"WebSocket {endpoint} test failed: {e}")182 return False183 184 async def test_performance_metrics(self) -> bool:185 """Test performance metrics endpoint"""186 try:187 async with self.session.get(f"{self.base_url}/pipeline_status") as response:188 data = await response.json()189 190 success = response.status == 200 and data.get("initialized", False)191 192 self.test_results["performance_metrics"] = {193 "success": success,194 "status": response.status,195 "data": data196 }197 198 if success:199 stats = data.get("stats", {})200 logger.info(f"Performance metrics: ✅ PASS")201 logger.info(f" GPU Memory: {stats.get('gpu_memory_used', 0):.1f} GB")202 logger.info(f" Video FPS: {stats.get('video_fps', 0):.1f}")203 logger.info(f" Avg Latency: {stats.get('avg_video_latency_ms', 0):.1f} ms")204 else:205 logger.info("Performance metrics: ❌ FAIL")206 207 return success208 209 except Exception as e:210 logger.error(f"Performance metrics test failed: {e}")211 self.test_results["performance_metrics"] = {"success": False, "error": str(e)}212 return False213 214 async def test_latency_benchmark(self) -> Dict[str, float]:215 """Benchmark system latency"""216 latencies = []217 218 try:219 # Warm up220 for _ in range(5):221 start_time = time.time()222 async with self.session.get(f"{self.base_url}/health") as response:223 await response.json()224 latencies.append((time.time() - start_time) * 1000)225 226 # Actual benchmark227 latencies = []228 for _ in range(20):229 start_time = time.time()230 async with self.session.get(f"{self.base_url}/pipeline_status") as response:231 await response.json()232 latencies.append((time.time() - start_time) * 1000)233 234 results = {235 "avg_latency_ms": np.mean(latencies),236 "min_latency_ms": np.min(latencies),237 "max_latency_ms": np.max(latencies),238 "p95_latency_ms": np.percentile(latencies, 95),239 "p99_latency_ms": np.percentile(latencies, 99)240 }241 242 self.test_results["latency_benchmark"] = results243 244 logger.info("Latency benchmark results:")245 logger.info(f" Average: {results['avg_latency_ms']:.1f} ms")246 logger.info(f" P95: {results['p95_latency_ms']:.1f} ms")247 logger.info(f" P99: {results['p99_latency_ms']:.1f} ms")248 249 return results250 251 except Exception as e:252 logger.error(f"Latency benchmark failed: {e}")253 return {}254 255 def test_system_requirements(self) -> Dict[str, Any]:256 """Test system requirements and capabilities"""257 results = {}258 259 try:260 # Check GPU availability261 try:262 import torch263 results["gpu_available"] = torch.cuda.is_available()264 if torch.cuda.is_available():265 results["gpu_name"] = torch.cuda.get_device_name(0)266 results["gpu_memory_gb"] = torch.cuda.get_device_properties(0).total_memory / 1024**3267 results["cuda_version"] = torch.version.cuda268 except ImportError:269 results["gpu_available"] = False270 271 # Check system resources272 memory = psutil.virtual_memory()273 results["system_memory_gb"] = memory.total / 1024**3274 results["cpu_count"] = psutil.cpu_count()275 276 # Check disk space277 disk = psutil.disk_usage('/')278 results["disk_free_gb"] = disk.free / 1024**3279 280 # Check required packages281 required_packages = [282 "torch", "torchvision", "torchaudio", "opencv-python", 283 "numpy", "fastapi", "websockets"284 ]285 286 missing_packages = []287 for package in required_packages:288 try:289 __import__(package.replace("-", "_"))290 except ImportError:291 missing_packages.append(package)292 293 results["missing_packages"] = missing_packages294 results["requirements_met"] = len(missing_packages) == 0295 296 self.test_results["system_requirements"] = results297 298 logger.info("System requirements:")299 logger.info(f" GPU: {'✅' if results['gpu_available'] else '❌'}")300 logger.info(f" Memory: {results['system_memory_gb']:.1f} GB")301 logger.info(f" CPU: {results['cpu_count']} cores")302 logger.info(f" Packages: {'✅' if results['requirements_met'] else '❌'}")303 304 return results305 306 except Exception as e:307 logger.error(f"System requirements check failed: {e}")308 return {"error": str(e)}309 310 async def run_comprehensive_test(self) -> Dict[str, Any]:311 """Run all tests and return comprehensive results"""312 logger.info("🧪 Starting comprehensive system test...")313 314 # System requirements (runs first, no server needed)315 self.test_system_requirements()316 317 # Server-dependent tests318 tests = [319 ("Health Check", self.test_health_endpoint()),320 ("Pipeline Initialization", self.test_pipeline_initialization()),321 ("Reference Image Upload", self.test_reference_image_upload()),322 ("WebSocket Connections", self.test_websocket_connections()),323 ("Performance Metrics", self.test_performance_metrics()),324 ]325 326 # Run tests sequentially327 for test_name, test_coro in tests:328 logger.info(f"Running: {test_name}...")329 try:330 result = await test_coro331 if not result:332 logger.warning(f"{test_name} failed - may affect subsequent tests")333 except Exception as e:334 logger.error(f"{test_name} threw exception: {e}")335 336 # Latency benchmark (runs last)337 logger.info("Running latency benchmark...")338 await self.test_latency_benchmark()339 340 # Calculate overall success rate341 successful_tests = sum(1 for result in self.test_results.values() 342 if isinstance(result, dict) and result.get("success", False))343 total_tests = len([r for r in self.test_results.values() if isinstance(r, dict) and "success" in r])344 345 overall_success = successful_tests / max(total_tests, 1) >= 0.8 # 80% success rate346 347 summary = {348 "overall_success": overall_success,349 "successful_tests": successful_tests,350 "total_tests": total_tests,351 "success_rate": successful_tests / max(total_tests, 1),352 "detailed_results": self.test_results353 }354 355 logger.info(f"🏁 Test completed: {successful_tests}/{total_tests} tests passed")356 logger.info(f"Overall result: {'✅ PASS' if overall_success else '❌ FAIL'}")357 358 return summary359 360async def main():361 """Main test runner"""362 import sys363 364 base_url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:7860"365 366 async with MirageSystemTester(base_url) as tester:367 results = await tester.run_comprehensive_test()368 369 # Save results to file370 results_file = Path("test_results.json")371 with open(results_file, "w") as f:372 json.dump(results, f, indent=2, default=str)373 374 logger.info(f"📊 Detailed results saved to: {results_file}")375 376 # Exit with appropriate code377 sys.exit(0 if results["overall_success"] else 1)378 379if __name__ == "__main__":380 asyncio.run(main())