forked from phoenix/litellm-mirror
test(tests/): add unit testing for proxy server endpoints
This commit is contained in:
Krrish Dholakia
parent
b2b41727ce
commit
f5ced089d6
11 changed files with 870 additions and 111 deletions
|
|
@ -1,58 +0,0 @@
|
|||
# What this tests ?
|
||||
## Tests /chat/completions by generating a key and then making a chat completions request
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
|
||||
async def generate_key(session):
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["gpt-4"],
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def chat_completion(session, key):
|
||||
url = "http://0.0.0.0:4000/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
],
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_gen():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await chat_completion(session=session, key=key)
|
||||
115
tests/test_health.py
Normal file
115
tests/test_health.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
# What this tests?
|
||||
## Tests /health + /routes endpoints.
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
|
||||
async def health(session, call_key):
|
||||
url = "http://0.0.0.0:4000/health"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {call_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with session.get(url, headers=headers) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(f"Response (Status code: {status}):")
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def generate_key(session):
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["gpt-4", "text-embedding-ada-002", "dall-e-2"],
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health():
|
||||
"""
|
||||
- Call /health
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# as admin #
|
||||
all_healthy_models = await health(session=session, call_key="sk-1234")
|
||||
total_model_count = (
|
||||
all_healthy_models["healthy_count"] + all_healthy_models["unhealthy_count"]
|
||||
)
|
||||
assert total_model_count > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_readiness():
|
||||
"""
|
||||
Check if 200
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
url = "http://0.0.0.0:4000/health/readiness"
|
||||
async with session.get(url) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_liveliness():
|
||||
"""
|
||||
Check if 200
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
url = "http://0.0.0.0:4000/health/liveliness"
|
||||
async with session.get(url) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routes():
|
||||
"""
|
||||
Check if 200
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
url = "http://0.0.0.0:4000/routes"
|
||||
async with session.get(url) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
183
tests/test_keys.py
Normal file
183
tests/test_keys.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
# What this tests ?
|
||||
## Tests /key endpoints.
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
|
||||
async def generate_key(session, i):
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["azure-models"],
|
||||
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(f"Response {i} (Status code: {status}):")
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
||||
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_gen():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [generate_key(session, i) for i in range(1, 11)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
async def update_key(session, get_key):
|
||||
"""
|
||||
Make sure only models user has access to are returned
|
||||
"""
|
||||
url = "http://0.0.0.0:4000/key/update"
|
||||
headers = {
|
||||
"Authorization": f"Bearer sk-1234",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {"key": get_key, "models": ["gpt-4"]}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def chat_completion(session, key, model="gpt-4"):
|
||||
url = "http://0.0.0.0:4000/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
],
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_update():
|
||||
"""
|
||||
Create key
|
||||
Update key with new model
|
||||
Test key w/ model
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session, i=0)
|
||||
key = key_gen["key"]
|
||||
await update_key(
|
||||
session=session,
|
||||
get_key=key,
|
||||
)
|
||||
await chat_completion(session=session, key=key)
|
||||
|
||||
|
||||
async def delete_key(session, get_key):
|
||||
"""
|
||||
Delete key
|
||||
"""
|
||||
url = "http://0.0.0.0:4000/key/delete"
|
||||
headers = {
|
||||
"Authorization": f"Bearer sk-1234",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {"keys": [get_key]}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_delete():
|
||||
"""
|
||||
Delete key
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session, i=0)
|
||||
key = key_gen["key"]
|
||||
await delete_key(
|
||||
session=session,
|
||||
get_key=key,
|
||||
)
|
||||
|
||||
|
||||
async def get_key_info(session, get_key, call_key):
|
||||
"""
|
||||
Make sure only models user has access to are returned
|
||||
"""
|
||||
url = f"http://0.0.0.0:4000/key/info?key={get_key}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {call_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with session.get(url, headers=headers) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
if call_key != get_key:
|
||||
return status
|
||||
else:
|
||||
print(f"call_key: {call_key}; get_key: {get_key}")
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_info():
|
||||
"""
|
||||
Get key info
|
||||
- as admin -> 200
|
||||
- as key itself -> 200
|
||||
- as random key -> 403
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session, i=0)
|
||||
key = key_gen["key"]
|
||||
# as admin #
|
||||
await get_key_info(session=session, get_key=key, call_key="sk-1234")
|
||||
# as key itself #
|
||||
await get_key_info(session=session, get_key=key, call_key=key)
|
||||
# as random key #
|
||||
key_gen = await generate_key(session=session, i=0)
|
||||
random_key = key_gen["key"]
|
||||
status = await get_key_info(session=session, get_key=key, call_key=random_key)
|
||||
assert status == 403
|
||||
190
tests/test_models.py
Normal file
190
tests/test_models.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
# What this tests ?
|
||||
## Tests /models and /model/* endpoints
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
|
||||
async def generate_key(session, models=[]):
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": models,
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def get_models(session, key):
|
||||
url = "http://0.0.0.0:4000/models"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with session.get(url, headers=headers) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_models():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await get_models(session=session, key=key)
|
||||
|
||||
|
||||
async def add_models(session, model_id="123"):
|
||||
url = "http://0.0.0.0:4000/model/new"
|
||||
headers = {
|
||||
"Authorization": f"Bearer sk-1234",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
data = {
|
||||
"model_name": "azure-gpt-3.5",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": "os.environ/AZURE_API_KEY",
|
||||
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
|
||||
"api_version": "2023-05-15",
|
||||
},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(f"Add models {response_text}")
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
async def get_model_info(session, key):
|
||||
"""
|
||||
Make sure only models user has access to are returned
|
||||
"""
|
||||
url = "http://0.0.0.0:4000/model/info"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with session.get(url, headers=headers) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def chat_completion(session, key):
|
||||
url = "http://0.0.0.0:4000/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {
|
||||
"model": "azure-gpt-3.5",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
],
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_models():
|
||||
"""
|
||||
Add model
|
||||
Call new model
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await add_models(session=session)
|
||||
await chat_completion(session=session, key=key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_models():
|
||||
"""
|
||||
Get models user has access to
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session, models=["gpt-4"])
|
||||
key = key_gen["key"]
|
||||
response = await get_model_info(session=session, key=key)
|
||||
models = [m["model_name"] for m in response["data"]]
|
||||
for m in models:
|
||||
assert m == "gpt-4"
|
||||
|
||||
|
||||
async def delete_model(session, model_id="123"):
|
||||
"""
|
||||
Make sure only models user has access to are returned
|
||||
"""
|
||||
url = "http://0.0.0.0:4000/model/delete"
|
||||
headers = {
|
||||
"Authorization": f"Bearer sk-1234",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {"id": model_id}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_models():
|
||||
"""
|
||||
Get models user has access to
|
||||
"""
|
||||
model_id = "12345"
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await add_models(session=session, model_id=model_id)
|
||||
await chat_completion(session=session, key=key)
|
||||
await delete_model(session=session, model_id=model_id)
|
||||
201
tests/test_openai_endpoints.py
Normal file
201
tests/test_openai_endpoints.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
# What this tests ?
|
||||
## Tests /chat/completions by generating a key and then making a chat completions request
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
|
||||
async def generate_key(session):
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["gpt-4", "text-embedding-ada-002", "dall-e-2"],
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def new_user(session):
|
||||
url = "http://0.0.0.0:4000/user/new"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["gpt-4", "text-embedding-ada-002", "dall-e-2"],
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def chat_completion(session, key):
|
||||
url = "http://0.0.0.0:4000/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
],
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_completion():
|
||||
"""
|
||||
- Create key
|
||||
Make chat completion call
|
||||
- Create user
|
||||
make chat completion call
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await chat_completion(session=session, key=key)
|
||||
key_gen = await new_user(session=session)
|
||||
key_2 = key_gen["key"]
|
||||
await chat_completion(session=session, key=key_2)
|
||||
|
||||
|
||||
async def completion(session, key):
|
||||
url = "http://0.0.0.0:4000/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {"model": "gpt-4", "prompt": "Hello!"}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion():
|
||||
"""
|
||||
- Create key
|
||||
Make chat completion call
|
||||
- Create user
|
||||
make chat completion call
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await completion(session=session, key=key)
|
||||
key_gen = await new_user(session=session)
|
||||
key_2 = key_gen["key"]
|
||||
await completion(session=session, key=key_2)
|
||||
|
||||
|
||||
async def embeddings(session, key):
|
||||
url = "http://0.0.0.0:4000/embeddings"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {
|
||||
"model": "text-embedding-ada-002",
|
||||
"input": ["hello world"],
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embeddings():
|
||||
"""
|
||||
- Create key
|
||||
Make embeddings call
|
||||
- Create user
|
||||
make embeddings call
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await embeddings(session=session, key=key)
|
||||
key_gen = await new_user(session=session)
|
||||
key_2 = key_gen["key"]
|
||||
await embeddings(session=session, key=key_2)
|
||||
|
||||
|
||||
async def image_generation(session, key):
|
||||
url = "http://0.0.0.0:4000/images/generations"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
data = {
|
||||
"model": "dall-e-2",
|
||||
"prompt": "A cute baby sea otter",
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_generation():
|
||||
"""
|
||||
- Create key
|
||||
Make embeddings call
|
||||
- Create user
|
||||
make embeddings call
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await generate_key(session=session)
|
||||
key = key_gen["key"]
|
||||
await image_generation(session=session, key=key)
|
||||
key_gen = await new_user(session=session)
|
||||
key_2 = key_gen["key"]
|
||||
await image_generation(session=session, key=key_2)
|
||||
|
|
@ -1,33 +0,0 @@
|
|||
# What this tests ?
|
||||
## Tests /key/generate by making 10 parallel requests, and asserting all are successful
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
|
||||
async def generate_key(session, i):
|
||||
url = "http://0.0.0.0:4000/key/generate"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["azure-models"],
|
||||
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(f"Response {i} (Status code: {status}):")
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_gen():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [generate_key(session, i) for i in range(1, 11)]
|
||||
await asyncio.gather(*tasks)
|
||||
102
tests/test_users.py
Normal file
102
tests/test_users.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
# What this tests ?
|
||||
## Tests /user endpoints.
|
||||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import time
|
||||
|
||||
|
||||
async def new_user(session, i, user_id=None):
|
||||
url = "http://0.0.0.0:4000/user/new"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {
|
||||
"models": ["azure-models"],
|
||||
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
|
||||
"duration": None,
|
||||
}
|
||||
|
||||
if user_id is not None:
|
||||
data["user_id"] = user_id
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(f"Response {i} (Status code: {status}):")
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
||||
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_new():
|
||||
"""
|
||||
Make 20 parallel calls to /user/new. Assert all worked.
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [new_user(session, i) for i in range(1, 11)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
|
||||
async def get_user_info(session, get_user, call_user):
|
||||
"""
|
||||
Make sure only models user has access to are returned
|
||||
"""
|
||||
url = f"http://0.0.0.0:4000/user/info?key={get_user}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {call_user}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with session.get(url, headers=headers) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
if call_user != get_user:
|
||||
return status
|
||||
else:
|
||||
print(f"call_user: {call_user}; get_user: {get_user}")
|
||||
raise Exception(f"Request did not return a 200 status code: {status}")
|
||||
return await response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_info():
|
||||
"""
|
||||
Get user info
|
||||
- as admin
|
||||
- as user themself
|
||||
- as random
|
||||
"""
|
||||
get_user = f"krrish_{time.time()}@berri.ai"
|
||||
async with aiohttp.ClientSession() as session:
|
||||
key_gen = await new_user(session, 0, user_id=get_user)
|
||||
key = key_gen["key"]
|
||||
## as admin ##
|
||||
await get_user_info(session=session, get_user=get_user, call_user="sk-1234")
|
||||
## as user themself ##
|
||||
await get_user_info(session=session, get_user=get_user, call_user=key)
|
||||
# as random user #
|
||||
key_gen = await new_user(session=session, i=0)
|
||||
random_key = key_gen["key"]
|
||||
status = await get_user_info(
|
||||
session=session, get_user=get_user, call_user=random_key
|
||||
)
|
||||
assert status == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_update():
|
||||
"""
|
||||
Create user
|
||||
Update user access to new model
|
||||
Make chat completion call
|
||||
"""
|
||||
pass
|
||||
Loading…
Add table
Add a link
Reference in a new issue