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

(test) proxy

Browse source
This commit is contained in:
ishaan-jaff 2023-12-11 17:57:58 -08:00
parent b6798c703a
commit 634d301cae
1 changed files with 17 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

37
litellm/tests/test_proxy_server.py
View file

@ -24,14 +24,6 @@ logging.basicConfig(
from fastapi.testclient import TestClient from fastapi.testclient import TestClient
from fastapi import FastAPI from fastapi import FastAPI
from litellm.proxy.proxy_server import router, save_worker_config, initialize # 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
filepath = os.path.dirname(os.path.abspath(__file__))
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
save_worker_config(config=config_fp, 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
@app.on_event("startup")
async def wrapper_startup_event():
initialize(config=config_fp)
# Your bearer token # Your bearer token
token = os.getenv("PROXY_MASTER_KEY") token = os.getenv("PROXY_MASTER_KEY")
@ -40,14 +32,16 @@ headers = {
"Authorization": f"Bearer {token}" "Authorization": f"Bearer {token}"
} }
# Here you create a fixture that will be used by your tests def get_client(config_fp):
# Make sure the fixture returns TestClient(app) filepath = os.path.dirname(os.path.abspath(__file__))
@pytest.fixture(autouse=True) config_fp = f"{filepath}/test_configs/{config_fp}"
def client(): initialize(config=config_fp)
with TestClient(app) as client: app = FastAPI()
yield client app.include_router(router) # Include your router in the test app
return TestClient(app)
def test_chat_completion(client): def test_chat_completion():
client = get_client(config_fp="test_config_no_auth.yaml")
global headers global headers
try: try:
# Your test data # Your test data
@ -73,7 +67,8 @@ def test_chat_completion(client):
# Run the test # Run the test
def test_chat_completion_azure(client): def test_chat_completion_azure():
client = get_client(config_fp="test_config_no_auth.yaml")
global headers global headers
try: try:
# Your test data # Your test data
@ -102,7 +97,8 @@ def test_chat_completion_azure(client):
# test_chat_completion_azure() # test_chat_completion_azure()
def test_embedding(client): def test_embedding():
client = get_client(config_fp="test_config_no_auth.yaml")
global headers global headers
try: try:
test_data = { test_data = {
@ -123,7 +119,8 @@ def test_embedding(client):
# test_embedding() # test_embedding()
# @pytest.mark.skip(reason="hitting yaml load issues on circle-ci") # @pytest.mark.skip(reason="hitting yaml load issues on circle-ci")
def test_add_new_model(client): def test_add_new_model():
client = get_client(config_fp="test_config_no_auth.yaml")
global headers global headers
try: try:
test_data = { test_data = {
@ -164,10 +161,10 @@ class MyCustomHandler(CustomLogger):
customHandler = MyCustomHandler() customHandler = MyCustomHandler()
def test_chat_completion_optional_params(client): def test_chat_completion_optional_params():
# [PROXY: PROD TEST] - DO NOT DELETE # [PROXY: PROD TEST] - DO NOT DELETE
# This tests if all the /chat/completion params are passed to litellm # This tests if all the /chat/completion params are passed to litellm
client = get_client(config_fp="test_config_no_auth.yaml")
try: try:
# Your test data # Your test data
litellm.set_verbose=True litellm.set_verbose=True