From 52a3f64846a781e33ba5893f1e5b65614adefe00 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 9 Jun 2026 14:00:02 -0700 Subject: [PATCH] fix' : --- tests/otel_tests/test_e2e_budgeting.py | 1 + .../test_budget_endpoints.py | 23 ++++++++----------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/tests/otel_tests/test_e2e_budgeting.py b/tests/otel_tests/test_e2e_budgeting.py index 0b40a081cb1..ab2874190c8 100644 --- a/tests/otel_tests/test_e2e_budgeting.py +++ b/tests/otel_tests/test_e2e_budgeting.py @@ -122,6 +122,7 @@ async def generate_key_with_budget_id(session, budget_id: str): headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} data = {"budget_id": budget_id} async with session.post(url, headers=headers, json=data) as response: + assert response.status == 200, await response.text() return await response.json() diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index d5edb08877c..cfce2c14e53 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -45,7 +45,6 @@ def admin_client_and_mocks(monkeypatch): yield client, mock_prisma, mock_table app.dependency_overrides.clear() - monkeypatch.setattr(ps, "prisma_client", ps.prisma_client) @pytest.fixture @@ -75,9 +74,7 @@ def client_and_mocks(monkeypatch): yield client, mock_prisma, mock_table - # teardown app.dependency_overrides.clear() - monkeypatch.setattr(ps, "prisma_client", ps.prisma_client) @pytest.mark.asyncio @@ -302,17 +299,15 @@ async def test_new_budget_invalid_model_max_budget(client_and_mocks, monkeypatch async def test_info_budget_success(admin_client_and_mocks): client, _, mock_table = admin_client_and_mocks - mock_row = types.SimpleNamespace( - budget_id="budget-info-1", - max_budget=10.0, - budget_duration="30d", - dict=lambda: { - "budget_id": "budget-info-1", - "max_budget": 10.0, - "budget_duration": "30d", - }, + mock_table.find_many = AsyncMock( + return_value=[ + { + "budget_id": "budget-info-1", + "max_budget": 10.0, + "budget_duration": "30d", + } + ] ) - mock_table.find_many = AsyncMock(return_value=[mock_row]) resp = client.post("/budget/info", json={"budgets": ["budget-info-1"]}) assert resp.status_code == 200, resp.text @@ -429,6 +424,6 @@ async def test_delete_budget_rejects_non_admin(client_and_mocks): resp = client.post("/budget/delete", json={"id": "budget-delete-1"}) assert resp.status_code == 400, resp.text assert CommonProxyErrors.not_allowed_access.value in str(resp.json()["detail"]) - mock_table.delete.assert_not_awaited() + mock_table.delete.assert_not_called() finally: app.dependency_overrides.clear()