Update test_ai_integration.py

test_ai_integration.py

@@ -151,6 +151,10 @@ def test_ai_integration():
                 trust_remote_code=True
             )
             
+            # Ensure WebSocket connection is active before proceeding
+            if not ai_accelerator_for_loading.storage.wait_for_connection():
+                raise RuntimeError("WebSocket connection lost - please retry")
+            
             # Calculate model size for proper VRAM allocation
             model_size = sum(p.numel() * p.element_size() for p in model.parameters())
             print(f"Model size: {model_size / (1024**3):.2f} GB")