""" Testing and Validation Suite for Mirage AI Avatar System Tests end-to-end functionality, latency, and performance """ import asyncio import time import aiohttp import json import numpy as np import cv2 import logging from pathlib import Path import subprocess import psutil from typing import Dict, Any, List logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class MirageSystemTester: """Comprehensive testing suite for the AI avatar system""" def __init__(self, base_url: str = "http://localhost:7860"): self.base_url = base_url self.session = None self.test_results = {} async def __aenter__(self): self.session = aiohttp.ClientSession() return self async def __aexit__(self, exc_type, exc_val, exc_tb): if self.session: await self.session.close() async def test_health_endpoint(self) -> bool: """Test basic health endpoint""" try: async with self.session.get(f"{self.base_url}/health") as response: data = await response.json() success = ( response.status == 200 and data.get("status") == "ok" and data.get("system") == "real-time-ai-avatar" ) self.test_results["health"] = { "success": success, "status": response.status, "data": data } logger.info(f"Health check: {'✅ PASS' if success else '❌ FAIL'}") return success except Exception as e: logger.error(f"Health check failed: {e}") self.test_results["health"] = {"success": False, "error": str(e)} return False async def test_pipeline_initialization(self) -> bool: """Test AI pipeline initialization""" try: start_time = time.time() async with self.session.post(f"{self.base_url}/initialize") as response: data = await response.json() init_time = time.time() - start_time success = ( response.status == 200 and data.get("status") in ["success", "already_initialized"] ) self.test_results["initialization"] = { "success": success, "status": response.status, "data": data, "init_time_seconds": init_time } logger.info(f"Pipeline init: {'✅ PASS' if success else '❌ FAIL'} ({init_time:.1f}s)") return success except Exception as e: logger.error(f"Pipeline initialization failed: {e}") self.test_results["initialization"] = {"success": False, "error": str(e)} return False async def test_reference_image_upload(self) -> bool: """Test reference image upload functionality""" try: # Create a test image test_image = np.zeros((512, 512, 3), dtype=np.uint8) cv2.circle(test_image, (256, 200), 50, (255, 255, 255), -1) # Face-like circle cv2.circle(test_image, (230, 180), 10, (0, 0, 0), -1) # Eye cv2.circle(test_image, (280, 180), 10, (0, 0, 0), -1) # Eye cv2.ellipse(test_image, (256, 220), (20, 10), 0, 0, 180, (0, 0, 0), 2) # Mouth # Encode as JPEG _, encoded = cv2.imencode('.jpg', test_image) image_data = encoded.tobytes() # Upload test image form_data = aiohttp.FormData() form_data.add_field('file', image_data, filename='test_face.jpg', content_type='image/jpeg') async with self.session.post(f"{self.base_url}/set_reference", data=form_data) as response: data = await response.json() success = ( response.status == 200 and data.get("status") == "success" ) self.test_results["reference_upload"] = { "success": success, "status": response.status, "data": data } logger.info(f"Reference upload: {'✅ PASS' if success else '❌ FAIL'}") return success except Exception as e: logger.error(f"Reference image upload failed: {e}") self.test_results["reference_upload"] = {"success": False, "error": str(e)} return False async def test_websocket_connections(self) -> bool: """Test WebSocket connections for audio and video""" try: import websockets # Test audio WebSocket audio_success = await self._test_websocket_endpoint("/audio") # Test video WebSocket video_success = await self._test_websocket_endpoint("/video") success = audio_success and video_success self.test_results["websockets"] = { "success": success, "audio_success": audio_success, "video_success": video_success } logger.info(f"WebSocket connections: {'✅ PASS' if success else '❌ FAIL'}") return success except Exception as e: logger.error(f"WebSocket test failed: {e}") self.test_results["websockets"] = {"success": False, "error": str(e)} return False async def _test_websocket_endpoint(self, endpoint: str) -> bool: """Test a specific WebSocket endpoint""" try: import websockets ws_url = self.base_url.replace("http://", "ws://") + endpoint async with websockets.connect(ws_url) as websocket: # Send test data if endpoint == "/audio": # Send 160ms of silence (16kHz, 16-bit) test_audio = np.zeros(int(16000 * 0.160), dtype=np.int16) await websocket.send(test_audio.tobytes()) else: # video # Send a small test JPEG test_frame = np.zeros((256, 256, 3), dtype=np.uint8) _, encoded = cv2.imencode('.jpg', test_frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) await websocket.send(encoded.tobytes()) # Wait for response response = await asyncio.wait_for(websocket.recv(), timeout=5.0) return len(response) > 0 except Exception as e: logger.error(f"WebSocket {endpoint} test failed: {e}") return False async def test_performance_metrics(self) -> bool: """Test performance metrics endpoint""" try: async with self.session.get(f"{self.base_url}/pipeline_status") as response: data = await response.json() success = response.status == 200 and data.get("initialized", False) self.test_results["performance_metrics"] = { "success": success, "status": response.status, "data": data } if success: stats = data.get("stats", {}) logger.info(f"Performance metrics: ✅ PASS") logger.info(f" GPU Memory: {stats.get('gpu_memory_used', 0):.1f} GB") logger.info(f" Video FPS: {stats.get('video_fps', 0):.1f}") logger.info(f" Avg Latency: {stats.get('avg_video_latency_ms', 0):.1f} ms") else: logger.info("Performance metrics: ❌ FAIL") return success except Exception as e: logger.error(f"Performance metrics test failed: {e}") self.test_results["performance_metrics"] = {"success": False, "error": str(e)} return False async def test_latency_benchmark(self) -> Dict[str, float]: """Benchmark system latency""" latencies = [] try: # Warm up for _ in range(5): start_time = time.time() async with self.session.get(f"{self.base_url}/health") as response: await response.json() latencies.append((time.time() - start_time) * 1000) # Actual benchmark latencies = [] for _ in range(20): start_time = time.time() async with self.session.get(f"{self.base_url}/pipeline_status") as response: await response.json() latencies.append((time.time() - start_time) * 1000) results = { "avg_latency_ms": np.mean(latencies), "min_latency_ms": np.min(latencies), "max_latency_ms": np.max(latencies), "p95_latency_ms": np.percentile(latencies, 95), "p99_latency_ms": np.percentile(latencies, 99) } self.test_results["latency_benchmark"] = results logger.info("Latency benchmark results:") logger.info(f" Average: {results['avg_latency_ms']:.1f} ms") logger.info(f" P95: {results['p95_latency_ms']:.1f} ms") logger.info(f" P99: {results['p99_latency_ms']:.1f} ms") return results except Exception as e: logger.error(f"Latency benchmark failed: {e}") return {} def test_system_requirements(self) -> Dict[str, Any]: """Test system requirements and capabilities""" results = {} try: # Check GPU availability try: import torch results["gpu_available"] = torch.cuda.is_available() if torch.cuda.is_available(): results["gpu_name"] = torch.cuda.get_device_name(0) results["gpu_memory_gb"] = torch.cuda.get_device_properties(0).total_memory / 1024**3 results["cuda_version"] = torch.version.cuda except ImportError: results["gpu_available"] = False # Check system resources memory = psutil.virtual_memory() results["system_memory_gb"] = memory.total / 1024**3 results["cpu_count"] = psutil.cpu_count() # Check disk space disk = psutil.disk_usage('/') results["disk_free_gb"] = disk.free / 1024**3 # Check required packages required_packages = [ "torch", "torchvision", "torchaudio", "opencv-python", "numpy", "fastapi", "websockets" ] missing_packages = [] for package in required_packages: try: __import__(package.replace("-", "_")) except ImportError: missing_packages.append(package) results["missing_packages"] = missing_packages results["requirements_met"] = len(missing_packages) == 0 self.test_results["system_requirements"] = results logger.info("System requirements:") logger.info(f" GPU: {'✅' if results['gpu_available'] else '❌'}") logger.info(f" Memory: {results['system_memory_gb']:.1f} GB") logger.info(f" CPU: {results['cpu_count']} cores") logger.info(f" Packages: {'✅' if results['requirements_met'] else '❌'}") return results except Exception as e: logger.error(f"System requirements check failed: {e}") return {"error": str(e)} async def run_comprehensive_test(self) -> Dict[str, Any]: """Run all tests and return comprehensive results""" logger.info("🧪 Starting comprehensive system test...") # System requirements (runs first, no server needed) self.test_system_requirements() # Server-dependent tests tests = [ ("Health Check", self.test_health_endpoint()), ("Pipeline Initialization", self.test_pipeline_initialization()), ("Reference Image Upload", self.test_reference_image_upload()), ("WebSocket Connections", self.test_websocket_connections()), ("Performance Metrics", self.test_performance_metrics()), ] # Run tests sequentially for test_name, test_coro in tests: logger.info(f"Running: {test_name}...") try: result = await test_coro if not result: logger.warning(f"{test_name} failed - may affect subsequent tests") except Exception as e: logger.error(f"{test_name} threw exception: {e}") # Latency benchmark (runs last) logger.info("Running latency benchmark...") await self.test_latency_benchmark() # Calculate overall success rate successful_tests = sum(1 for result in self.test_results.values() if isinstance(result, dict) and result.get("success", False)) total_tests = len([r for r in self.test_results.values() if isinstance(r, dict) and "success" in r]) overall_success = successful_tests / max(total_tests, 1) >= 0.8 # 80% success rate summary = { "overall_success": overall_success, "successful_tests": successful_tests, "total_tests": total_tests, "success_rate": successful_tests / max(total_tests, 1), "detailed_results": self.test_results } logger.info(f"🏁 Test completed: {successful_tests}/{total_tests} tests passed") logger.info(f"Overall result: {'✅ PASS' if overall_success else '❌ FAIL'}") return summary async def main(): """Main test runner""" import sys base_url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:7860" async with MirageSystemTester(base_url) as tester: results = await tester.run_comprehensive_test() # Save results to file results_file = Path("test_results.json") with open(results_file, "w") as f: json.dump(results, f, indent=2, default=str) logger.info(f"📊 Detailed results saved to: {results_file}") # Exit with appropriate code sys.exit(0 if results["overall_success"] else 1) if __name__ == "__main__": asyncio.run(main())