Kye Gomez
4 weeks ago
committed by
GitHub
parent
a564fd27e4
commit
5ed5af20a7
@ -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 test_create_agent():
|
||||
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(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:
|
||||
logger.success("Successfully created agent")
|
||||
return response.json()["agent_id"]
|
||||
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 None
|
||||
|
||||
return False, None
|
||||
except Exception as e:
|
||||
logger.exception("Exception during agent creation")
|
||||
return False, None
|
||||
|
||||
def test_list_agents():
|
||||
"""Test listing all agents."""
|
||||
logger.info("Testing agent listing")
|
||||
def test_list_user_agents(session: TestSession) -> bool:
|
||||
"""Test listing user's agents."""
|
||||
logger.info("Testing user agent listing")
|
||||
|
||||
response = requests.get(f"{BASE_URL}/agents")
|
||||
logger.debug(f"List response: {response.json()}")
|
||||
try:
|
||||
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):
|
||||
"""Test deleting an agent."""
|
||||
logger.info("Testing agent deletion")
|
||||
|
||||
response = requests.delete(f"{BASE_URL}/agent/{agent_id}")
|
||||
logger.debug(f"Delete response: {response.json()}")
|
||||
|
||||
# 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.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():
|
||||
"""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
|
||||
# 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
|
||||
|
||||
# Wait a bit for agent to be ready
|
||||
time.sleep(1)
|
||||
# Test agent operations
|
||||
if not test_agent_operations(session, agent_id):
|
||||
logger.error("Agent operations failed")
|
||||
return False
|
||||
|
||||
# Run other tests
|
||||
test_list_agents()
|
||||
test_completion(agent_id)
|
||||
test_delete_agent(agent_id)
|
||||
# Test completion
|
||||
if not test_completion(session, agent_id):
|
||||
logger.error("Completion test failed")
|
||||
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)
|
||||
|
||||
Loading…
Reference in new issue