forked from phoenix/litellm-mirror
(ci/cd) make prisma tests async
This commit is contained in:
ishaan-jaff
parent
e9350ad010
commit
0ebd0653c5
1 changed files with 71 additions and 41 deletions
|
|
@ -14,6 +14,7 @@ import pytest, logging
|
|||
import litellm
|
||||
from litellm import embedding, completion, completion_cost, Timeout
|
||||
from litellm import RateLimitError
|
||||
from httpx import AsyncClient
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
|
|
@ -74,19 +75,28 @@ def event_loop():
|
|||
# Here you create a fixture that will be used by your tests
|
||||
# Make sure the fixture returns TestClient(app)
|
||||
@pytest.fixture(scope="function")
|
||||
def client():
|
||||
from litellm.proxy.proxy_server import cleanup_router_config_variables, initialize
|
||||
async def client():
|
||||
from litellm.proxy.proxy_server import (
|
||||
cleanup_router_config_variables,
|
||||
initialize,
|
||||
ProxyLogging,
|
||||
proxy_logging_obj,
|
||||
)
|
||||
|
||||
cleanup_router_config_variables() # rest proxy before test
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache={})
|
||||
proxy_logging_obj._init_litellm_callbacks() # INITIALIZE LITELLM CALLBACKS ON SERVER STARTUP <- do this to catch any logging errors on startup, not when calls are being made
|
||||
|
||||
asyncio.run(initialize(config=config_fp, debug=True))
|
||||
await initialize(config=config_fp, debug=True)
|
||||
app = FastAPI()
|
||||
app.include_router(router) # Include your router in the test app
|
||||
|
||||
return TestClient(app)
|
||||
async with AsyncClient(app=app, base_url="http://testserver") as client:
|
||||
yield client
|
||||
|
||||
|
||||
def test_add_new_key(client):
|
||||
@pytest.mark.parametrize("anyio_backend", ["asyncio"])
|
||||
@pytest.mark.anyio
|
||||
async def test_add_new_key(client):
|
||||
try:
|
||||
# Your test data
|
||||
test_data = {
|
||||
|
|
@ -99,13 +109,13 @@ def test_add_new_key(client):
|
|||
token = os.getenv("PROXY_MASTER_KEY")
|
||||
|
||||
headers = {"Authorization": f"Bearer {token}"}
|
||||
response = client.post("/key/generate", json=test_data, headers=headers)
|
||||
response = await client.post("/key/generate", json=test_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["key"].startswith("sk-")
|
||||
|
||||
def _post_data():
|
||||
async def _post_data():
|
||||
json_data = {
|
||||
"model": "azure-model",
|
||||
"messages": [
|
||||
|
|
@ -115,20 +125,22 @@ def test_add_new_key(client):
|
|||
}
|
||||
],
|
||||
}
|
||||
response = client.post(
|
||||
response = await client.post(
|
||||
"/chat/completions",
|
||||
json=json_data,
|
||||
headers={"Authorization": f"Bearer {result['key']}"},
|
||||
)
|
||||
return response
|
||||
|
||||
_post_data()
|
||||
await _post_data()
|
||||
print(f"Received response: {result}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
|
||||
|
||||
|
||||
def test_update_new_key(client):
|
||||
@pytest.mark.parametrize("anyio_backend", ["asyncio"])
|
||||
@pytest.mark.anyio
|
||||
async def test_update_new_key(client):
|
||||
try:
|
||||
# Your test data
|
||||
test_data = {
|
||||
|
|
@ -141,20 +153,20 @@ def test_update_new_key(client):
|
|||
token = os.getenv("PROXY_MASTER_KEY")
|
||||
|
||||
headers = {"Authorization": f"Bearer {token}"}
|
||||
response = client.post("/key/generate", json=test_data, headers=headers)
|
||||
response = await client.post("/key/generate", json=test_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["key"].startswith("sk-")
|
||||
|
||||
def _post_data():
|
||||
async def _post_data():
|
||||
json_data = {"models": ["bedrock-models"], "key": result["key"]}
|
||||
response = client.post("/key/update", json=json_data, headers=headers)
|
||||
response = await client.post("/key/update", json=json_data, headers=headers)
|
||||
print(f"response text: {response.text}")
|
||||
assert response.status_code == 200
|
||||
return response
|
||||
|
||||
_post_data()
|
||||
await _post_data()
|
||||
print(f"Received response: {result}")
|
||||
except Exception as e:
|
||||
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
|
||||
|
|
@ -163,20 +175,31 @@ def test_update_new_key(client):
|
|||
# # Run the test - only runs via pytest
|
||||
|
||||
|
||||
def test_add_new_key_max_parallel_limit(client):
|
||||
@pytest.mark.parametrize("anyio_backend", ["asyncio"])
|
||||
@pytest.mark.anyio
|
||||
async def test_add_new_key_max_parallel_limit(client):
|
||||
try:
|
||||
import anyio
|
||||
|
||||
print("ANY IO BACKENDS")
|
||||
print(anyio.get_all_backends())
|
||||
# Your test data
|
||||
test_data = {"duration": "20m", "max_parallel_requests": 1}
|
||||
test_data = {
|
||||
"duration": "20m",
|
||||
"max_parallel_requests": 1,
|
||||
"metadata": {"type": "ishaan-test"},
|
||||
}
|
||||
# Your bearer token
|
||||
token = os.getenv("PROXY_MASTER_KEY")
|
||||
|
||||
headers = {"Authorization": f"Bearer {token}"}
|
||||
response = client.post("/key/generate", json=test_data, headers=headers)
|
||||
|
||||
response = await client.post("/key/generate", json=test_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
def _post_data():
|
||||
async def _post_data():
|
||||
json_data = {
|
||||
"model": "azure-model",
|
||||
"messages": [
|
||||
|
|
@ -186,32 +209,38 @@ def test_add_new_key_max_parallel_limit(client):
|
|||
}
|
||||
],
|
||||
}
|
||||
response = client.post(
|
||||
|
||||
response = await client.post(
|
||||
"/chat/completions",
|
||||
json=json_data,
|
||||
headers={"Authorization": f"Bearer {result['key']}"},
|
||||
)
|
||||
return response
|
||||
|
||||
def _run_in_parallel():
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
future1 = executor.submit(_post_data)
|
||||
future2 = executor.submit(_post_data)
|
||||
async def _run_in_parallel():
|
||||
try:
|
||||
futures = [_post_data() for _ in range(2)]
|
||||
responses = await asyncio.gather(*futures)
|
||||
print("response1 status: ", responses[0].status_code)
|
||||
print("response2 status: ", responses[1].status_code)
|
||||
|
||||
# Obtain the results from the futures
|
||||
response1 = future1.result()
|
||||
response2 = future2.result()
|
||||
if response1.status_code == 429 or response2.status_code == 429:
|
||||
if any(response.status_code == 429 for response in responses):
|
||||
pass
|
||||
else:
|
||||
raise Exception()
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
_run_in_parallel()
|
||||
await _run_in_parallel()
|
||||
|
||||
# assert responses[0].status_code == 200 or responses[1].status_code == 200
|
||||
except Exception as e:
|
||||
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
|
||||
|
||||
|
||||
def test_add_new_key_max_parallel_limit_streaming(client):
|
||||
@pytest.mark.parametrize("anyio_backend", ["asyncio"])
|
||||
@pytest.mark.anyio
|
||||
async def test_add_new_key_max_parallel_limit_streaming(client):
|
||||
try:
|
||||
# Your test data
|
||||
test_data = {"duration": "20m", "max_parallel_requests": 1}
|
||||
|
|
@ -219,12 +248,12 @@ def test_add_new_key_max_parallel_limit_streaming(client):
|
|||
token = os.getenv("PROXY_MASTER_KEY")
|
||||
|
||||
headers = {"Authorization": f"Bearer {token}"}
|
||||
response = client.post("/key/generate", json=test_data, headers=headers)
|
||||
response = await client.post("/key/generate", json=test_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
def _post_data():
|
||||
async def _post_data():
|
||||
json_data = {
|
||||
"model": "azure-model",
|
||||
"messages": [
|
||||
|
|
@ -235,25 +264,26 @@ def test_add_new_key_max_parallel_limit_streaming(client):
|
|||
],
|
||||
"stream": True,
|
||||
}
|
||||
response = client.post(
|
||||
response = await client.post(
|
||||
"/chat/completions",
|
||||
json=json_data,
|
||||
headers={"Authorization": f"Bearer {result['key']}"},
|
||||
)
|
||||
return response
|
||||
|
||||
def _run_in_parallel():
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
future1 = executor.submit(_post_data)
|
||||
future2 = executor.submit(_post_data)
|
||||
async def _run_in_parallel():
|
||||
try:
|
||||
futures = [_post_data() for _ in range(2)]
|
||||
responses = await asyncio.gather(*futures)
|
||||
print("response1 status: ", responses[0].status_code)
|
||||
print("response2 status: ", responses[1].status_code)
|
||||
|
||||
# Obtain the results from the futures
|
||||
response1 = future1.result()
|
||||
response2 = future2.result()
|
||||
if response1.status_code == 429 or response2.status_code == 429:
|
||||
if any(response.status_code == 429 for response in responses):
|
||||
pass
|
||||
else:
|
||||
raise Exception()
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
_run_in_parallel()
|
||||
except Exception as e:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue