From e125a98605eeb802d177e3ab01b36edd51e55c69 Mon Sep 17 00:00:00 2001
From: lei_lei <96427312+leilei3167@users.noreply.github.com>
Date: Sat, 29 Aug 2026 11:57:09 +0000
Subject: [PATCH] test: tighten spend-drain helper types

---
 .../proxy/proxy_server/test_lifecycle.py          |  7 ++++---
 .../prisma_and_spend/test_spend_functions.py      | 15 ++++++++++-----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
index 3ff99fb81b5..83d3d271c52 100644
--- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
+++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py
@@ -177,18 +177,19 @@ async def test_proxy_shutdown_drains_spend_logs_before_disconnecting(monkeypatch
     directly — is otherwise discarded on worker recycle with no log line.
     Ordering is the behavior, so assert drain then disconnect.
     """
-    calls: list = []  # mutable-ok: records call order, which is the assertion
+    calls: list[str] = []  # mutable-ok: records call order, which is the assertion
 
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
 
-    async def _record_drain():
+    async def _record_drain() -> None:
         calls.append("drain_spend")
 
     monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", _record_drain, raising=False)
 
-    async def _record_flush(client, accumulator):
+    async def _record_flush(client: object, accumulator: object) -> None:
+        del accumulator
         calls.append("flush_gateway")
         assert client is fake_prisma
 
diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py
index 20f651c2e35..81c892ae742 100644
--- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py
+++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py
@@ -351,8 +351,9 @@ async def test_drain_spend_logs_queue_flushes_rows_queued_while_draining(
 
     written: list[str] = []
 
-    async def _write(*args: Any, **kwargs: Any) -> None:
-        written.extend(row["request_id"] for row in kwargs["data"])
+    async def _write(*, data: list[dict[str, object]], skip_duplicates: bool) -> None:
+        del skip_duplicates
+        written.extend(row["request_id"] for row in data)
         if len(written) == 1:
             mock_prisma_client.spend_log_transactions.append(
                 make_spend_log_row(request_id="r2")
@@ -392,12 +393,13 @@ async def test_drain_spend_logs_queue_stops_monitor_and_keeps_its_popped_rows(
     written: list[str] = []
     write_calls = {"n": 0}
 
-    async def _write(*args: Any, **kwargs: Any) -> None:
+    async def _write(*, data: list[dict[str, object]], skip_duplicates: bool) -> None:
+        del skip_duplicates
         write_calls["n"] += 1
         if write_calls["n"] == 1:
             write_started.set()
             await asyncio.Event().wait()
-        written.extend(row["request_id"] for row in kwargs["data"])
+        written.extend(row["request_id"] for row in data)
 
     mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_write)
 
@@ -440,7 +442,10 @@ async def test_drain_spend_logs_queue_gives_up_after_max_passes(
     proxy_logging.failure_handler = AsyncMock()
     mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")]
 
-    async def _write_and_refill(*args: Any, **kwargs: Any) -> None:
+    async def _write_and_refill(
+        *, data: list[dict[str, object]], skip_duplicates: bool
+    ) -> None:
+        del data, skip_duplicates
         mock_prisma_client.spend_log_transactions.append(make_spend_log_row())
 
     mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(