ai-mirow-projects
/
swarms
mirror of https://github.com/kyegomez/swarms
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update agent_api_test.py

Browse Source
668^2
Kye Gomez 3 weeks ago committed by GitHub
parent 886ca11370
commit 95a7d59623
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 84 additions and 58 deletions
Show Stats Download Patch File Download Diff File

98
api/agent_api_test.py
Unescape View File

@ -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,6 +22,7 @@ def check_api_server() -> bool:
logger.error(f"Error checking API server: {str(e)}")
return False
class TestSession:
"""Manages test session state and authentication."""
@ -35,6 +36,7 @@ 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")
@ -42,7 +44,7 @@ def create_test_user(session: TestSession) -> Tuple[bool, str]:
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:
@ -58,7 +60,10 @@ 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")
@ -66,7 +71,7 @@ def create_additional_api_key(session: TestSession) -> Tuple[bool, str]:
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:
@ -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,14 +96,12 @@ 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:
@ -106,18 +112,18 @@ 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:
@ -125,13 +131,18 @@ def test_list_user_agents(session: TestSession) -> bool:
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}")
@ -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,14 +187,14 @@ 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:
@ -190,10 +206,11 @@ 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")
@ -203,13 +220,15 @@ def cleanup_test_resources(session: TestSession):
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,24 +236,31 @@ 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")
@ -244,7 +270,6 @@ def run_test_workflow():
return False
session = TestSession()
success = True
try:
# Create user
@ -283,12 +308,13 @@ def run_test_workflow():
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)

Write Preview
Loading…
Cancel
Save