scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 175

test(proxy): persist the failure's spend row through the real writer instead of a mocked one

Browse source
This commit is contained in:
yucheng 2026-11-05 01:35:08 -08:00
parent 40a3624bc3
commit 5fa85cc858
1 changed files with 17 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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"