From d566421bc23c6f9ec214476ee81abe932d3a10da Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 16:47:31 -0800
Subject: [PATCH] (test) test 8 key/gen, user_auth scenarios

---
 litellm/tests/test_key_generate_dynamodb.py | 75 +++++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 05b9daadf2d..609f38b8ff1 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -7,3 +7,78 @@
 # 6. Make a call with unexpired key - expect to pass
 # 7. Make a call with key under budget, expect to pass
 # 8. Make a call with key over budget, expect to fail
+
+
+# function to call to generate key - async def new_user(data: NewUserRequest):
+# function to validate a request - async def user_auth(request: Request):
+
+import sys, os
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import os, io
+
+# this file is to test litellm/proxy
+
+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
+
+
+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
+
+
+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_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_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_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