From 8455aa4bd4e0bacc05541fa2b6ecb77af14fc359 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 3 Feb 2024 11:06:51 -0800 Subject: [PATCH] (test) setting user_role --- litellm/tests/test_key_generate_prisma.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 140d638e88..734a0b1148 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -47,6 +47,7 @@ from litellm.proxy.proxy_server import ( spend_user_fn, spend_key_fn, view_spend_logs, + user_info, ) from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token from litellm._logging import verbose_proxy_logger @@ -106,9 +107,17 @@ def test_generate_and_call_with_valid_key(prisma_client): await litellm.proxy.proxy_server.prisma_client.connect() from litellm.proxy.proxy_server import user_api_key_cache - request = NewUserRequest() + request = NewUserRequest(user_role="app_owner") key = await new_user(request) print(key) + user_id = key.user_id + + # check /user/info to verify user_role was set correctly + new_user_info = await user_info(user_id=user_id) + new_user_info = new_user_info["user_info"] + print("new_user_info=", new_user_info) + assert new_user_info.user_role == "app_owner" + assert new_user_info.user_id == user_id generated_key = key.key bearer_token = "Bearer " + generated_key