Spaces:
Running
Running
File size: 6,454 Bytes
501c69f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 |
import requests
import json
import time
from typing import Dict, Any
class APITester:
def __init__(self, base_url: str = "http://localhost:8000"):
self.base_url = base_url
self.session = requests.Session()
def test_health_check(self) -> None:
"""Test the health check endpoint."""
print("\n=== Testing Health Check Endpoint ===")
try:
response = self.session.get(f"{self.base_url}/health_check")
print(f"Status Code: {response.status_code}")
print(f"Response: {response.json()}")
assert response.status_code == 200
print("✅ Health check test passed!")
except Exception as e:
print(f"❌ Health check test failed: {str(e)}")
def test_models(self) -> None:
"""Test the models endpoint."""
print("\n=== Testing Models Endpoint ===")
try:
response = self.session.get(f"{self.base_url}/models")
print(f"Status Code: {response.status_code}")
data = response.json()
print(f"Number of models available: {len(data['data'])}")
print("Sample models:")
for model in data['data'][:5]: # Show first 5 models
print(f"- {model['id']}")
assert response.status_code == 200
print("✅ Models endpoint test passed!")
except Exception as e:
print(f"❌ Models endpoint test failed: {str(e)}")
def test_chat_completions_non_streaming(self) -> None:
"""Test the chat completions endpoint without streaming."""
print("\n=== Testing Chat Completions Endpoint (Non-Streaming) ===")
payload = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me a short joke about programming."}
],
"temperature": 0.7,
"max_tokens": 150,
"stream": False
}
try:
response = self.session.post(
f"{self.base_url}/chat/completions",
json=payload
)
print(f"Status Code: {response.status_code}")
if response.status_code == 200:
data = response.json()
print("Response content:")
print(data['choices'][0]['message']['content'])
assert response.status_code == 200
print("✅ Chat completions (non-streaming) test passed!")
except Exception as e:
print(f"❌ Chat completions (non-streaming) test failed: {str(e)}")
def test_chat_completions_streaming(self) -> None:
"""Test the chat completions endpoint with streaming."""
print("\n=== Testing Chat Completions Endpoint (Streaming) ===")
payload = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Write 5 lines about India"}
],
"temperature": 0.7,
"max_tokens": 150,
"stream": True
}
try:
with self.session.post(
f"{self.base_url}/chat/completions",
json=payload,
stream=True,
headers={"Accept": "text/event-stream"}
) as response:
print(f"Status Code: {response.status_code}")
print("Streaming response:")
buffer = ""
for chunk in response.iter_lines():
if chunk:
chunk = chunk.decode('utf-8')
if chunk.startswith('data: '):
chunk = chunk[6:] # Remove 'data: ' prefix
if chunk.strip() == '[DONE]':
break
try:
data = json.loads(chunk)
if 'choices' in data and len(data['choices']) > 0:
if 'delta' in data['choices'][0] and 'content' in data['choices'][0]['delta']:
content = data['choices'][0]['delta']['content']
print(content, end='', flush=True)
time.sleep(0.1) # Add a small delay to simulate real-time streaming
except json.JSONDecodeError:
continue
print("\n✅ Chat completions (streaming) test passed!")
except Exception as e:
print(f"❌ Chat completions (streaming) test failed: {str(e)}")
def test_developer_info(self) -> None:
"""Test the developer info endpoint."""
print("\n=== Testing Developer Info Endpoint ===")
try:
response = self.session.get(f"{self.base_url}/developer_info")
print(f"Status Code: {response.status_code}")
print("Developer Info:")
print(json.dumps(response.json(), indent=2))
assert response.status_code == 200
print("✅ Developer info test passed!")
except Exception as e:
print(f"❌ Developer info test failed: {str(e)}")
def run_all_tests(self) -> None:
"""Run all tests sequentially."""
tests = [
self.test_health_check,
self.test_models,
self.test_chat_completions_non_streaming,
self.test_chat_completions_streaming,
self.test_developer_info
]
print("🚀 Starting API Tests...")
start_time = time.time()
for test in tests:
test()
end_time = time.time()
duration = end_time - start_time
print(f"\n============================")
print(f"🏁 All tests completed in {duration:.2f} seconds")
print(f"============================")
def main():
# Initialize tester with your API's base URL
tester = APITester("http://localhost:8000")
# Run all tests
tester.run_all_tests()
if __name__ == "__main__":
main() |