diff --git a/api/agent_api_test.py b/api/agent_api_test.py index 066efc4f..2ad4e059 100644 --- a/api/agent_api_test.py +++ b/api/agent_api_test.py @@ -1,107 +1,294 @@ import requests from loguru import logger import time - -# Configure loguru -logger.add( - "api_tests_{time}.log", - rotation="100 MB", - level="DEBUG", - format="{time} {level} {message}", -) +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: + 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())}"} + ) + + 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) - logger.debug(f"Create response: {response.json()}") + try: + response = requests.post( + 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) + logger.success(f"Created agent with ID: {agent_id}") + return True, agent_id + else: + logger.error(f"Failed to create agent: {response.text}") + return False, None + except Exception as e: + logger.exception("Exception during agent creation") + return False, None - if response.status_code == 200: - logger.success("Successfully created agent") - return response.json()["agent_id"] - else: - logger.error(f"Failed to create agent: {response.text}") - return 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 + ) + + 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}") + return False + except Exception as e: + logger.exception("Exception during agent listing") + return False -def test_list_agents(): - """Test listing all agents.""" - logger.info("Testing agent listing") +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 - response = requests.get(f"{BASE_URL}/agents") - logger.debug(f"List response: {response.json()}") - - if response.status_code == 200: - logger.success(f"Found {len(response.json())} agents") - else: - logger.error(f"Failed to list agents: {response.text}") - - -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 } - response = requests.post( - f"{BASE_URL}/agent/completions", json=payload - ) - logger.debug(f"Completion response: {response.json()}") - - if response.status_code == 200: - logger.success("Successfully got completion") - else: - logger.error(f"Failed to get completion: {response.text}") + try: + response = requests.post( + f"{BASE_URL}/agent/completions", + headers=session.headers, + json=payload + ) + + if response.status_code == 200: + 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") + + # Delete test agents + for agent_id in session.test_agents: + try: + response = requests.delete( + f"{BASE_URL}/agent/{agent_id}", + 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.exception(f"Exception deleting agent {agent_id}") -def test_delete_agent(agent_id): - """Test deleting an agent.""" - logger.info("Testing agent deletion") + # 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") - response = requests.delete(f"{BASE_URL}/agent/{agent_id}") - logger.debug(f"Delete response: {response.json()}") - - if response.status_code == 200: - logger.success("Successfully deleted agent") - else: - logger.error(f"Failed to delete agent: {response.text}") - - -def run_tests(): - """Run all tests in sequence.""" +def run_test_workflow(): + """Run complete test workflow.""" logger.info("Starting API tests") - - # Create agent and get ID - agent_id = test_create_agent() - if not agent_id: - logger.error("Cannot continue tests without agent ID") - return - - # Wait a bit for agent to be ready - time.sleep(1) - - # Run other tests - test_list_agents() - test_completion(agent_id) - test_delete_agent(agent_id) - - logger.info("Tests completed") - + + # 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: + 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)