146 lines
4.0 KiB
Python
146 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script to verify environment variable passing to agent pods
|
|
"""
|
|
import requests
|
|
import json
|
|
import sys
|
|
|
|
API_URL = "http://localhost:8000"
|
|
|
|
def test_create_mysql_agent_with_env():
|
|
"""Test creating a MySQL agent with environment variables"""
|
|
|
|
payload = {
|
|
"name": "test-mysql-agent",
|
|
"template": "mysql_agent",
|
|
"env": {
|
|
"MYSQL_HOST": "test-mysql-server.mysql.database.azure.com",
|
|
"MYSQL_PORT": "3306",
|
|
"MYSQL_USER": "testuser",
|
|
"MYSQL_PASSWORD": "testpass",
|
|
"MYSQL_DATABASE": "testdb",
|
|
"OPENAI_API_KEY": "sk-test-key"
|
|
},
|
|
"config": {
|
|
"cpu_request": "100m",
|
|
"memory_request": "128Mi"
|
|
}
|
|
}
|
|
|
|
print("Creating MySQL agent with environment variables...")
|
|
print(f"Payload: {json.dumps(payload, indent=2)}")
|
|
|
|
try:
|
|
response = requests.post(f"{API_URL}/agents", json=payload)
|
|
print(f"\nStatus Code: {response.status_code}")
|
|
print(f"Response: {json.dumps(response.json(), indent=2)}")
|
|
|
|
if response.status_code == 200:
|
|
print("\n✅ Agent created successfully!")
|
|
return True
|
|
else:
|
|
print("\n❌ Agent creation failed!")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ Error: {str(e)}")
|
|
return False
|
|
|
|
def test_get_agent_status():
|
|
"""Test getting agent status"""
|
|
|
|
agent_name = "test-mysql-agent"
|
|
|
|
print(f"\nGetting status for agent: {agent_name}")
|
|
|
|
try:
|
|
response = requests.get(f"{API_URL}/agents/{agent_name}/status")
|
|
print(f"Status Code: {response.status_code}")
|
|
print(f"Response: {json.dumps(response.json(), indent=2)}")
|
|
|
|
if response.status_code == 200:
|
|
print("\n✅ Status retrieved successfully!")
|
|
return True
|
|
else:
|
|
print("\n❌ Failed to get status!")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ Error: {str(e)}")
|
|
return False
|
|
|
|
def test_get_template_info():
|
|
"""Test getting template information"""
|
|
|
|
template = "mysql_agent"
|
|
|
|
print(f"\nGetting template info for: {template}")
|
|
|
|
try:
|
|
response = requests.get(f"{API_URL}/templates/{template}")
|
|
print(f"Status Code: {response.status_code}")
|
|
print(f"Response: {json.dumps(response.json(), indent=2)}")
|
|
|
|
if response.status_code == 200:
|
|
print("\n✅ Template info retrieved successfully!")
|
|
return True
|
|
else:
|
|
print("\n❌ Failed to get template info!")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ Error: {str(e)}")
|
|
return False
|
|
|
|
def cleanup_test_agent():
|
|
"""Cleanup test agent"""
|
|
|
|
agent_name = "test-mysql-agent"
|
|
|
|
print(f"\nCleaning up test agent: {agent_name}")
|
|
|
|
try:
|
|
response = requests.delete(f"{API_URL}/agents/{agent_name}")
|
|
print(f"Status Code: {response.status_code}")
|
|
|
|
if response.status_code == 200:
|
|
print("✅ Test agent cleaned up!")
|
|
else:
|
|
print("⚠️ Cleanup may have failed (agent might not exist)")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Cleanup error: {str(e)}")
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 60)
|
|
print("Testing Agent Manager API - Environment Variables")
|
|
print("=" * 60)
|
|
|
|
# Test 1: Get template info
|
|
test_get_template_info()
|
|
|
|
print("\n" + "=" * 60)
|
|
|
|
# Test 2: Create agent with env vars
|
|
test_create_mysql_agent_with_env()
|
|
|
|
print("\n" + "=" * 60)
|
|
|
|
# Wait a bit for pod to be created
|
|
import time
|
|
print("\nWaiting 5 seconds for pod creation...")
|
|
time.sleep(5)
|
|
|
|
# Test 3: Get agent status
|
|
test_get_agent_status()
|
|
|
|
print("\n" + "=" * 60)
|
|
|
|
# Cleanup
|
|
cleanup_test_agent()
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Tests completed!")
|
|
print("=" * 60)
|