phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 02:34:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(ci/cd) make prisma tests async

Browse source
This commit is contained in:
ishaan-jaff 2024-01-06 11:42:51 +05:30
parent e9350ad010
commit 0ebd0653c5
1 changed files with 71 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

112
litellm/tests/test_proxy_server_keys.py
View file

@ -14,6 +14,7 @@ import pytest, logging
import litellm import litellm
from litellm import embedding, completion, completion_cost, Timeout from litellm import embedding, completion, completion_cost, Timeout
from litellm import RateLimitError from litellm import RateLimitError
from httpx import AsyncClient
# Configure logging # Configure logging
logging.basicConfig( logging.basicConfig(
@ -74,19 +75,28 @@ def event_loop():
# Here you create a fixture that will be used by your tests # Here you create a fixture that will be used by your tests
# Make sure the fixture returns TestClient(app) # Make sure the fixture returns TestClient(app)
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
def client(): async def client():
from litellm.proxy.proxy_server import cleanup_router_config_variables, initialize from litellm.proxy.proxy_server import (
cleanup_router_config_variables,
initialize,
ProxyLogging,
proxy_logging_obj,
)
cleanup_router_config_variables() # rest proxy before test 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 = FastAPI()
app.include_router(router) # Include your router in the test app app.include_router(router) # Include your router in the test app
async with AsyncClient(app=app, base_url="http://testserver") as client:
return TestClient(app) 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: try:
# Your test data # Your test data
test_data = { test_data = {
@ -99,13 +109,13 @@ def test_add_new_key(client):
token = os.getenv("PROXY_MASTER_KEY") token = os.getenv("PROXY_MASTER_KEY")
headers = {"Authorization": f"Bearer {token}"} 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}") print(f"response: {response.text}")
assert response.status_code == 200 assert response.status_code == 200
result = response.json() result = response.json()
assert result["key"].startswith("sk-") assert result["key"].startswith("sk-")
def _post_data(): async def _post_data():
json_data = { json_data = {
"model": "azure-model", "model": "azure-model",
"messages": [ "messages": [
@ -115,20 +125,22 @@ def test_add_new_key(client):
} }
], ],
} }
response = client.post( response = await client.post(
"/chat/completions", "/chat/completions",
json=json_data, json=json_data,
headers={"Authorization": f"Bearer {result['key']}"}, headers={"Authorization": f"Bearer {result['key']}"},
) )
return response return response
_post_data() await _post_data()
print(f"Received response: {result}") print(f"Received response: {result}")
except Exception as e: except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(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: try:
# Your test data # Your test data
test_data = { test_data = {
@ -141,20 +153,20 @@ def test_update_new_key(client):
token = os.getenv("PROXY_MASTER_KEY") token = os.getenv("PROXY_MASTER_KEY")
headers = {"Authorization": f"Bearer {token}"} 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}") print(f"response: {response.text}")
assert response.status_code == 200 assert response.status_code == 200
result = response.json() result = response.json()
assert result["key"].startswith("sk-") assert result["key"].startswith("sk-")
def _post_data(): async def _post_data():
json_data = {"models": ["bedrock-models"], "key": result["key"]} 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}") print(f"response text: {response.text}")
assert response.status_code == 200 assert response.status_code == 200
return response return response
_post_data() await _post_data()
print(f"Received response: {result}") print(f"Received response: {result}")
except Exception as e: except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(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 # # 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: try:
import anyio
print("ANY IO BACKENDS")
print(anyio.get_all_backends())
# Your test data # 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 # Your bearer token
token = os.getenv("PROXY_MASTER_KEY") token = os.getenv("PROXY_MASTER_KEY")
headers = {"Authorization": f"Bearer {token}"} 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}") print(f"response: {response.text}")
assert response.status_code == 200 assert response.status_code == 200
result = response.json() result = response.json()
def _post_data(): async def _post_data():
json_data = { json_data = {
"model": "azure-model", "model": "azure-model",
"messages": [ "messages": [
@ -186,32 +209,38 @@ def test_add_new_key_max_parallel_limit(client):
} }
], ],
} }
response = client.post(
response = await client.post(
"/chat/completions", "/chat/completions",
json=json_data, json=json_data,
headers={"Authorization": f"Bearer {result['key']}"}, headers={"Authorization": f"Bearer {result['key']}"},
) )
return response return response
def _run_in_parallel(): async def _run_in_parallel():
with ThreadPoolExecutor(max_workers=2) as executor: try:
future1 = executor.submit(_post_data) futures = [_post_data() for _ in range(2)]
future2 = executor.submit(_post_data) 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 if any(response.status_code == 429 for response in responses):
response1 = future1.result()
response2 = future2.result()
if response1.status_code == 429 or response2.status_code == 429:
pass pass
else: else:
raise Exception() 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: except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(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: try:
# Your test data # Your test data
test_data = {"duration": "20m", "max_parallel_requests": 1} 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") token = os.getenv("PROXY_MASTER_KEY")
headers = {"Authorization": f"Bearer {token}"} 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}") print(f"response: {response.text}")
assert response.status_code == 200 assert response.status_code == 200
result = response.json() result = response.json()
def _post_data(): async def _post_data():
json_data = { json_data = {
"model": "azure-model", "model": "azure-model",
"messages": [ "messages": [
@ -235,25 +264,26 @@ def test_add_new_key_max_parallel_limit_streaming(client):
], ],
"stream": True, "stream": True,
} }
response = client.post( response = await client.post(
"/chat/completions", "/chat/completions",
json=json_data, json=json_data,
headers={"Authorization": f"Bearer {result['key']}"}, headers={"Authorization": f"Bearer {result['key']}"},
) )
return response return response
def _run_in_parallel(): async def _run_in_parallel():
with ThreadPoolExecutor(max_workers=2) as executor: try:
future1 = executor.submit(_post_data) futures = [_post_data() for _ in range(2)]
future2 = executor.submit(_post_data) 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 if any(response.status_code == 429 for response in responses):
response1 = future1.result()
response2 = future2.result()
if response1.status_code == 429 or response2.status_code == 429:
pass pass
else: else:
raise Exception() raise Exception()
except Exception as e:
pass
_run_in_parallel() _run_in_parallel()
except Exception as e: except Exception as e: