From 95a7d5962346d608ffc31c42356e64dd62af7a3e Mon Sep 17 00:00:00 2001 From: Kye Gomez <98760976+kyegomez@users.noreply.github.com> Date: Mon, 16 Dec 2024 18:27:59 -0800 Subject: [PATCH] Update agent_api_test.py --- api/agent_api_test.py | 142 +++++++++++++++++++++++++----------------- 1 file changed, 84 insertions(+), 58 deletions(-) diff --git a/api/agent_api_test.py b/api/agent_api_test.py index 2ad4e059..b8008697 100644 --- a/api/agent_api_test.py +++ b/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", - headers=session.headers, - json=payload + f"{BASE_URL}/agent", 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", - headers=session.headers + f"{BASE_URL}/users/me/agents", 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}") - except Exception as e: + logger.warning( + 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']}") - except Exception as e: - logger.exception(f"Exception revoking API key {key['name']}") - except Exception as e: + logger.warning( + f"Failed to revoke API key {key['name']}" + ) + 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)