scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

test: initialize websocket snapshot fixture without mutation

Browse source
This commit is contained in:
Anirudh Dharmarajan 2026-09-17 21:18:46 -06:00
parent 2e07c579e9
commit 13ea8be031
No known key found for this signature in database
GPG key ID: 1E74D51F530702D8
1 changed files with 6 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -2421,8 +2421,11 @@ def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store):
@pytest.mark.asyncio
@pytest.mark.parametrize("redact", [False, True], ids=["stored", "redacted"])
@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs")
async def test_websocket_snapshot_spend_serialization_preserves_transport(mock_should_store, redact):
@patch(
"litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs",
return_value=True,
)
async def test_websocket_snapshot_spend_serialization_preserves_transport(_mock_should_store, redact):
from starlette.websockets import WebSocket
from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot
@ -2432,7 +2435,6 @@ async def test_websocket_snapshot_spend_serialization_preserves_transport(mock_s
def __dict__(self):
raise AssertionError("Spend logging must not traverse the WebSocket app graph")
mock_should_store.return_value = True
send = AsyncMock()
websocket = WebSocket(
scope={"type": "websocket", "app": AppTraversalTrap()},
@ -2442,15 +2444,12 @@ async def test_websocket_snapshot_spend_serialization_preserves_transport(mock_s
await websocket.accept()
data = {
"model": "gpt-4o-realtime-preview",
"messages": [{"role": "user", "content": "private input"}],
"messages": [{"role": "user", "content": "<MASKED>"}],
"metadata": {"trace_id": "realtime-trace", "websocket": "user metadata"},
"websocket": websocket,
"proxy_server_request": {},
}
refresh_proxy_server_request_body_snapshot(data)
# Guardrails can replace messages after the initial snapshot.
data["messages"] = [{"role": "user", "content": "<MASKED>"}]
refresh_proxy_server_request_body_snapshot(data)
stored = json.loads(
_get_proxy_server_request_for_spend_logs_payload(