Skip to content

Commit

Permalink
Update agent_api_test.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@kyegomez
kyegomez authored Dec 17, 2024
1 parent 886ca11 commit 95a7d59
Showing 1 changed file with 84 additions and 58 deletions.
142 changes: 84 additions & 58 deletions api/agent_api_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,11 @@
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:
Expand All @@ -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
Expand All @@ -35,16 +36,17 @@ 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())}"}
json={"username": f"test_user_{int(time.time())}"},
)

if response.status_code == 200:
data = response.json()
session.user_id = data["user_id"]
Expand All @@ -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"]
Expand All @@ -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")

Expand All @@ -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)
Expand All @@ -106,81 +112,91 @@ 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(
f"{BASE_URL}/agent/{agent_id}",
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")

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(
Expand All @@ -190,105 +206,115 @@ 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
if session.user_id:
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)

0 comments on commit 95a7d59

Please sign in to comment.