97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
import httpx
|
|
import time
|
|
|
|
BASE_URL = "http://127.0.0.1:8000/api/v1"
|
|
API_KEY = "storyline-secret-key-123"
|
|
|
|
client = httpx.Client(timeout=30.0)
|
|
|
|
def test_health():
|
|
print("Testing Health Check...")
|
|
response = client.get(f"{BASE_URL}/storyline/health")
|
|
print(f"Status: {response.status_code}, Body: {response.json()}")
|
|
|
|
def test_chat_unauthorized():
|
|
print("\nTesting Chat (No Auth)...")
|
|
response = client.post(f"{BASE_URL}/storyline/chat", json={"prompt": "Hello"})
|
|
print(f"Status: {response.status_code}, Body: {response.json()}")
|
|
|
|
def test_chat_authorized():
|
|
print("\nTesting Chat (Authorized)...")
|
|
headers = {"X-API-Key": API_KEY}
|
|
response = client.post(
|
|
f"{BASE_URL}/storyline/chat",
|
|
json={"prompt": "How do I use Storyline?", "context": "Module 1"},
|
|
headers=headers
|
|
)
|
|
print(f"Status: {response.status_code}, Body: {response.json()}")
|
|
|
|
def test_rate_limit():
|
|
print("\nTesting Rate Limit (making multiple calls)...")
|
|
headers = {"X-API-Key": API_KEY}
|
|
for i in range(22):
|
|
response = httpx.post(
|
|
f"{BASE_URL}/storyline/chat",
|
|
json={"prompt": f"Test {i}"},
|
|
headers=headers
|
|
)
|
|
if response.status_code == 429:
|
|
print(f"Rate limit hit at call {i+1}!")
|
|
return
|
|
print("Rate limit not hit (this is unexpected if limit is 20).")
|
|
|
|
def test_gemini():
|
|
print("\nTesting Gemini Endpoint...")
|
|
headers = {"X-API-Key": API_KEY}
|
|
response = client.post(f"{BASE_URL}/gemini/chat", json={"prompt": "Hello Gemini"}, headers=headers)
|
|
print(f"Status: {response.status_code}, Body: {response.json()}")
|
|
|
|
def test_openai():
|
|
print("\nTesting OpenAI Endpoint...")
|
|
headers = {"X-API-Key": API_KEY}
|
|
response = client.post(f"{BASE_URL}/openai/chat", json={"prompt": "Hello OpenAI"}, headers=headers)
|
|
print(f"Status: {response.status_code}, Body: {response.json()}")
|
|
|
|
def test_module_management():
|
|
print("\nTesting Module Management...")
|
|
headers = {"X-API-Key": API_KEY}
|
|
|
|
# 1. Create a module
|
|
print("Creating module 'Articulate 101'...")
|
|
response = client.post(
|
|
f"{BASE_URL}/admin/modules",
|
|
json={"name": "Articulate 101"},
|
|
headers=headers
|
|
)
|
|
if response.status_code == 400:
|
|
print("Module already exists, getting list...")
|
|
response = client.get(f"{BASE_URL}/admin/modules", headers=headers)
|
|
module_data = response.json()[0]
|
|
else:
|
|
module_data = response.json()
|
|
|
|
new_module_key = module_data["secret_key"]
|
|
print(f"Module created/found. Secret Key: {new_module_key}")
|
|
|
|
# 2. Use the new module key to call Gemini
|
|
print("\nTesting Gemini with NEW module key...")
|
|
module_headers = {"X-API-Key": new_module_key}
|
|
response = client.post(
|
|
f"{BASE_URL}/gemini/chat",
|
|
json={"prompt": "Can you see this with my module key?"},
|
|
headers=module_headers
|
|
)
|
|
print(f"Status: {response.status_code}, Body: {response.json()}")
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
test_health()
|
|
test_chat_unauthorized()
|
|
test_chat_authorized()
|
|
test_gemini()
|
|
test_openai()
|
|
test_module_management()
|
|
# test_rate_limit()
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|