Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -22,7 +22,7 @@ logging.basicConfig(
     format='%(asctime)s - %(levelname)s - %(message)s'
 )
 
-# HTTP connection manager with persistent connection
+# HTTP connection manager with retry handling
 @contextlib.contextmanager
 def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
     storage = None
@@ -32,35 +32,24 @@ def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
         nonlocal storage
         if storage:
             try:
-                if hasattr(storage, 'is_connected') and storage.is_connected():
+                if storage.is_connected():
                     return True
                 storage.close()
             except:
                 pass
-        storage = HTTPGPUStorage()  # Create storage instance
-        try:
-            # Set connection parameters after creation
-            storage.set_timeout(timeout)
-            storage.set_max_retries(max_retries)
-            storage.set_retry_delay(retry_delay)
-            
-            # Enable persistent connection if available
-            if hasattr(storage, 'enable_keep_alive'):
-                storage.enable_keep_alive()
-            return storage.connect()
-        except Exception as e:
-            logging.error(f"Connection configuration error: {e}")
-            return False
+        storage = HTTPGPUStorage()  # Create storage instance with default configuration
+        return storage._create_session()
     
     # Initial connection with improved error handling
     for attempt in range(max_retries):
         try:
             if try_connect():
                 logging.info("Successfully connected to GPU storage server via HTTP")
-                # Simple connection test by checking server status
-                if storage.check_server_status():
+                # Test connection by making a simple request
+                status = storage.get_connection_status()
+                if status["connected"]:
                     break
-                logging.warning("Connection established but server not responding")
+                logging.warning("Connection established but not responsive")
             else:
                 logging.warning(f"HTTP connection attempt {attempt + 1} failed, retrying in {retry_delay}s...")
                 time.sleep(retry_delay * (1.5 ** attempt))  # Exponential backoff
@@ -285,14 +274,14 @@ def test_ai_integration_http():
         for attempt in range(max_connection_attempts):
             try:
                 if (components['storage'] and 
-                    components['storage'].wait_for_connection(timeout=10.0)):
+                    components['storage'].is_connected()):
                     shared_storage = components['storage']
                     logging.info("Successfully reused existing HTTP connection")
                     break
                 else:
                     logging.warning("Existing connection unavailable, creating new HTTP connection...")
-                    with http_storage_manager(timeout=30.0) as new_storage:
-                        if new_storage and new_storage.wait_for_connection(timeout=10.0):
+                    with http_storage_manager() as new_storage:
+                        if new_storage and new_storage.is_connected():
                             components['storage'] = new_storage
                             shared_storage = new_storage
                             logging.info("Successfully established new HTTP connection")
@@ -335,9 +324,9 @@ def test_ai_integration_http():
             max_retry = 3
             for retry in range(max_retry):
                 try:
-                    if not shared_storage.wait_for_connection(timeout=5.0):
+                    if not shared_storage.is_connected():
                         logging.warning(f"Connection check failed for chip {i}, attempt {retry + 1}")
-                        shared_storage.reconnect()  # Attempt to reconnect
+                        shared_storage._create_session()  # Attempt to reconnect
                         time.sleep(1)
                         continue