Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -30,25 +30,38 @@ def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
     
     def try_connect():
         nonlocal storage
-        if storage:
-            try:
+        try:
+            if storage:
                 if storage.is_connected():
-                    return True
+                    # Verify session is active
+                    if storage.session_token is not None:
+                        return True
                 storage.close()
-            except:
-                pass
-        storage = HTTPGPUStorage()  # Create storage instance with default configuration
-        return storage._create_session()
+            
+            # Create new storage instance
+            storage = HTTPGPUStorage()
+            
+            # Initialize session
+            if storage._create_session():
+                # Verify session was created
+                if storage.session_token is not None and not storage._closing:
+                    return True
+            return False
+        except Exception as e:
+            logging.error(f"Connection error: {e}")
+            return False
     
     # 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")
-                # Test connection by making a simple request
-                status = storage.get_connection_status()
-                if status["connected"]:
-                    break
+                # Verify the connection is active
+                if storage.is_connected():
+                    # Test the connection with a basic operation
+                    test_key = "_connection_test"
+                    if storage.cache_data(test_key, {"test": True}):
+                        break
                 logging.warning("Connection established but not responsive")
             else:
                 logging.warning(f"HTTP connection attempt {attempt + 1} failed, retrying in {retry_delay}s...")
@@ -65,7 +78,6 @@ def http_storage_manager(max_retries=5, retry_delay=2, timeout=30.0):
             raise RuntimeError(error_msg)
     
     try:
-        # Yield the storage connection
         yield storage
     except Exception as e:
         logging.error(f"HTTP operation failed: {e}")
@@ -176,7 +188,7 @@ def test_ai_integration_http():
             model_data = {
                 "model_name": model_id,
                 "model_type": "florence-2-large",
-                "parameters": 771000000,  # Approximate parameter count
+                "parameters": 771000000, 
                 "architecture": "vision-language",
                 "loaded_at": time.time()
             }