From d59ac23e827f82e8031f3d202d74fae27ada1242 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 16:13:55 -0800
Subject: [PATCH 01/23] (fix) proxy, dynamo - allow users to set ssl_verify
 False

---
 litellm/proxy/_types.py         |  1 +
 litellm/proxy/db/dynamo_db.py   | 30 ++++++++++++++++++++++++++----
 litellm/proxy/proxy_config.yaml |  3 ++-
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 77f9aabced..039db119cb 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -193,6 +193,7 @@ class DynamoDBArgs(LiteLLMBase):
     billing_mode: Literal["PROVISIONED_THROUGHPUT", "PAY_PER_REQUEST"]
     read_capacity_units: Optional[int] = None
     write_capacity_units: Optional[int] = None
+    ssl_verify: Optional[bool] = None
     region_name: str
     user_table_name: str = "LiteLLM_UserTable"
     key_table_name: str = "LiteLLM_VerificationToken"
diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py
index a5c4b8371f..73e2005497 100644
--- a/litellm/proxy/db/dynamo_db.py
+++ b/litellm/proxy/db/dynamo_db.py
@@ -71,9 +71,15 @@ class DynamoDBWrapper(CustomDB):
         from aiodynamo.models import ReturnValues
         from aiodynamo.http.aiohttp import AIOHTTP
         from aiohttp import ClientSession
+        import aiohttp
 
         verbose_proxy_logger.debug("DynamoDB Wrapper - Attempting to connect")
-        async with ClientSession() as session:
+        # before making ClientSession check if ssl_verify=False
+        if self.database_arguments.ssl_verify == False:
+            client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
+        else:
+            client_session = ClientSession()
+        async with client_session as session:
             client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
             ## User
             try:
@@ -145,8 +151,13 @@ class DynamoDBWrapper(CustomDB):
         from aiodynamo.models import ReturnValues
         from aiodynamo.http.aiohttp import AIOHTTP
         from aiohttp import ClientSession
+        import aiohttp
 
-        async with ClientSession() as session:
+        if self.database_arguments.ssl_verify == False:
+            client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
+        else:
+            client_session = ClientSession()
+        async with client_session as session:
             client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
             table = None
             if table_name == "user":
@@ -178,8 +189,14 @@ class DynamoDBWrapper(CustomDB):
         from aiodynamo.models import ReturnValues
         from aiodynamo.http.aiohttp import AIOHTTP
         from aiohttp import ClientSession
+        import aiohttp
 
-        async with ClientSession() as session:
+        if self.database_arguments.ssl_verify == False:
+            client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
+        else:
+            client_session = ClientSession()
+
+        async with client_session as session:
             client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
             table = None
             key_name = None
@@ -232,8 +249,13 @@ class DynamoDBWrapper(CustomDB):
         from aiodynamo.models import ReturnValues
         from aiodynamo.http.aiohttp import AIOHTTP
         from aiohttp import ClientSession
+        import aiohttp
 
-        async with ClientSession() as session:
+        if self.database_arguments.ssl_verify == False:
+            client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False))
+        else:
+            client_session = ClientSession()
+        async with client_session as session:
             client = Client(AIOHTTP(session), Credentials.auto(), self.region_name)
             table = None
             key_name = None
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 1216b68dca..c464deaebd 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -55,7 +55,8 @@ general_settings:
   database_type: "dynamo_db" 
   database_args: { # 👈  all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190
     "billing_mode": "PAY_PER_REQUEST", 
-    "region_name": "us-west-2" 
+    "region_name": "us-west-2",
+    "ssl_verify": False
   }
 
 

From 5b201bbb4df64779ddcb1d987f80ba3e1cc161b0 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 16:20:29 -0800
Subject: [PATCH 02/23] (test) add new key gen tests

---
 litellm/tests/test_key_generate_dynamodb.py | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 litellm/tests/test_key_generate_dynamodb.py

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
new file mode 100644
index 0000000000..05b9daadf2
--- /dev/null
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -0,0 +1,9 @@
+# Test the following scenarios:
+# 1. Generate a Key, and use it to make a call
+# 2. Make a call with invalid key, expect it to fail
+# 3. Make a call to a key with invalid model - expect to fail
+# 4. Make a call to a key with valid model - expect to pass
+# 5. Make a call with expired key - expect to fail
+# 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

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 03/23] (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 05b9daadf2..609f38b8ff 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

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 04/23] (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

From 466fff4aaf411eee1e4d3dc3ee233bbc5eed8483 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 17:35:39 -0800
Subject: [PATCH 05/23] (fix) when budget, spend is None comparison

---
 litellm/proxy/proxy_server.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 353c3157db..1d0245c51c 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -348,10 +348,11 @@ async def user_api_key_auth(
                 if valid_token.spend is not None and valid_token.user_id is not None:
                     user_max_budget = user_id_information.max_budget
                     user_current_spend = user_id_information.spend
-                    if user_current_spend > user_max_budget:
-                        raise Exception(
-                            f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
-                        )
+                    if user_max_budget is not None and user_current_spend is not None:
+                        if user_current_spend > user_max_budget:
+                            raise Exception(
+                                f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
+                            )
 
             # Check 3. If token is expired
             if valid_token.expires is not None:

From adc921f0be7e762d210f669cfeaa56290d8611b3 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 17:36:33 -0800
Subject: [PATCH 06/23] (test) /key/gen

---
 litellm/tests/test_key_generate_dynamodb.py | 45 +++++++++++++++------
 1 file changed, 32 insertions(+), 13 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 4ba73eeb95..3b8fb86899 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -32,15 +32,14 @@ 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()
+db_args = {
+    "ssl_verify": False,
+    "billing_mode": "PAY_PER_REQUEST",
+    "region_name": "us-west-2",
+}
 custom_db_client = DBClient(
     custom_db_type="dynamo_db",
-    custom_db_args=db_args_dict,
+    custom_db_args=db_args,
 )
 
 request_data = {
@@ -54,6 +53,7 @@ request_data = {
 def test_generate_and_call_with_valid_key():
     # 1. Generate a Key, and use it to make a call
     setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
     try:
 
         async def test():
@@ -63,6 +63,31 @@ def test_generate_and_call_with_valid_key():
 
             generated_key = key.key
             bearer_token = "Bearer " + generated_key
+            from starlette.datastructures import URL
+
+            request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
+
+            # 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
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    try:
+
+        async def test():
+            generated_key = "bad-key"
+            bearer_token = "Bearer " + generated_key
 
             request = Request(scope={"type": "http"})
 
@@ -77,12 +102,6 @@ def test_generate_and_call_with_valid_key():
         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_model():
 #     # 3. Make a call to a key with an invalid model - expect to fail
 #     key = new_user(ValidNewUserRequest())

From 2c79b8fc1b1165871f8403d2c6c3938c32391d6d Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 17:40:32 -0800
Subject: [PATCH 07/23] (test) invalid key auth

---
 litellm/tests/test_key_generate_dynamodb.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 3b8fb86899..ea4254a221 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -31,6 +31,7 @@ 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
+from starlette.datastructures import URL
 
 db_args = {
     "ssl_verify": False,
@@ -63,7 +64,6 @@ def test_generate_and_call_with_valid_key():
 
             generated_key = key.key
             bearer_token = "Bearer " + generated_key
-            from starlette.datastructures import URL
 
             request = Request(scope={"type": "http"})
             request._url = URL(url="/chat/completions")
@@ -90,16 +90,19 @@ def test_call_with_invalid_key():
             bearer_token = "Bearer " + generated_key
 
             request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
 
             # 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)
+            pytest.fail(f"This should have failed!")
 
         asyncio.run(test())
     # result = user_auth(ValidRequest(key))
     # assert result is True
     except Exception as e:
-        pytest.fail(f"An exception occurred - {str(e)}")
+        print("Got Exception", e)
+        print(str(e))
+        pass
 
 
 # def test_call_with_invalid_model():

From 23ae967326e1645d835497b003b2578fa3901f74 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 17:41:34 -0800
Subject: [PATCH 08/23] (ci/cd) add aiodynamo to testing

---
 .circleci/config.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 26d714f9d3..307110a2d2 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -40,6 +40,7 @@ jobs:
             pip install "httpx==0.24.1"
             pip install "gunicorn==21.2.0"
             pip install "anyio==3.7.1"
+            pip install "aiodynamo==23.10.1"
             pip install "asyncio==3.4.3"
             pip install "PyGithub==1.59.1"
       - save_cache:

From d8a11b7ee60b6c562f7860b850f60a34c9a761e7 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 17:42:47 -0800
Subject: [PATCH 09/23] (test) cleanup

---
 litellm/tests/test_key_generate_dynamodb.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index ea4254a221..b2463bf97f 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -73,8 +73,6 @@ def test_generate_and_call_with_valid_key():
             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)}")
 
@@ -97,8 +95,6 @@ def test_call_with_invalid_key():
             pytest.fail(f"This should have failed!")
 
         asyncio.run(test())
-    # result = user_auth(ValidRequest(key))
-    # assert result is True
     except Exception as e:
         print("Got Exception", e)
         print(str(e))

From 121266b05b05c95a7e94018a10d70b07ef3688ee Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 19:35:34 -0800
Subject: [PATCH 10/23] (test) /key/gen fails when crossing budget

---
 litellm/tests/test_key_generate_dynamodb.py | 71 +++++++++++++++++++--
 1 file changed, 64 insertions(+), 7 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index b2463bf97f..e4469d29f6 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -27,7 +27,7 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 import pytest, logging, asyncio
 import litellm, asyncio
-from litellm.proxy.proxy_server import new_user, user_api_key_auth
+from litellm.proxy.proxy_server import new_user, user_api_key_auth, user_update
 
 from litellm.proxy._types import NewUserRequest, DynamoDBArgs
 from litellm.proxy.utils import DBClient
@@ -92,7 +92,7 @@ def test_call_with_invalid_key():
 
             # use generated key to auth in
             result = await user_api_key_auth(request=request, api_key=bearer_token)
-            pytest.fail(f"This should have failed!")
+            pytest.fail(f"This should have failed!. IT's an invalid key")
 
         asyncio.run(test())
     except Exception as e:
@@ -129,11 +129,68 @@ def test_call_with_invalid_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
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    try:
+
+        async def test():
+            request = NewUserRequest(max_budget=0.00001)
+            key = await new_user(request)
+            print(key)
+
+            generated_key = key.key
+            user_id = key.user_id
+            bearer_token = "Bearer " + generated_key
+
+            request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
+
+            # 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)
+
+            # update spend using track_cost callback, make 2nd request, it should fail
+            from litellm.proxy.proxy_server import track_cost_callback
+            from litellm import ModelResponse, Choices, Message, Usage
+
+            resp = ModelResponse(
+                id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
+                choices=[
+                    Choices(
+                        finish_reason=None,
+                        index=0,
+                        message=Message(
+                            content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
+                            role="assistant",
+                        ),
+                    )
+                ],
+                model="gpt-35-turbo",  # azure always has model written like this
+                usage=Usage(prompt_tokens=21, completion_tokens=17, total_tokens=38),
+            )
+            await track_cost_callback(
+                kwargs={
+                    "stream": False,
+                    "litellm_params": {
+                        "metadata": {
+                            "user_api_key": generated_key,
+                            "user_api_key_user_id": user_id,
+                        }
+                    },
+                },
+                completion_response=resp,
+            )
+
+            # 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:
+        assert "ExceededBudget:" in str(e)
+        print("Got Exception", e)
 
 
 # def test_call_with_key_over_budget():

From b06ebb886c9819c16d8116dbfd6ff9d6267fb174 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 19:53:43 -0800
Subject: [PATCH 11/23] (ci/cd) undo change from PR 1464

---
 litellm/proxy/proxy_server.py | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 1d0245c51c..b383317b3b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1331,11 +1331,7 @@ async def startup_event():
         verbose_proxy_logger.debug(f"custom_db_client connecting - {custom_db_client}")
         await custom_db_client.connect()
 
-    if prisma_client is not None:
-        if master_key is None:
-            raise ValueError(
-                "Using Proxy Auth, but Master Key not set, please set `LITELLM_MASTER_KEY` in your environment or `master_key` in your config.yaml"
-            )
+    if prisma_client is not None and master_key is not None:
         # add master key to db
         await generate_key_helper_fn(
             duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
@@ -1343,11 +1339,7 @@ async def startup_event():
     verbose_proxy_logger.debug(
         f"custom_db_client client - Inserting master key {custom_db_client}. Master_key: {master_key}"
     )
-    if custom_db_client is not None:
-        if master_key is None:
-            raise ValueError(
-                "Using Proxy Auth, but Master Key not set, please set `LITELLM_MASTER_KEY` in your environment or `master_key` in your config.yaml"
-            )
+    if custom_db_client is not None and master_key is not None:
         # add master key to db
         await generate_key_helper_fn(
             duration=None, models=[], aliases={}, config={}, spend=0, token=master_key

From a9e86cfa690c4ca6a927ac40a9450c01bb136a51 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 19:54:38 -0800
Subject: [PATCH 12/23] (test) all tests should pass now

---
 litellm/tests/test_key_generate_dynamodb.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index e4469d29f6..96da4e9df2 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -168,7 +168,7 @@ def test_call_with_key_under_budget():
                     )
                 ],
                 model="gpt-35-turbo",  # azure always has model written like this
-                usage=Usage(prompt_tokens=21, completion_tokens=17, total_tokens=38),
+                usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
             )
             await track_cost_callback(
                 kwargs={
@@ -186,10 +186,11 @@ def test_call_with_key_under_budget():
             # 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)
+            pytest.fail(f"This should have failed!. They key crossed it's budget")
 
         asyncio.run(test())
     except Exception as e:
-        assert "ExceededBudget:" in str(e)
+        pass
         print("Got Exception", e)
 
 

From 4df3bfac12417e5589b9e944f4e7b3927b3d652e Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 19:59:18 -0800
Subject: [PATCH 13/23] (test) assert error raised is ExceededBudget

---
 litellm/tests/test_key_generate_dynamodb.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 96da4e9df2..559cc9ff52 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -190,8 +190,9 @@ def test_call_with_key_under_budget():
 
         asyncio.run(test())
     except Exception as e:
-        pass
-        print("Got Exception", e)
+        error_detail = e.detail
+        assert "Authentication Error, ExceededBudget:" in error_detail
+        print(vars(e))
 
 
 # def test_call_with_key_over_budget():

From 39d5af5c85661a54f2087f1746bb4c4d9d0f1be4 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 20:01:04 -0800
Subject: [PATCH 14/23] (test) assert correct exceptions

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

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 559cc9ff52..5363bcd09a 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -97,7 +97,8 @@ def test_call_with_invalid_key():
         asyncio.run(test())
     except Exception as e:
         print("Got Exception", e)
-        print(str(e))
+        print(e.detail)
+        assert "Authentication Error" in e.detail
         pass
 
 

From cc51651b4921817f29e03a5d811009d314e0a258 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 20:43:20 -0800
Subject: [PATCH 15/23] (test) /key/generate - valid, invalid models

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

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index 5363bcd09a..ff2e9940b4 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -33,6 +33,7 @@ from litellm.proxy._types import NewUserRequest, DynamoDBArgs
 from litellm.proxy.utils import DBClient
 from starlette.datastructures import URL
 
+
 db_args = {
     "ssl_verify": False,
     "billing_mode": "PAY_PER_REQUEST",
@@ -87,7 +88,7 @@ def test_call_with_invalid_key():
             generated_key = "bad-key"
             bearer_token = "Bearer " + generated_key
 
-            request = Request(scope={"type": "http"})
+            request = Request(scope={"type": "http"}, receive=None)
             request._url = URL(url="/chat/completions")
 
             # use generated key to auth in
@@ -102,18 +103,70 @@ def test_call_with_invalid_key():
         pass
 
 
-# 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
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    try:
+
+        async def test():
+            request = NewUserRequest(models=["mistral"])
+            key = await new_user(request)
+            print(key)
+
+            generated_key = key.key
+            bearer_token = "Bearer " + generated_key
+
+            request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
+
+            async def return_body():
+                return b'{"model": "gemini-pro-vision"}'
+
+            request.body = return_body
+
+            # use generated key to auth in
+            result = await user_api_key_auth(request=request, api_key=bearer_token)
+            pytest.fail(f"This should have failed!. IT's an invalid model")
+
+        asyncio.run(test())
+    except Exception as e:
+        assert (
+            e.detail
+            == "Authentication Error, API Key not allowed to access model. This token can only access models=['mistral']. Tried to access gemini-pro-vision"
+        )
+        pass
 
 
-# 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
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    try:
+
+        async def test():
+            request = NewUserRequest(models=["mistral"])
+            key = await new_user(request)
+            print(key)
+
+            generated_key = key.key
+            bearer_token = "Bearer " + generated_key
+
+            request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
+
+            async def return_body():
+                return b'{"model": "mistral"}'
+
+            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)}")
 
 
 # def test_call_with_expired_key():

From cca852ce429291d8d77ca794a1a0e93acb1e867a Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 16:14:35 -0800
Subject: [PATCH 16/23] =?UTF-8?q?bump:=20version=201.17.10=20=E2=86=92=201?=
 =?UTF-8?q?.17.11?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pyproject.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ab56d67e4c..34fb6b776d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "1.17.10"
+version = "1.17.11"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"
@@ -61,7 +61,7 @@ requires = ["poetry-core", "wheel"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "1.17.10"
+version = "1.17.11"
 version_files = [
     "pyproject.toml:^version"
 ]

From 40364daac778535514963892f1b49a535f723e65 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 16 Jan 2024 16:23:29 -0800
Subject: [PATCH 17/23] fix(vertex_ai.py): raise exception if vertex ai missing
 required dependendencies

---
 litellm/llms/vertex_ai.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py
index 77c20d30a5..57cf6c2a0e 100644
--- a/litellm/llms/vertex_ai.py
+++ b/litellm/llms/vertex_ai.py
@@ -216,6 +216,14 @@ def completion(
             status_code=400,
             message="vertexai import failed please run `pip install google-cloud-aiplatform`",
         )
+
+    if not (
+        hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")
+    ):
+        raise VertexAIError(
+            status_code=400,
+            message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
+        )
     try:
         from vertexai.preview.language_models import (
             ChatModel,

From f8351717f15b1b384610b04e8f2162930290ee9d Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 16:27:21 -0800
Subject: [PATCH 18/23] (ci/cd) use fixture

---
 litellm/tests/test_proxy_startup.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/litellm/tests/test_proxy_startup.py b/litellm/tests/test_proxy_startup.py
index c033505bbd..33e317ccf8 100644
--- a/litellm/tests/test_proxy_startup.py
+++ b/litellm/tests/test_proxy_startup.py
@@ -24,7 +24,14 @@ from litellm.proxy.proxy_server import (
 )
 
 
-def test_proxy_gunicorn_startup_direct_config():
+@pytest.fixture
+def cleanup_router_config():
+    from litellm.proxy.proxy_server import cleanup_router_config_variables
+
+    cleanup_router_config_variables()
+
+
+def test_proxy_gunicorn_startup_direct_config(cleanup_router_config):
     """
     gunicorn startup requires the config to be passed in via environment variables
 
@@ -33,9 +40,6 @@ def test_proxy_gunicorn_startup_direct_config():
     Test both approaches
     """
     try:
-        from litellm.proxy.proxy_server import cleanup_router_config_variables
-
-        cleanup_router_config_variables()
         filepath = os.path.dirname(os.path.abspath(__file__))
         # test with worker_config = config yaml
         config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
@@ -49,11 +53,8 @@ def test_proxy_gunicorn_startup_direct_config():
             pytest.fail(f"An exception occurred - {str(e)}")
 
 
-def test_proxy_gunicorn_startup_config_dict():
+def test_proxy_gunicorn_startup_config_dict(cleanup_router_config):
     try:
-        from litellm.proxy.proxy_server import cleanup_router_config_variables
-
-        cleanup_router_config_variables()
         filepath = os.path.dirname(os.path.abspath(__file__))
         # test with worker_config = config yaml
         config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"

From d1897ca1562c390d659402d0cd7dac2f066b5f77 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 16:53:04 -0800
Subject: [PATCH 19/23] (ci/cd) undo change from PR 1464

---
 litellm/tests/test_proxy_startup.py | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/litellm/tests/test_proxy_startup.py b/litellm/tests/test_proxy_startup.py
index 33e317ccf8..650e2f8a7a 100644
--- a/litellm/tests/test_proxy_startup.py
+++ b/litellm/tests/test_proxy_startup.py
@@ -24,14 +24,7 @@ from litellm.proxy.proxy_server import (
 )
 
 
-@pytest.fixture
-def cleanup_router_config():
-    from litellm.proxy.proxy_server import cleanup_router_config_variables
-
-    cleanup_router_config_variables()
-
-
-def test_proxy_gunicorn_startup_direct_config(cleanup_router_config):
+def test_proxy_gunicorn_startup_direct_config():
     """
     gunicorn startup requires the config to be passed in via environment variables
 
@@ -53,7 +46,7 @@ def test_proxy_gunicorn_startup_direct_config(cleanup_router_config):
             pytest.fail(f"An exception occurred - {str(e)}")
 
 
-def test_proxy_gunicorn_startup_config_dict(cleanup_router_config):
+def test_proxy_gunicorn_startup_config_dict():
     try:
         filepath = os.path.dirname(os.path.abspath(__file__))
         # test with worker_config = config yaml

From 179860863a7a0dbd4670619151460c1b2930ce6b Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 19:42:34 -0800
Subject: [PATCH 20/23] (v0) fix

---
 litellm/proxy/proxy_server.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index b383317b3b..0503d2aad8 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -507,7 +507,9 @@ async def track_cost_callback(
             if user_api_key and (
                 prisma_client is not None or custom_db_client is not None
             ):
-                await update_database(token=user_api_key, response_cost=response_cost)
+                await update_database(
+                    token=user_api_key, response_cost=response_cost, user_id=user_id
+                )
         elif kwargs["stream"] == False:  # for non streaming responses
             response_cost = litellm.completion_cost(
                 completion_response=completion_response

From 8d75968650cddbe4e862ffb41517de75d1fa9615 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 20:07:51 -0800
Subject: [PATCH 21/23] =?UTF-8?q?bump:=20version=201.17.11=20=E2=86=92=201?=
 =?UTF-8?q?.17.12?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pyproject.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 34fb6b776d..ba2722d733 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "1.17.11"
+version = "1.17.12"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"
@@ -61,7 +61,7 @@ requires = ["poetry-core", "wheel"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "1.17.11"
+version = "1.17.12"
 version_files = [
     "pyproject.toml:^version"
 ]

From bae9cd74396a819629ed1d38b52b4c9e6c5b6348 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 20:09:43 -0800
Subject: [PATCH 22/23] (docs) proxy /key/gen

---
 docs/my-website/sidebars.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js
index facda0e98d..0735870e96 100644
--- a/docs/my-website/sidebars.js
+++ b/docs/my-website/sidebars.js
@@ -107,8 +107,8 @@ const sidebars = {
         "proxy/user_keys",
         "proxy/load_balancing", 
         "proxy/virtual_keys",
-        "proxy/ui",
         "proxy/users",
+        "proxy/ui",
         "proxy/model_management",
         "proxy/reliability",
         "proxy/caching",

From 9c214d67b0e071753e37992f5c4124539fe3c650 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 16 Jan 2024 20:56:01 -0800
Subject: [PATCH 23/23] (test) /key/gen and budgets

---
 litellm/tests/test_key_generate_dynamodb.py | 99 +++++++++++++++------
 1 file changed, 74 insertions(+), 25 deletions(-)

diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py
index ff2e9940b4..207c8835a7 100644
--- a/litellm/tests/test_key_generate_dynamodb.py
+++ b/litellm/tests/test_key_generate_dynamodb.py
@@ -3,10 +3,8 @@
 # 2. Make a call with invalid key, expect it to fail
 # 3. Make a call to a key with invalid model - expect to fail
 # 4. Make a call to a key with valid model - expect to pass
-# 5. Make a call with expired key - expect to fail
-# 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
+# 5. Make a call with key over budget, expect to fail
+# 6. Make a streaming chat/completions call with key over budget, expect to fail
 
 
 # function to call to generate key - async def new_user(data: NewUserRequest):
@@ -169,22 +167,8 @@ def test_call_with_valid_model():
         pytest.fail(f"An exception occurred - {str(e)}")
 
 
-# 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
+def test_call_with_key_over_budget():
+    # 5. Make a call with a key over budget, expect to fail
     setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
     setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
     try:
@@ -249,8 +233,73 @@ def test_call_with_key_under_budget():
         print(vars(e))
 
 
-# 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_stream():
+    # 6. Make a call with a key over budget, expect to fail
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    from litellm._logging import verbose_proxy_logger
+    import logging
+
+    litellm.set_verbose = True
+    verbose_proxy_logger.setLevel(logging.DEBUG)
+    try:
+
+        async def test():
+            request = NewUserRequest(max_budget=0.00001)
+            key = await new_user(request)
+            print(key)
+
+            generated_key = key.key
+            user_id = key.user_id
+            bearer_token = "Bearer " + generated_key
+
+            request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
+
+            # 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)
+
+            # update spend using track_cost callback, make 2nd request, it should fail
+            from litellm.proxy.proxy_server import track_cost_callback
+            from litellm import ModelResponse, Choices, Message, Usage
+
+            resp = ModelResponse(
+                id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
+                choices=[
+                    Choices(
+                        finish_reason=None,
+                        index=0,
+                        message=Message(
+                            content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
+                            role="assistant",
+                        ),
+                    )
+                ],
+                model="gpt-35-turbo",  # azure always has model written like this
+                usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
+            )
+            await track_cost_callback(
+                kwargs={
+                    "stream": True,
+                    "complete_streaming_response": resp,
+                    "litellm_params": {
+                        "metadata": {
+                            "user_api_key": generated_key,
+                            "user_api_key_user_id": user_id,
+                        }
+                    },
+                },
+                completion_response=ModelResponse(),
+            )
+
+            # 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)
+            pytest.fail(f"This should have failed!. They key crossed it's budget")
+
+        asyncio.run(test())
+    except Exception as e:
+        error_detail = e.detail
+        assert "Authentication Error, ExceededBudget:" in error_detail
+        print(vars(e))