(ci/cd) fixes

This commit is contained in:
ishaan-jaff 2024-01-16 16:03:42 -08:00
parent 90c8698b64
commit e4ba779d6a
2 changed files with 22 additions and 12 deletions

View file

@ -2727,7 +2727,7 @@ async def shutdown_event():
def cleanup_router_config_variables(): def cleanup_router_config_variables():
global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, use_background_health_checks, health_check_interval global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, use_background_health_checks, health_check_interval, prisma_client, custom_db_client
# Set all variables to None # Set all variables to None
master_key = None master_key = None
@ -2737,6 +2737,8 @@ def cleanup_router_config_variables():
user_custom_auth_path = None user_custom_auth_path = None
use_background_health_checks = None use_background_health_checks = None
health_check_interval = None health_check_interval = None
prisma_client = None
custom_db_client = None
app.include_router(router) app.include_router(router)

View file

@ -1,4 +1,4 @@
# What this tests # What this tests
## This tests the proxy server startup ## This tests the proxy server startup
import sys, os, json import sys, os, json
import traceback import traceback
@ -20,45 +20,53 @@ from litellm.proxy.proxy_server import (
initialize, initialize,
startup_event, startup_event,
llm_model_list, llm_model_list,
shutdown_event shutdown_event,
) )
def test_proxy_gunicorn_startup_direct_config(): def test_proxy_gunicorn_startup_direct_config():
""" """
gunicorn startup requires the config to be passed in via environment variables gunicorn startup requires the config to be passed in via environment variables
We support saving either the config or the dict as an environment variable. We support saving either the config or the dict as an environment variable.
Test both approaches Test both approaches
""" """
try: try:
from litellm.proxy.proxy_server import cleanup_router_config_variables
cleanup_router_config_variables()
filepath = os.path.dirname(os.path.abspath(__file__)) filepath = os.path.dirname(os.path.abspath(__file__))
# test with worker_config = config yaml # test with worker_config = config yaml
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
os.environ["WORKER_CONFIG"] = config_fp os.environ["WORKER_CONFIG"] = config_fp
asyncio.run(startup_event()) asyncio.run(startup_event())
asyncio.run(shutdown_event()) asyncio.run(shutdown_event())
except Exception as e: except Exception as e:
if "Already connected to the query engine" in str(e): if "Already connected to the query engine" in str(e):
pass pass
else: else:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
def test_proxy_gunicorn_startup_config_dict(): def test_proxy_gunicorn_startup_config_dict():
try: try:
from litellm.proxy.proxy_server import cleanup_router_config_variables
cleanup_router_config_variables()
filepath = os.path.dirname(os.path.abspath(__file__)) filepath = os.path.dirname(os.path.abspath(__file__))
# test with worker_config = config yaml # test with worker_config = config yaml
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
# test with worker_config = dict # test with worker_config = dict
worker_config = {"config": config_fp} worker_config = {"config": config_fp}
os.environ["WORKER_CONFIG"] = json.dumps(worker_config) os.environ["WORKER_CONFIG"] = json.dumps(worker_config)
asyncio.run(startup_event()) asyncio.run(startup_event())
asyncio.run(shutdown_event()) asyncio.run(shutdown_event())
except Exception as e: except Exception as e:
if "Already connected to the query engine" in str(e): if "Already connected to the query engine" in str(e):
pass pass
else: else:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")
# test_proxy_gunicorn_startup() # test_proxy_gunicorn_startup()