Update test_ai_integration.py

test_ai_integration.py

@@ -111,8 +111,8 @@ def test_ai_integration():
         with websocket_manager() as storage:
             components['storage'] = storage  # Save storage reference
             
-            # Initialize virtual GPU stack with unlimited WebSocket storage
-            chip_for_loading = Chip(chip_id=0, vram_size_gb=None)  # Unlimited storage
+            # Initialize virtual GPU stack with unlimited WebSocket storage and shared connection
+            chip_for_loading = Chip(chip_id=0, vram_size_gb=None, storage=storage)  # Pass shared storage
             components['chips'].append(chip_for_loading)
             
             # Initialize VRAM with WebSocket storage
@@ -120,7 +120,7 @@ def test_ai_integration():
             vram.storage = storage  # Share WebSocket connection
             components['vram'] = vram
             
-            # Set up AI accelerator with proper initialization
+            # Set up AI accelerator - note it already has the shared storage
             ai_accelerator_for_loading = chip_for_loading.ai_accelerator
             ai_accelerator_for_loading.vram = vram  # Use WebSocket-backed VRAM
             ai_accelerator_for_loading.initialize_tensor_cores()  # Ensure tensor cores are ready