test type of budget errors

litellm/tests/test_key_generate_prisma.py

@@ -95,6 +95,8 @@ from litellm.proxy._types import (
     NewCustomerRequest,
     NewTeamRequest,
     NewUserRequest,
+    ProxyErrorTypes,
+    ProxyException,
     UpdateKeyRequest,
     UpdateTeamRequest,
     UserAPIKeyAuth,
@@ -491,6 +493,8 @@ def test_call_with_user_over_budget(prisma_client):
     except Exception as e:
         error_detail = e.message
         assert "Budget has been exceeded" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -597,6 +601,8 @@ def test_call_with_end_user_over_budget(prisma_client):
     except Exception as e:
         error_detail = e.message
         assert "Budget has been exceeded! Current" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -686,6 +692,8 @@ def test_call_with_proxy_over_budget(prisma_client):
         else:
             error_detail = traceback.format_exc()
         assert "Budget has been exceeded" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -765,6 +773,8 @@ def test_call_with_user_over_budget_stream(prisma_client):
     except Exception as e:
         error_detail = e.message
         assert "Budget has been exceeded" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -1401,6 +1411,8 @@ def test_call_with_key_over_budget(prisma_client):
         else:
             error_detail = str(e)
         assert "Budget has been exceeded" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -1514,6 +1526,8 @@ def test_call_with_key_over_budget_no_cache(prisma_client):
         else:
             error_detail = str(e)
         assert "Budget has been exceeded" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -1629,6 +1643,8 @@ def test_call_with_key_over_model_budget(prisma_client):
         traceback.print_exc()
         error_detail = e.message
         assert "Budget has been exceeded!" in error_detail
+        assert isinstance(e, ProxyException)
+        assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
 
 
@@ -1795,6 +1811,7 @@ async def test_call_with_key_over_budget_stream(prisma_client):
         print("Got Exception", e)
         error_detail = e.message
         assert "Budget has been exceeded" in error_detail
+
         print(vars(e))