Upload test_ai_integration.py

test_ai_integration.py

@@ -68,8 +68,7 @@ def websocket_manager(max_retries=5, retry_delay=2, timeout=30.0):
             raise RuntimeError(error_msg)
     
     try:
-        # Set up keep-alive mechanism
-        storage.set_keep_alive(True)
+        # Yield the storage connection
         yield storage
     except Exception as e:
         logging.error(f"WebSocket operation failed: {e}")
@@ -82,7 +81,6 @@ def websocket_manager(max_retries=5, retry_delay=2, timeout=30.0):
     finally:
         if storage:
             try:
-                storage.set_keep_alive(False)
                 storage.close()
             except:
                 pass
@@ -243,7 +241,6 @@ def test_ai_integration():
                 if (components['storage'] and 
                     components['storage'].wait_for_connection(timeout=10.0)):
                     shared_storage = components['storage']
-                    shared_storage.set_keep_alive(True)  # Enable keep-alive
                     logging.info("Successfully reused existing WebSocket connection")
                     break
                 else:
@@ -252,7 +249,6 @@ def test_ai_integration():
                         if new_storage and new_storage.wait_for_connection(timeout=10.0):
                             components['storage'] = new_storage
                             shared_storage = new_storage
-                            shared_storage.set_keep_alive(True)  # Enable keep-alive
                             logging.info("Successfully established new WebSocket connection")
                             break
             except Exception as e: