From e2a79101134b622d37f37e403fde93937d8d389b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 29 Jan 2024 16:53:27 -0800 Subject: [PATCH] (test) dynamoDB backwards compat --- litellm/tests/test_key_generate_dynamodb.py | 28 +++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index f482f35437f..4a715b34232 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/litellm/tests/test_key_generate_dynamodb.py @@ -472,3 +472,31 @@ def test_call_with_key_over_budget_stream(custom_db_client): error_detail = e.message assert "Authentication Error, ExceededTokenBudget:" in error_detail print(vars(e)) + + +def test_dynamo_db_migration(custom_db_client): + # Tests the temporary patch we have in place + setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + try: + + async def test(): + bearer_token = ( + "Bearer " + "sk-elJDL2pOEjcAuC7zD4psAg" + ) # this works with ishaan's db, it's a never expiring key + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body(): + return b'{"model": "azure-models"}' + + request.body = return_body + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print("result from user auth with new key", result) + + asyncio.run(test()) + except Exception as e: + pytest.fail(f"An exception occurred - {str(e)}")