test_call_with_end_user_over_budget

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))