Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -32,25 +32,21 @@ def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
         nonlocal storage
         if storage:
             try:
-                if storage.is_connected():
+                if hasattr(storage, 'is_connected') and storage.is_connected():
                     return True
                 storage.close()
             except:
                 pass
-        storage = HTTPGPUStorage(keep_alive=True)  # Enable keep-alive
+        storage = HTTPGPUStorage()  # Create storage instance
         try:
-            storage.configure({
-                'timeout': timeout,
-                'retry_strategy': {
-                    'max_retries': max_retries,
-                    'retry_delay': retry_delay,
-                    'backoff_factor': 1.5
-                },
-                'connection_pool': {
-                    'max_size': 10,
-                    'max_retries': 3
-                }
-            })
+            # 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}")
@@ -61,8 +57,10 @@ def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
         try:
             if try_connect():
                 logging.info("Successfully connected to GPU storage server via HTTP")
-                storage.ping()  # Verify connection is responsive
-                break
+                # Simple connection test by checking server status
+                if storage.check_server_status():
+                    break
+                logging.warning("Connection established but server not responding")
             else:
                 logging.warning(f"HTTP connection attempt {attempt + 1} failed, retrying in {retry_delay}s...")
                 time.sleep(retry_delay * (1.5 ** attempt))  # Exponential backoff
@@ -97,28 +95,26 @@ def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
         
 # Enhanced cleanup handler with connection management
 def cleanup_resources():
-    # Get all active HTTP connections
-    active_connections = HTTPGPUStorage.get_active_connections()
-    
-    # Properly close each connection
-    for conn in active_connections:
-        try:
-            if conn and conn.is_connected():
-                conn.flush()  # Ensure all pending operations are completed
-                conn.close()
-        except Exception as e:
-            logging.error(f"Error closing HTTP connection: {e}")
+    try:
+        # Get the current storage instance if it exists
+        from http_storage import HTTPGPUStorage
+        current_storage = HTTPGPUStorage.get_current_instance()
+        if current_storage is not None:
+            try:
+                # Ensure all pending operations are completed
+                if hasattr(current_storage, 'sync'):
+                    current_storage.sync()
+                # Close the connection
+                current_storage.close()
+            except Exception as e:
+                logging.error(f"Error closing HTTP storage: {e}")
+    except Exception as e:
+        logging.error(f"Error in storage cleanup: {e}")
     
     # Clear VRAM and other resources
     import gc
     gc.collect()
     
-    try:
-        # Force close any remaining connections
-        HTTPGPUStorage.close_all_connections()
-    except Exception as e:
-        logging.error(f"Error in final connection cleanup: {e}")
-    
 # Register enhanced cleanup handler
 atexit.register(cleanup_resources)