From 086f827af83fbe491c933b33efd6433bb500ff59 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 17:21:50 -0800
Subject: [PATCH] (v0) add test for basic key gen with dynamo

---
 litellm/tests/test_key_generate_dynamodb.py | 119 +++++++++++++-------
 1 file changed, 80 insertions(+), 39 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 609f38b8ff..4ba73eeb95 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -15,6 +15,7 @@
 import sys, os
 import traceback
 from dotenv import load_dotenv
+from fastapi import Request
 
 load_dotenv()
 import os, io
@@ -25,60 +26,100 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import pytest, logging, asyncio
-import litellm
-from litellm.proxy.proxy_server import new_user, user_auth
+import litellm, asyncio
+from litellm.proxy.proxy_server import new_user, user_api_key_auth
+
+from litellm.proxy._types import NewUserRequest, DynamoDBArgs
+from litellm.proxy.utils import DBClient
+
+db_args = DynamoDBArgs(
+    ssl_verify=False,
+    billing_mode="PAY_PER_REQUEST",
+    region_name="us-west-2",
+)
+db_args_dict = db_args.model_dump()
+custom_db_client = DBClient(
+    custom_db_type="dynamo_db",
+    custom_db_args=db_args_dict,
+)
+
+request_data = {
+    "model": "azure-gpt-3.5",
+    "messages": [
+        {"role": "user", "content": "this is my new test. respond in 50 lines"}
+    ],
+}
 
 
 def test_generate_and_call_with_valid_key():
     # 1. Generate a Key, and use it to make a call
-    key = new_user(ValidNewUserRequest())
-    result = user_auth(ValidRequest(key))
-    assert result is True
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    try:
+
+        async def test():
+            request = NewUserRequest()
+            key = await new_user(request)
+            print(key)
+
+            generated_key = key.key
+            bearer_token = "Bearer " + generated_key
+
+            request = Request(scope={"type": "http"})
+
+            # 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())
+    # result = user_auth(ValidRequest(key))
+    # assert result is True
+    except Exception as e:
+        pytest.fail(f"An exception occurred - {str(e)}")
 
 
-def test_call_with_invalid_key():
-    # 2. Make a call with invalid key, expect it to fail
-    result = user_auth(InvalidKeyRequest())
-    assert result is False
+# def test_call_with_invalid_key():
+#     # 2. Make a call with invalid key, expect it to fail
+#     result = user_auth(InvalidKeyRequest())
+#     assert result is False
 
 
-def test_call_with_invalid_model():
-    # 3. Make a call to a key with an invalid model - expect to fail
-    key = new_user(ValidNewUserRequest())
-    result = user_auth(InvalidModelRequest(key))
-    assert result is False
+# def test_call_with_invalid_model():
+#     # 3. Make a call to a key with an invalid model - expect to fail
+#     key = new_user(ValidNewUserRequest())
+#     result = user_auth(InvalidModelRequest(key))
+#     assert result is False
 
 
-def test_call_with_valid_model():
-    # 4. Make a call to a key with a valid model - expect to pass
-    key = new_user(ValidNewUserRequest())
-    result = user_auth(ValidModelRequest(key))
-    assert result is True
+# def test_call_with_valid_model():
+#     # 4. Make a call to a key with a valid model - expect to pass
+#     key = new_user(ValidNewUserRequest())
+#     result = user_auth(ValidModelRequest(key))
+#     assert result is True
 
 
-def test_call_with_expired_key():
-    # 5. Make a call with an expired key - expect to fail
-    key = new_user(ExpiredKeyRequest())
-    result = user_auth(ValidRequest(key))
-    assert result is False
+# def test_call_with_expired_key():
+#     # 5. Make a call with an expired key - expect to fail
+#     key = new_user(ExpiredKeyRequest())
+#     result = user_auth(ValidRequest(key))
+#     assert result is False
 
 
-def test_call_with_unexpired_key():
-    # 6. Make a call with an unexpired key - expect to pass
-    key = new_user(UnexpiredKeyRequest())
-    result = user_auth(ValidRequest(key))
-    assert result is True
+# def test_call_with_unexpired_key():
+#     # 6. Make a call with an unexpired key - expect to pass
+#     key = new_user(UnexpiredKeyRequest())
+#     result = user_auth(ValidRequest(key))
+#     assert result is True
 
 
-def test_call_with_key_under_budget():
-    # 7. Make a call with a key under budget, expect to pass
-    key = new_user(KeyUnderBudgetRequest())
-    result = user_auth(ValidRequest(key))
-    assert result is True
+# def test_call_with_key_under_budget():
+#     # 7. Make a call with a key under budget, expect to pass
+#     key = new_user(KeyUnderBudgetRequest())
+#     result = user_auth(ValidRequest(key))
+#     assert result is True
 
 
-def test_call_with_key_over_budget():
-    # 8. Make a call with a key over budget, expect to fail
-    key = new_user(KeyOverBudgetRequest())
-    result = user_auth(ValidRequest(key))
-    assert result is False
+# def test_call_with_key_over_budget():
+#     # 8. Make a call with a key over budget, expect to fail
+#     key = new_user(KeyOverBudgetRequest())
+#     result = user_auth(ValidRequest(key))
+#     assert result is False