Update agent_api_test.py

api/agent_api_test.py

@@ -3,11 +3,11 @@ from loguru import logger
 import time
 from typing import Dict, Optional, Tuple
 from uuid import UUID
-from datetime import datetime
 import sys
 
 BASE_URL = "http://localhost:8000/v1"
 
+
 def check_api_server() -> bool:
     """Check if the API server is running and accessible."""
     try:
@@ -22,9 +22,10 @@ def check_api_server() -> bool:
         logger.error(f"Error checking API server: {str(e)}")
         return False
 
+
 class TestSession:
     """Manages test session state and authentication."""
-    
+
     def __init__(self):
         self.user_id: Optional[UUID] = None
         self.api_key: Optional[str] = None
@@ -35,16 +36,17 @@ class TestSession:
         """Get headers with authentication."""
         return {"api-key": self.api_key} if self.api_key else {}
 
+
 def create_test_user(session: TestSession) -> Tuple[bool, str]:
     """Create a test user and store credentials in session."""
     logger.info("Creating test user")
-    
+
     try:
         response = requests.post(
             f"{BASE_URL}/users",
-            json={"username": f"test_user_{int(time.time())}"}
+            json={"username": f"test_user_{int(time.time())}"},
         )
-        
+
         if response.status_code == 200:
             data = response.json()
             session.user_id = data["user_id"]
@@ -58,17 +60,20 @@ def create_test_user(session: TestSession) -> Tuple[bool, str]:
         logger.exception("Exception during user creation")
         return False, str(e)
 
-def create_additional_api_key(session: TestSession) -> Tuple[bool, str]:
+
+def create_additional_api_key(
+    session: TestSession,
+) -> Tuple[bool, str]:
     """Test creating an additional API key."""
     logger.info("Creating additional API key")
-    
+
     try:
         response = requests.post(
             f"{BASE_URL}/users/{session.user_id}/api-keys",
             headers=session.headers,
-            json={"name": "Test Key"}
+            json={"name": "Test Key"},
         )
-        
+
         if response.status_code == 200:
             logger.success("Created additional API key")
             return True, response.json()["key"]
@@ -79,7 +84,10 @@ def create_additional_api_key(session: TestSession) -> Tuple[bool, str]:
         logger.exception("Exception during API key creation")
         return False, str(e)
 
-def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
+
+def test_create_agent(
+    session: TestSession,
+) -> Tuple[bool, Optional[UUID]]:
     """Test creating a new agent."""
     logger.info("Testing agent creation")
 
@@ -88,16 +96,14 @@ def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
         "system_prompt": "You are a helpful assistant",
         "model_name": "gpt-4",
         "description": "Test agent",
-        "tags": ["test", "automated"]
+        "tags": ["test", "automated"],
     }
 
     try:
         response = requests.post(
-            f"{BASE_URL}/agent",
+            f"{BASE_URL}/agent", headers=session.headers, json=payload
-            headers=session.headers,
-            json=payload
         )
-        
+
         if response.status_code == 200:
             agent_id = response.json()["agent_id"]
             session.test_agents.append(agent_id)
@@ -106,35 +112,40 @@ def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
         else:
             logger.error(f"Failed to create agent: {response.text}")
             return False, None
-    except Exception as e:
+    except Exception:
         logger.exception("Exception during agent creation")
         return False, None
 
+
 def test_list_user_agents(session: TestSession) -> bool:
     """Test listing user's agents."""
     logger.info("Testing user agent listing")
 
     try:
         response = requests.get(
-            f"{BASE_URL}/users/me/agents",
+            f"{BASE_URL}/users/me/agents", headers=session.headers
-            headers=session.headers
         )
-        
+
         if response.status_code == 200:
             agents = response.json()
             logger.success(f"Found {len(agents)} user agents")
             return True
         else:
-            logger.error(f"Failed to list user agents: {response.text}")
+            logger.error(
+                f"Failed to list user agents: {response.text}"
+            )
             return False
-    except Exception as e:
+    except Exception:
         logger.exception("Exception during agent listing")
         return False
 
-def test_agent_operations(session: TestSession, agent_id: UUID) -> bool:
+
+def test_agent_operations(
+    session: TestSession, agent_id: UUID
+) -> bool:
     """Test various operations on an agent."""
     logger.info(f"Testing operations for agent {agent_id}")
-    
+
     # Test update
     try:
         update_response = requests.patch(
@@ -142,28 +153,33 @@ def test_agent_operations(session: TestSession, agent_id: UUID) -> bool:
             headers=session.headers,
             json={
                 "description": "Updated description",
-                "tags": ["test", "updated"]
+                "tags": ["test", "updated"],
-            }
+            },
         )
         if update_response.status_code != 200:
-            logger.error(f"Failed to update agent: {update_response.text}")
+            logger.error(
+                f"Failed to update agent: {update_response.text}"
+            )
             return False
-        
+
         # Test metrics
         metrics_response = requests.get(
             f"{BASE_URL}/agent/{agent_id}/metrics",
-            headers=session.headers
+            headers=session.headers,
         )
         if metrics_response.status_code != 200:
-            logger.error(f"Failed to get agent metrics: {metrics_response.text}")
+            logger.error(
+                f"Failed to get agent metrics: {metrics_response.text}"
+            )
             return False
-            
+
         logger.success("Successfully performed agent operations")
         return True
-    except Exception as e:
+    except Exception:
         logger.exception("Exception during agent operations")
         return False
 
+
 def test_completion(session: TestSession, agent_id: UUID) -> bool:
     """Test running a completion."""
     logger.info("Testing completion")
@@ -171,16 +187,16 @@ def test_completion(session: TestSession, agent_id: UUID) -> bool:
     payload = {
         "prompt": "What is the weather like today?",
         "agent_id": agent_id,
-        "max_tokens": 100
+        "max_tokens": 100,
     }
 
     try:
         response = requests.post(
             f"{BASE_URL}/agent/completions",
             headers=session.headers,
-            json=payload
+            json=payload,
         )
-        
+
         if response.status_code == 200:
             completion_data = response.json()
             logger.success(
@@ -190,26 +206,29 @@ def test_completion(session: TestSession, agent_id: UUID) -> bool:
         else:
             logger.error(f"Failed to get completion: {response.text}")
             return False
-    except Exception as e:
+    except Exception:
         logger.exception("Exception during completion")
         return False
 
+
 def cleanup_test_resources(session: TestSession):
     """Clean up all test resources."""
     logger.info("Cleaning up test resources")
-    
+
     # Delete test agents
     for agent_id in session.test_agents:
         try:
             response = requests.delete(
                 f"{BASE_URL}/agent/{agent_id}",
-                headers=session.headers
+                headers=session.headers,
             )
             if response.status_code == 200:
                 logger.debug(f"Deleted agent {agent_id}")
             else:
-                logger.warning(f"Failed to delete agent {agent_id}: {response.text}")
+                logger.warning(
-        except Exception as e:
+                    f"Failed to delete agent {agent_id}: {response.text}"
+                )
+        except Exception:
             logger.exception(f"Exception deleting agent {agent_id}")
 
     # Revoke API keys
@@ -217,78 +236,85 @@ def cleanup_test_resources(session: TestSession):
         try:
             response = requests.get(
                 f"{BASE_URL}/users/{session.user_id}/api-keys",
-                headers=session.headers
+                headers=session.headers,
             )
             if response.status_code == 200:
                 for key in response.json():
                     try:
                         revoke_response = requests.delete(
                             f"{BASE_URL}/users/{session.user_id}/api-keys/{key['key']}",
-                            headers=session.headers
+                            headers=session.headers,
                         )
                         if revoke_response.status_code == 200:
-                            logger.debug(f"Revoked API key {key['name']}")
+                            logger.debug(
+                                f"Revoked API key {key['name']}"
+                            )
                         else:
-                            logger.warning(f"Failed to revoke API key {key['name']}")
+                            logger.warning(
-                    except Exception as e:
+                                f"Failed to revoke API key {key['name']}"
-                        logger.exception(f"Exception revoking API key {key['name']}")
+                            )
-        except Exception as e:
+                    except Exception:
+                        logger.exception(
+                            f"Exception revoking API key {key['name']}"
+                        )
+        except Exception:
             logger.exception("Exception getting API keys for cleanup")
 
+
 def run_test_workflow():
     """Run complete test workflow."""
     logger.info("Starting API tests")
-    
+
     # Check if API server is running first
     if not check_api_server():
         return False
-        
+
     session = TestSession()
-    success = True
+
-    
     try:
         # Create user
         user_success, message = create_test_user(session)
         if not user_success:
             logger.error(f"User creation failed: {message}")
             return False
-        
+
         # Create additional API key
         key_success, key = create_additional_api_key(session)
         if not key_success:
             logger.error(f"API key creation failed: {key}")
             return False
-        
+
         # Create agent
         agent_success, agent_id = test_create_agent(session)
         if not agent_success or not agent_id:
             logger.error("Agent creation failed")
             return False
-        
+
         # Test user agent listing
         if not test_list_user_agents(session):
             logger.error("Agent listing failed")
             return False
-        
+
         # Test agent operations
         if not test_agent_operations(session, agent_id):
             logger.error("Agent operations failed")
             return False
-        
+
         # Test completion
         if not test_completion(session, agent_id):
             logger.error("Completion test failed")
             return False
-        
+
         logger.success("All tests completed successfully")
         return True
-        
+
-    except Exception as e:
+    except Exception:
         logger.exception("Exception during test workflow")
         return False
     finally:
         cleanup_test_resources(session)
 
+
 if __name__ == "__main__":
     success = run_test_workflow()
     sys.exit(0 if success else 1)