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