From ee8b1cfabcccc05bfca139114088cfb1408c4de2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 13 Nov 2025 16:26:02 -0800 Subject: [PATCH] test_call_with_end_user_over_budget --- tests/proxy_unit_tests/test_key_generate_prisma.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 06972c32a6..ca8014f76f 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -132,6 +132,11 @@ def prisma_client(): ### add connection pool + pool timeout args params = {"connection_limit": 100, "pool_timeout": 60} database_url = os.getenv("DATABASE_URL") + + # If DATABASE_URL is not set, use a default test database URL + if not database_url: + database_url = "postgresql://postgres:postgres@localhost:5432/circle_test" + modified_url = append_query_params(database_url, params) os.environ["DATABASE_URL"] = modified_url @@ -666,7 +671,8 @@ def test_call_with_end_user_over_budget(prisma_client): except Exception as e: print(f"raised error: {e}, traceback: {traceback.format_exc()}") error_detail = e.message - assert "Budget has been exceeded! Current" in error_detail + assert "ExceededBudget: End User=" in error_detail + assert "over budget" in error_detail assert isinstance(e, ProxyException) assert e.type == ProxyErrorTypes.budget_exceeded print(vars(e))