From 2ffd0570422c828cf66d29bd5033070a4441e511 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 16 Apr 2024 18:13:40 -0700 Subject: [PATCH] test(test_models.py): ensure only admin can call `/health` --- litellm/proxy/proxy_server.py | 2 +- litellm/router.py | 2 +- tests/test_models.py | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04ef98376..6deaee894 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2622,7 +2622,7 @@ class ProxyConfig: # router settings _router_settings = config_data.get("router_settings", {}) - llm_router.set_settings(**_router_settings) + llm_router.update_settings(**_router_settings) except Exception as e: verbose_proxy_logger.error( "{}\nTraceback:{}".format(str(e), traceback.format_exc()) diff --git a/litellm/router.py b/litellm/router.py index c58d62673..abbd6343b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2366,7 +2366,7 @@ class Router: _settings_to_return[var] = _all_vars[var] return _settings_to_return - def set_settings(self, **kwargs): + def update_settings(self, **kwargs): # only the following settings are allowed to be configured _allowed_settings = [ "routing_strategy_args", diff --git a/tests/test_models.py b/tests/test_models.py index 4f7c81938..b009ded6c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -280,6 +280,7 @@ async def test_add_model_run_health(): async with aiohttp.ClientSession() as session: key_gen = await generate_key(session=session) key = key_gen["key"] + master_key = "sk-1234" model_id = str(uuid.uuid4()) model_name = f"azure-model-health-check-{model_id}" print("adding model", model_name) @@ -295,7 +296,7 @@ async def test_add_model_run_health(): print("calling /health?model=", model_name) _health_info = await get_model_health( - session=session, key=key, model_name=model_name + session=session, key=master_key, model_name=model_name ) _healthy_endpooint = _health_info["healthy_endpoints"][0]