From 5fa85cc8585596a37cd8ae0a769b88ec6c431d80 Mon Sep 17 00:00:00 2001
From: yucheng <yucheng@berri.ai>
Date: Thu, 5 Nov 2026 01:35:08 -0800
Subject: [PATCH] test(proxy): persist the failure's spend row through the real
 writer instead of a mocked one

---
 .../hooks/test_proxy_track_cost_callback.py   | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
index 31cd4dc41c0..e4f9a9df207 100644
--- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
+++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
@@ -39,6 +39,7 @@ from litellm.proxy.utils import (
     REJECTED_BEFORE_REQUEST_SETUP,
     REJECTED_BEFORE_REQUEST_SETUP_KEY,
     BilledFailedAttempt,
+    PrismaClient,
     ProxyUpdateSpend,
 )
 from litellm.types.router import Deployment, LiteLLM_Params
@@ -4058,6 +4059,16 @@ async def test_a_poll_of_an_already_charged_job_still_charges_its_own_guardrails
     assert charge_spend.calls == ((0.43, 0.43, reservation), *(() if second_call is None else (second_call,)))
 
 
+class _SpendLogQueue:
+    """The database client's spend-log queue, where the writer puts a failure's spend row."""
+
+    get_request_status = PrismaClient.get_request_status
+
+    def __init__(self) -> None:
+        self.spend_log_transactions: Final[list[Mapping[str, object]]] = []
+        self._spend_log_transactions_lock: Final = asyncio.Lock()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"])
 async def test_async_post_call_failure_hook_scrubs_the_prompt_out_of_the_persisted_error(
@@ -4065,9 +4076,11 @@ async def test_async_post_call_failure_hook_scrubs_the_prompt_out_of_the_persist
 ) -> None:
     secret: Final = "my ssn is 123-45-6789"
     monkeypatch.setattr(litellm, "turn_off_message_logging", redact)
-    spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
+    spend_logs: Final = _SpendLogQueue()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", spend_logs)
+    monkeypatch.setattr(PrismaClient, "spend_log_queue_bytes", 0)
 
-    await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook(
+    await ProxyDBLogger(spend_writer=DBSpendUpdateWriter).async_post_call_failure_hook(
         request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": secret}], "metadata": {}},
         original_exception=litellm.BadRequestError(
             message=f"Unsupported content: {secret}", model="gpt-4o", llm_provider="openai"
@@ -4076,9 +4089,8 @@ async def test_async_post_call_failure_hook_scrubs_the_prompt_out_of_the_persist
         traceback_str=f"Traceback: {secret}",
     )
 
-    persisted: Final = spend_writer.update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"][
-        "error_information"
-    ]
+    (row,) = spend_logs.spend_log_transactions
+    persisted: Final = json.loads(str(row["metadata"]))["error_information"]
     assert (secret in json.dumps(persisted)) is not redact
     assert "Unsupported content: " in persisted["error_message"]
     assert persisted["error_class"] == "BadRequestError"