From 917f92800de03322682a00ce68dc8e91aee7ad86 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Mar 2024 15:24:42 -0700 Subject: [PATCH] test(test_key_generate_prisma.py): fix tests --- litellm/tests/test_key_generate_prisma.py | 43 ++++++++++++++++++++--- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 91c6d24db..62f6c38a9 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -788,7 +788,19 @@ def test_delete_key_auth(prisma_client): delete_key_request = KeyRequest(keys=[generated_key]) # delete the key - result_delete_key = await delete_key_fn(data=delete_key_request) + bearer_token = "Bearer sk-1234" + + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" + + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) print("result from delete key", result_delete_key) assert result_delete_key == {"deleted_keys": [generated_key]} @@ -803,6 +815,7 @@ def test_delete_key_auth(prisma_client): ) # use generated key to auth in + bearer_token = "Bearer " + generated_key result = await user_api_key_auth(request=request, api_key=bearer_token) print("got result", result) pytest.fail(f"This should have failed!. IT's an invalid key") @@ -847,9 +860,19 @@ def test_generate_and_call_key_info(prisma_client): # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) + bearer_token = "Bearer sk-1234" - # delete the key - await delete_key_fn(data=delete_key_request) + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" + + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) asyncio.run(test()) except Exception as e: @@ -928,7 +951,19 @@ def test_generate_and_update_key(prisma_client): delete_key_request = KeyRequest(keys=[generated_key]) # delete the key - await delete_key_fn(data=delete_key_request) + bearer_token = "Bearer sk-1234" + + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" + + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) asyncio.run(test()) except Exception as e: