From e4ba779d6a1a32ce554dc0ac0019eef50cab5ded Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 16 Jan 2024 16:03:42 -0800 Subject: [PATCH] (ci/cd) fixes --- litellm/proxy/proxy_server.py | 4 +++- litellm/tests/test_proxy_startup.py | 30 ++++++++++++++++++----------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d3d63dd59f..353c3157db 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2727,7 +2727,7 @@ async def shutdown_event(): 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 master_key = None @@ -2737,6 +2737,8 @@ def cleanup_router_config_variables(): user_custom_auth_path = None use_background_health_checks = None health_check_interval = None + prisma_client = None + custom_db_client = None app.include_router(router) diff --git a/litellm/tests/test_proxy_startup.py b/litellm/tests/test_proxy_startup.py index 1183e8c616..c033505bbd 100644 --- a/litellm/tests/test_proxy_startup.py +++ b/litellm/tests/test_proxy_startup.py @@ -1,4 +1,4 @@ -# What this tests +# What this tests ## This tests the proxy server startup import sys, os, json import traceback @@ -20,45 +20,53 @@ from litellm.proxy.proxy_server import ( initialize, startup_event, llm_model_list, - shutdown_event + shutdown_event, ) + def test_proxy_gunicorn_startup_direct_config(): """ 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 """ - try: + try: + from litellm.proxy.proxy_server import cleanup_router_config_variables + + cleanup_router_config_variables() 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" os.environ["WORKER_CONFIG"] = config_fp asyncio.run(startup_event()) asyncio.run(shutdown_event()) 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 else: pytest.fail(f"An exception occurred - {str(e)}") + 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__)) - # test with worker_config = config yaml + # test with worker_config = config 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} os.environ["WORKER_CONFIG"] = json.dumps(worker_config) asyncio.run(startup_event()) asyncio.run(shutdown_event()) 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 else: pytest.fail(f"An exception occurred - {str(e)}") -# test_proxy_gunicorn_startup() \ No newline at end of file +# test_proxy_gunicorn_startup()