phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix: fix proxy testing

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-05 11:13:09 -08:00
parent a602d59645
commit b46c73a46e
4 changed files with 73 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

35
litellm/tests/test_proxy_server.py
View file

@ -18,11 +18,22 @@ from litellm import RateLimitError
# test /chat/completion request to the proxy
from fastapi.testclient import TestClient
from fastapi import FastAPI
from litellm.proxy.proxy_server import router # Replace with the actual module where your FastAPI router is defined
from litellm.proxy.proxy_server import router, save_worker_config, initialize # Replace with the actual module where your FastAPI router is defined
save_worker_config(config=None, model=None, alias=None, api_base=None, api_version=None, debug=False, temperature=None, max_tokens=None, request_timeout=600, max_budget=None, telemetry=False, drop_params=True, add_function_to_prompt=False, headers=None, save=False, use_queue=False)
app = FastAPI()
app.include_router(router) # Include your router in the test app
client = TestClient(app)
def test_chat_completion():
@app.on_event("startup")
async def wrapper_startup_event(): # required to reset config on app init - b/c pytest collects across multiple files - which sets the fastapi client + WORKER CONFIG to whatever was collected last
initialize(config=None, model=None, alias=None, api_base=None, api_version=None, debug=False, temperature=None, max_tokens=None, request_timeout=600, max_budget=None, telemetry=False, drop_params=True, add_function_to_prompt=False, headers=None, save=False, use_queue=False)
# Here you create a fixture that will be used by your tests
# Make sure the fixture returns TestClient(app)
@pytest.fixture(autouse=True)
def client():
with TestClient(app) as client:
yield client
def test_chat_completion(client):
try:
# Your test data
test_data = {
@ -37,18 +48,16 @@ def test_chat_completion():
}
print("testing proxy server")
response = client.post("/v1/chat/completions", json=test_data)
print(f"response - {response.text}")
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test
# test_chat_completion()
def test_chat_completion_azure():
def test_chat_completion_azure(client):
try:
# Your test data
test_data = {
@ -69,13 +78,13 @@ def test_chat_completion_azure():
print(f"Received response: {result}")
assert len(result["choices"][0]["message"]["content"]) > 0
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test
# test_chat_completion_azure()
def test_embedding():
def test_embedding(client):
try:
test_data = {
"model": "azure/azure-embedding-model",
@ -89,13 +98,13 @@ def test_embedding():
print(len(result["data"][0]["embedding"]))
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
except Exception as e:
pytest.fail("LiteLLM Proxy test failed. Exception", e)
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test
# test_embedding()
def test_add_new_model():
def test_add_new_model(client):
try:
test_data = {
"model_name": "test_openai_models",
@ -135,7 +144,7 @@ class MyCustomHandler(CustomLogger):
customHandler = MyCustomHandler()
def test_chat_completion_optional_params():
def test_chat_completion_optional_params(client):
# [PROXY: PROD TEST] - DO NOT DELETE
# This tests if all the /chat/completion params are passed to litellm