From 56dd4e06accc47a83b031c241f532c3b9a07ce1b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 13:29:59 -0700 Subject: [PATCH] test(proxy): satisfy the test-quality gate for the window spend writer tests --- .../test_window_spend_update_queue.py | 12 +-- .../db/test_budget_window_spend_writer.py | 43 ++++----- .../proxy/db/test_db_spend_update_writer.py | 88 +++++++------------ 3 files changed, 58 insertions(+), 85 deletions(-) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py index 25c685f6f29..b1ecda57afa 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_window_spend_update_queue.py @@ -1,10 +1,6 @@ import json -import os -import sys from datetime import datetime, timedelta, timezone -sys.path.insert(0, os.path.abspath("../../../../..")) - import pytest from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( @@ -207,9 +203,7 @@ def test_aggregation_survives_the_redis_json_round_trip(): [(_txn("k1", WINDOW_A, 1.0),), (_txn("k1", WINDOW_B, 2.0),)] ) - reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - [json.loads(json.dumps(aggregated))] - ) + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) assert reloaded == aggregated @@ -269,9 +263,7 @@ def test_request_ids_survive_the_redis_json_round_trip(): [(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)] ) - reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - [json.loads(json.dumps(aggregated))] - ) + reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))]) assert reloaded[0]["request_ids"] == ("req-1",) assert reloaded[0]["spend"] == 1.0 diff --git a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py index 4d7d86130fb..a849317c930 100644 --- a/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py +++ b/tests/test_litellm/proxy/db/test_budget_window_spend_writer.py @@ -1,12 +1,8 @@ import math -import os -import sys from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone from typing import Any -sys.path.insert(0, os.path.abspath("../../../..")) - import pytest from litellm.proxy.db.budget_window_spend_writer import ( @@ -134,7 +130,9 @@ def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | No "window_start": "2026-08-01T00:00:00.000000", "spend": spend, "request_ids": request_ids, - "started_at": None if started_at is None else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), + "started_at": None + if started_at is None + else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"), } @@ -171,7 +169,7 @@ async def test_missing_row_is_seeded_from_spend_logs_once(): assert aggregate.calls[0]["entity_id"] == "k1" assert aggregate.calls[0]["window_start"] == WINDOW_A - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[ENTITY_TYPE] == "key" assert params[ENTITY_ID] == "k1" assert params[WINDOW_DURATION] == "30d" @@ -194,7 +192,7 @@ async def test_existing_row_is_never_reseeded(): ) assert aggregate.calls == [] - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(1.0) assert params[INCREMENT] == pytest.approx(1.0) @@ -233,7 +231,7 @@ async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded(): spend_logs_aggregate=aggregate, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(9.25) assert params[INCREMENT] == pytest.approx(0.25) @@ -249,7 +247,7 @@ async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),), ) - (query, _), = db.batcher.calls + ((query, _),) = db.batcher.calls normalized = " ".join(query.split()) assert ( 'spend = CASE WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start ' @@ -270,7 +268,7 @@ async def test_upsert_never_interpolates_values_into_the_sql(): spend_logs_aggregate=aggregate, ) - (query, params), = db.batcher.calls + ((query, params),) = db.batcher.calls assert "DROP TABLE" not in query assert params[ENTITY_ID] == "'; DROP TABLE x; --" @@ -293,7 +291,10 @@ async def test_upserts_are_ordered_by_primary_key_then_window_start(): spend_logs_aggregate=aggregate, ) - ordered = [(params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) for _, params in db.batcher.calls] + ordered = [ + (params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) + for _, params in db.batcher.calls + ] assert ordered == [ ("key", "k1", "7d", datetime(2026, 8, 1)), ("key", "k2", "30d", datetime(2026, 8, 1)), @@ -316,7 +317,7 @@ async def test_existing_row_lookup_sends_every_primary_key_as_array_params(): spend_logs_aggregate=aggregate, ) - (query, params), = db.query_raw_calls + ((query, params),) = db.query_raw_calls assert "unnest($1::text[], $2::text[], $3::text[])" in query assert params == (("key", "team"), ("k1", "t1"), ("30d", "7d")) @@ -354,7 +355,7 @@ async def test_unknown_entity_type_contributes_no_seed(): spend_logs_aggregate=no_such_column, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(1.0) @@ -371,7 +372,7 @@ async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing() spend_logs_aggregate=unavailable, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(1.0) @@ -389,7 +390,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o new_window_start=WINDOW_B, ) - (query, params), = db.execute_raw_calls + ((query, params),) = db.execute_raw_calls normalized = " ".join(query.split()) assert "SET window_start = ($4::timestamptz AT TIME ZONE 'UTC'), spend = 0" in normalized assert "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3" in normalized @@ -447,7 +448,7 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed spend_logs_aggregate=already_flushed, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.000141) @@ -463,7 +464,7 @@ async def test_new_row_still_covers_spend_that_predates_the_batch(): spend_logs_aggregate=spend_logs, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.500047) @@ -482,7 +483,7 @@ async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed() spend_logs_aggregate=spend_logs, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.500047) @@ -499,7 +500,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet(): spend_logs_aggregate=nothing_flushed, ) - (_, params), = db.batcher.calls + ((_, params),) = db.batcher.calls assert params[INSERT_SPEND] == pytest.approx(0.000141) @@ -523,7 +524,7 @@ async def test_seed_aggregate_sql_excludes_the_request_ids_only_within_the_batch ) assert total == pytest.approx(1.25) - (query, params), = db.query_raw_calls + ((query, params),) = db.query_raw_calls normalized = " ".join(query.split()) assert expected_column in normalized assert "NOT (request_id = ANY($3::text[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized @@ -557,7 +558,7 @@ async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_boun ) assert total == pytest.approx(1.25) - (query, params), = db.query_raw_calls + ((query, params),) = db.query_raw_calls assert "request_id" not in query assert params == ("e1", WINDOW_A) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 8165dbbe3d1..e07d326e8be 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -67,9 +67,7 @@ async def test_daily_spend_tracking_with_disabled_spend_logs(): assert db_writer.add_spend_log_transaction_to_daily_user_transaction.called # Verify the payload passed to add_spend_log_transaction_to_daily_user_transaction - call_args = ( - db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] - ) + call_args = db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1] assert "payload" in call_args assert call_args["payload"]["spend"] == 0.1 assert call_args["payload"]["model"] == "gpt-4" @@ -409,7 +407,7 @@ async def test_update_daily_spend_sorting(): # fields, but entity_id is sufficient to test sorting. daily_spend_transactions = { f"test_key_{i}": { - "user_id": f"user{60-i}", # user60 ... user11, reverse order + "user_id": f"user{60 - i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", "model": "gpt-4", @@ -988,9 +986,9 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i transaction_dict = call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): - assert ( - transaction["request_id"] == request_id - ), f"request_id should be {request_id} but got {transaction.get('request_id')}" + assert transaction["request_id"] == request_id, ( + f"request_id should be {request_id} but got {transaction.get('request_id')}" + ) @pytest.mark.asyncio @@ -1216,21 +1214,15 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_calls_common } writer.daily_agent_spend_update_queue.add_update = AsyncMock() - original_common_helper = ( - writer._common_add_spend_log_transaction_to_daily_transaction - ) - writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( - wraps=original_common_helper - ) + original_common_helper = writer._common_add_spend_log_transaction_to_daily_transaction + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock(wraps=original_common_helper) await writer.add_spend_log_transaction_to_daily_agent_transaction( payload=payload, prisma_client=mock_prisma, ) - assert ( - writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 - ) + assert writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 @pytest.mark.asyncio @@ -1385,6 +1377,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): Test that when batch upsert fails, the exception is properly re-raised after logging. This ensures that error handling continues to work correctly upstream. """ + def raise_connection_lost(): raise ValueError("Database connection lost") @@ -1565,9 +1558,7 @@ async def test_update_database_creates_single_task(): patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch( - "litellm.proxy.db.db_spend_update_writer.asyncio.create_task" - ) as mock_create_task, + patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task, ): await db_writer.update_database( token="test-token", @@ -1666,9 +1657,7 @@ async def test_daily_agent_receives_deepcopied_payload(): db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( - side_effect=capture_agent_payload - ) + db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(side_effect=capture_agent_payload) db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() @@ -1730,8 +1719,8 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer = AsyncMock() mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() # Return all-None tuple (no data to commit); the pipeline yields 6 slots - mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = ( - AsyncMock(return_value=(None, None, None, None, None, None, None)) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) ) db_writer.redis_update_buffer = mock_redis_update_buffer @@ -2159,9 +2148,7 @@ async def test_update_database_does_not_deepcopy_on_request_path(): db_writer._update_org_db = AsyncMock() db_writer._update_tag_db = AsyncMock() db_writer._update_agent_db = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock( - side_effect=capture_batch_payload - ) + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock(side_effect=capture_batch_payload) db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() @@ -2253,9 +2240,7 @@ async def test_spend_update_path_never_queries_user_cache_with_none_user_id(): db_writer = DBSpendUpdateWriter() strict_redis_backed_cache = MagicMock() - strict_redis_backed_cache.async_get_cache = AsyncMock( - side_effect=DataError("Invalid input of type: 'NoneType'") - ) + strict_redis_backed_cache.async_get_cache = AsyncMock(side_effect=DataError("Invalid input of type: 'NoneType'")) with ( patch.object(litellm, "max_budget", 0), @@ -2383,8 +2368,7 @@ async def test_daily_transaction_carries_compression_saved_tokens(): cache_write_cost = model_info.get("cache_creation_input_token_cost") or input_cost assert transaction["compression_savings_spend"] == pytest.approx(7600 * input_cost) assert transaction["prompt_caching_savings_spend"] == pytest.approx( - 40 * max(input_cost - cache_read_cost, 0.0) - - 15 * (cache_write_cost - input_cost) + 40 * max(input_cost - cache_read_cost, 0.0) - 15 * (cache_write_cost - input_cost) ) assert transaction["compression_savings_spend"] > 0 assert transaction["prompt_caching_savings_spend"] > 0 @@ -2472,9 +2456,7 @@ class _WindowSpendFakePrisma: def _window_spend_upserts(db): - return [ - params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query - ] + return [params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query] @pytest.mark.asyncio @@ -2492,9 +2474,7 @@ async def test_window_spend_queue_is_flushed_without_redis_buffer(): ) ) db = _WindowSpendFakeDB( - existing_rows=[ - {"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"} - ] + existing_rows=[{"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"}] ) await db_writer._commit_spend_updates_to_db_without_redis_buffer( @@ -2554,9 +2534,7 @@ async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_wi db_writer.redis_update_buffer = mock_redis_update_buffer db_writer.pod_lock_manager = AsyncMock() db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) - db = _WindowSpendFakeDB( - existing_rows=[{"entity_type": "team", "entity_id": "team-1", "window_duration": "7d"}] - ) + db = _WindowSpendFakeDB(existing_rows=[{"entity_type": "team", "entity_id": "team-1", "window_duration": "7d"}]) await db_writer._commit_spend_updates_to_db_with_redis( prisma_client=_WindowSpendFakePrisma(db), @@ -2627,9 +2605,12 @@ async def test_update_database_returns_the_spend_log_request_id(): db_writer._enqueue_tool_usage_transaction = AsyncMock() with ( - patch("litellm.proxy.proxy_server.disable_spend_logs", False), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ) ): request_id = await db_writer.update_database( token="test-token", @@ -2655,10 +2636,13 @@ async def test_update_database_returns_none_when_the_payload_cannot_be_built(): db_writer = DBSpendUpdateWriter() with ( - patch("litellm.proxy.proxy_server.disable_spend_logs", False), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch( + patch.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam + "litellm.proxy.proxy_server", + disable_spend_logs=False, + prisma_client=MagicMock(), + litellm_proxy_budget_name="test-budget", + ), + patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", side_effect=Exception("payload boom"), ), @@ -2979,9 +2963,7 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey "call_type, expects_flush", [("aresponses", True), ("responses", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( - call_type: str, expects_flush: bool -): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. @@ -3011,9 +2993,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls( pytest.param("", True, id="injected-before-a-deployment-was-chosen"), ], ) -async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected( - injected_deployment, attributed -): +async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected(injected_deployment, attributed): """Retries, same-group failover and cross-model-group fallbacks all reuse one metadata bucket and one litellm_call_id, so a marker written by the leg that injected is visible to every sibling and nothing request-scoped can tell them apart.