Update agent_api_test.py

api/agent_api_test.py

@@ -1,107 +1,294 @@
 import requests
 from loguru import logger
 import time
+from typing import Dict, Optional, Tuple
+from uuid import UUID
+from datetime import datetime
+import sys
 
-# Configure loguru
+BASE_URL = "http://localhost:8000/v1"
-logger.add(
+
-    "api_tests_{time}.log",
+def check_api_server() -> bool:
-    rotation="100 MB",
+    """Check if the API server is running and accessible."""
-    level="DEBUG",
+    try:
-    format="{time} {level} {message}",
+        response = requests.get(f"{BASE_URL}/docs")
+        return response.status_code == 200
+    except requests.exceptions.ConnectionError:
+        logger.error("API server is not running at {BASE_URL}")
+        logger.error("Please start the API server first with:")
+        logger.error("    python main.py")
+        return False
+    except Exception as e:
+        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
+        self.test_agents: list[UUID] = []
+
+    @property
+    def headers(self) -> Dict[str, str]:
+        """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())}"}
         )
         
-BASE_URL = "http://localhost:8000/v1"
+        if response.status_code == 200:
+            data = response.json()
+            session.user_id = data["user_id"]
+            session.api_key = data["api_key"]
+            logger.success(f"Created user with ID: {session.user_id}")
+            return True, "Success"
+        else:
+            logger.error(f"Failed to create user: {response.text}")
+            return False, response.text
+    except Exception as e:
+        logger.exception("Exception during user creation")
+        return False, str(e)
+
+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"}
+        )
         
+        if response.status_code == 200:
+            logger.success("Created additional API key")
+            return True, response.json()["key"]
+        else:
+            logger.error(f"Failed to create API key: {response.text}")
+            return False, response.text
+    except Exception as e:
+        logger.exception("Exception during API key creation")
+        return False, str(e)
 
-def test_create_agent():
+def test_create_agent(session: TestSession) -> Tuple[bool, Optional[UUID]]:
     """Test creating a new agent."""
     logger.info("Testing agent creation")
 
     payload = {
-        "agent_name": "Test Agent",
+        "agent_name": f"Test Agent {int(time.time())}",
         "system_prompt": "You are a helpful assistant",
         "model_name": "gpt-4",
         "description": "Test agent",
-        "tags": ["test"],
+        "tags": ["test", "automated"]
     }
 
-    response = requests.post(f"{BASE_URL}/agent", json=payload)
+    try:
-    logger.debug(f"Create response: {response.json()}")
+        response = requests.post(
+            f"{BASE_URL}/agent",
+            headers=session.headers,
+            json=payload
+        )
         
         if response.status_code == 200:
-        logger.success("Successfully created agent")
+            agent_id = response.json()["agent_id"]
-        return response.json()["agent_id"]
+            session.test_agents.append(agent_id)
+            logger.success(f"Created agent with ID: {agent_id}")
+            return True, agent_id
         else:
             logger.error(f"Failed to create agent: {response.text}")
-        return None
+            return False, None
-
+    except Exception as e:
+        logger.exception("Exception during agent creation")
+        return False, None
 
-def test_list_agents():
+def test_list_user_agents(session: TestSession) -> bool:
-    """Test listing all agents."""
+    """Test listing user's agents."""
-    logger.info("Testing agent listing")
+    logger.info("Testing user agent listing")
 
-    response = requests.get(f"{BASE_URL}/agents")
+    try:
-    logger.debug(f"List response: {response.json()}")
+        response = requests.get(
+            f"{BASE_URL}/users/me/agents",
+            headers=session.headers
+        )
         
         if response.status_code == 200:
-        logger.success(f"Found {len(response.json())} agents")
+            agents = response.json()
+            logger.success(f"Found {len(agents)} user agents")
+            return True
         else:
-        logger.error(f"Failed to list agents: {response.text}")
+            logger.error(f"Failed to list user agents: {response.text}")
+            return False
+    except Exception as e:
+        logger.exception("Exception during agent listing")
+        return False
+
+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(
+            f"{BASE_URL}/agent/{agent_id}",
+            headers=session.headers,
+            json={
+                "description": "Updated description",
+                "tags": ["test", "updated"]
+            }
+        )
+        if update_response.status_code != 200:
+            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
+        )
+        if metrics_response.status_code != 200:
+            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:
+        logger.exception("Exception during agent operations")
+        return False
 
-def test_completion(agent_id):
+def test_completion(session: TestSession, agent_id: UUID) -> bool:
     """Test running a completion."""
     logger.info("Testing completion")
 
     payload = {
         "prompt": "What is the weather like today?",
         "agent_id": agent_id,
+        "max_tokens": 100
     }
 
+    try:
         response = requests.post(
-        f"{BASE_URL}/agent/completions", json=payload
+            f"{BASE_URL}/agent/completions",
+            headers=session.headers,
+            json=payload
         )
-    logger.debug(f"Completion response: {response.json()}")
         
         if response.status_code == 200:
-        logger.success("Successfully got completion")
+            completion_data = response.json()
+            logger.success(
+                f"Got completion, used {completion_data['token_usage']['total_tokens']} tokens"
+            )
+            return True
         else:
             logger.error(f"Failed to get completion: {response.text}")
+            return False
+    except Exception as e:
+        logger.exception("Exception during completion")
+        return False
 
+def cleanup_test_resources(session: TestSession):
+    """Clean up all test resources."""
+    logger.info("Cleaning up test resources")
     
-def test_delete_agent(agent_id):
+    # Delete test agents
-    """Test deleting an agent."""
+    for agent_id in session.test_agents:
-    logger.info("Testing agent deletion")
+        try:
-
+            response = requests.delete(
-    response = requests.delete(f"{BASE_URL}/agent/{agent_id}")
+                f"{BASE_URL}/agent/{agent_id}",
-    logger.debug(f"Delete response: {response.json()}")
+                headers=session.headers
-
+            )
             if response.status_code == 200:
-        logger.success("Successfully deleted agent")
+                logger.debug(f"Deleted agent {agent_id}")
             else:
-        logger.error(f"Failed to delete agent: {response.text}")
+                logger.warning(f"Failed to delete agent {agent_id}: {response.text}")
+        except Exception as e:
+            logger.exception(f"Exception deleting agent {agent_id}")
 
+    # Revoke API keys
+    if session.user_id:
+        try:
+            response = requests.get(
+                f"{BASE_URL}/users/{session.user_id}/api-keys",
+                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
+                        )
+                        if revoke_response.status_code == 200:
+                            logger.debug(f"Revoked API key {key['name']}")
+                        else:
+                            logger.warning(f"Failed to revoke API key {key['name']}")
+                    except Exception as e:
+                        logger.exception(f"Exception revoking API key {key['name']}")
+        except Exception as e:
+            logger.exception("Exception getting API keys for cleanup")
 
-def run_tests():
+def run_test_workflow():
-    """Run all tests in sequence."""
+    """Run complete test workflow."""
     logger.info("Starting API tests")
     
-    # Create agent and get ID
+    # Check if API server is running first
-    agent_id = test_create_agent()
+    if not check_api_server():
-    if not agent_id:
+        return False
-        logger.error("Cannot continue tests without agent ID")
+        
-        return
+    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
         
-    # Wait a bit for agent to be ready
+        # Test agent operations
-    time.sleep(1)
+        if not test_agent_operations(session, agent_id):
+            logger.error("Agent operations failed")
+            return False
         
-    # Run other tests
+        # Test completion
-    test_list_agents()
+        if not test_completion(session, agent_id):
-    test_completion(agent_id)
+            logger.error("Completion test failed")
-    test_delete_agent(agent_id)
+            return False
         
-    logger.info("Tests completed")
+        logger.success("All tests completed successfully")
+        return True
         
+    except Exception as e:
+        logger.exception("Exception during test workflow")
+        return False
+    finally:
+        cleanup_test_resources(session)
 
 if __name__ == "__main__":
-    run_tests()
+    success = run_test_workflow()
+    sys.exit(0 if success else 1)