diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 5607d6c5b..36864138a 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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))