fix(user_api_key_auth.py): handle older user_role's

Fixes issue where older user_role's (e.g. app_user) weren't being recognized. + Adds testing for it
This commit is contained in:
Krrish Dholakia 2024-08-05 08:57:06 -07:00
parent 51211dc2a6
commit f610fba58f
3 changed files with 77 additions and 7 deletions

View file

@ -93,3 +93,26 @@ async def test_check_blocked_team():
request._url = URL(url="/chat/completions")
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
@pytest.mark.parametrize(
"user_role", ["app_user", "internal_user", "proxy_admin_viewer"]
)
def test_returned_user_api_key_auth(user_role):
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj
user_id_information = [{"user_role": user_role}]
new_obj = _return_user_api_key_auth_obj(
user_id_information,
api_key="hello-world",
parent_otel_span=None,
valid_token_dict={},
route="/chat/completion",
)
if user_role in list(LitellmUserRoles.__annotations__.keys()):
assert new_obj.user_role == user_role
else:
assert new_obj.user_role == "internal_user"