diff --git a/tests/e2e/budgets/test_budget_crud_e2e.py b/tests/e2e/budgets/test_budget_crud_e2e.py
index 3a2d1dd4fc8..fe4104328ad 100644
--- a/tests/e2e/budgets/test_budget_crud_e2e.py
+++ b/tests/e2e/budgets/test_budget_crud_e2e.py
@@ -29,10 +29,10 @@ def test_budget_crud_roundtrip(client: BudgetClient, resources: ResourceManager)
     # Attach the budget to a key and confirm the key reflects it.
     key = client.generate_key(extra_params={"budget_id": budget_id})
     resources.defer(lambda: client.delete_key(key))
-    info = client.key_info(key)
-    linked = info.get("litellm_budget_table") or {}
-    assert info.get("budget_id") == budget_id or linked.get("max_budget") == 12.5, (
-        f"key does not reflect attached budget: {info.get('budget_id')}, {linked}"
+    info = client.gateway.key_info(key)
+    linked = info.litellm_budget_table
+    assert info.budget_id == budget_id or (linked is not None and linked.max_budget == 12.5), (
+        f"key does not reflect attached budget: {info.budget_id}, {linked}"
     )
 
 
@@ -46,7 +46,7 @@ def test_budget_duration_schedules_reset_on_key(client: BudgetClient, resources:
     key = client.generate_key(max_budget=10.0, extra_params={"budget_duration": "30d"})
     resources.defer(lambda: client.delete_key(key))
 
-    reset_at = client.key_info(key).get("budget_reset_at")
+    reset_at = client.gateway.key_info(key).budget_reset_at
     assert reset_at, "budget_duration did not set budget_reset_at on the key"
 
     # budget_duration schedules a FUTURE reset. Don't assume now+30d exactly: the