From df41566be9949dae06cc2c63db9a30d3b76dc06b Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Mon, 26 Oct 2026 19:05:55 -0700 Subject: [PATCH] test(passthrough): bill the speak relay tests through the real success handler The relay tests reuse the hang-up PR's helper, which runs the real success handler and returns what a spend logger receives, instead of mocking our own handler. tests/unit/AGENTS.md rules out doubles on our own code. The empty request body test is dropped, since the spend a session logs is now the assertion --- ...pgram_speak_passthrough_logging_handler.py | 2 +- .../test_pass_through_endpoints.py | 204 +++++++++--------- 2 files changed, 97 insertions(+), 109 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_speak_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_speak_passthrough_logging_handler.py index 6186d8d53fd..82e0d388607 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_speak_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_deepgram_speak_passthrough_logging_handler.py @@ -113,7 +113,7 @@ class _CapturingLogger(CustomLogger): super().__init__() self.payloads: list[StandardLoggingPayload] = [] # mutable-ok: the logger appends each payload it receives - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: self.payloads.append(kwargs["standard_logging_object"]) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index de9db053909..797e743ef2c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7,7 +7,7 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Coroutine, Mapping from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime @@ -67,10 +67,10 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, - WEBSOCKET_CLIENT_TEXT_COUNT_KEY, EndpointType, PassthroughStandardLoggingPayload, ) +from litellm.types.utils import StandardLoggingPayload from tests._master_key import MASTER_KEY from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -7717,6 +7717,66 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text upstream_ws.close.assert_awaited_once() +class _CapturingSpendLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.payloads: list[StandardLoggingPayload] = [] # mutable-ok: the logger appends each payload it receives + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.payloads.append(kwargs["standard_logging_object"]) + + +@dataclass(frozen=True, slots=True) +class _RelayedSession: + spend_logs: tuple[StandardLoggingPayload, ...] + failure_hook: AsyncMock + + +async def _relay_and_capture_spend( + monkeypatch: pytest.MonkeyPatch, + upstream_ws: "FakeUpstreamWebSocket | SpeakingUpstreamWebSocket", + websocket: MagicMock, + target: str, + endpoint: str, + count_client_text: Callable[[str | bytes], int] | None = None, +) -> _RelayedSession: + """Runs the relay with the real success handler and returns what a spend logger receives""" + capturing_logger: Final = _CapturingSpendLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger]) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + queued: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: the fake worker keeps each coroutine + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + return_value=FakeUpstreamConnect(upstream_ws), + ), + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, + ): + proxy_logging.pre_call_hook = AsyncMock(return_value={}) + proxy_logging.post_call_success_hook = AsyncMock() + proxy_logging.post_call_failure_hook = AsyncMock() + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: queued.append(async_coroutine) + ) + await websocket_passthrough_request( + websocket=websocket, + target=target, + custom_headers={"Authorization": "Token dg-provider-key"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint=endpoint, + accept_websocket=False, + count_client_text=count_client_text, + ) + for coroutine in queued: + await coroutine + return _RelayedSession( + spend_logs=tuple(capturing_logger.payloads), failure_hook=proxy_logging.post_call_failure_hook + ) + + @pytest.mark.asyncio async def test_websocket_passthrough_vertex_live_setup_ack_names_the_model_but_is_not_billed_as_usage(): """Vertex Live keeps its special first frame: the setup acknowledgement is forwarded verbatim, read for the @@ -7820,7 +7880,6 @@ SPEAK_CLOSE_FRAME: Final = json.dumps({"type": "Close"}) CLIENT_BINARY_FRAME: Final = b"\x00\xffclient-binary" SPEAK_SESSION_FRAMES: Final = ({"text": SPEAK_FRAME}, {"bytes": CLIENT_BINARY_FRAME}, {"text": FLUSH_FRAME}) SPEAK_SESSION_CLIENT_LENGTH: Final = sum(map(len, (SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME))) -RELAY_HANG_GUARD_SECONDS: Final = 5.0 class SpeakingUpstreamWebSocket: @@ -7877,67 +7936,35 @@ def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRA return websocket -async def _relay_deepgram_speak( - upstream_ws: SpeakingUpstreamWebSocket, - websocket: MagicMock, - count_client_text: Callable[[str | bytes], int] | None, -) -> AsyncMock: - """Runs the generic relay the way the Deepgram speak route does and returns the success handler mock""" - with ( - _patched_websocket_passthrough_environment(upstream_ws), - patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam - "litellm.proxy.pass_through_endpoints.pass_through_endpoints." - "pass_through_endpoint_logging.pass_through_async_success_handler", - new=AsyncMock(), - ) as success_handler, - ): - await asyncio.wait_for( - websocket_passthrough_request( - websocket=websocket, - target=DEEPGRAM_SPEAK_TARGET, - custom_headers={"Authorization": "Token dg-provider-key"}, - user_api_key_dict=UserAPIKeyAuth(), - forward_headers=False, - endpoint="/deepgram/v1/speak", - accept_websocket=False, - count_client_text=count_client_text, - ), - timeout=RELAY_HANG_GUARD_SECONDS, - ) - return success_handler +def _aura_2_usd_per_character() -> float: + return litellm.model_cost["deepgram/aura-2"]["input_cost_per_character"] @pytest.mark.asyncio -async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed() -> None: +@pytest.mark.usefixtures("local_model_cost_map") +async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed( + monkeypatch: pytest.MonkeyPatch, +) -> None: """Audio reaches the client while the session is open, every client frame reaches the upstream unchanged, and - the success handler gets only the summed count of the text and binary frames, never the frames""" + the session is billed for the summed count of the text and binary frames""" upstream_ws: Final = SpeakingUpstreamWebSocket() websocket: Final = _speak_client() - success_handler: Final = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len) + session: Final = await _relay_and_capture_spend( + monkeypatch, upstream_ws, websocket, DEEPGRAM_SPEAK_TARGET, "/deepgram/v1/speak", count_client_text=len + ) assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME] assert [call.args[0] for call in websocket.send_bytes.await_args_list] == [SPEAK_AUDIO_CHUNK] assert [call.args[0] for call in websocket.send_text.await_args_list] == [FLUSHED_FRAME] - assert success_handler.call_args.kwargs["request_body"] == { - WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH - } + assert [log["response_cost"] for log in session.spend_logs] == [ + pytest.approx(SPEAK_SESSION_CLIENT_LENGTH * _aura_2_usd_per_character()) + ] websocket.close.assert_awaited_once_with() @pytest.mark.asyncio -async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_body() -> None: - """Listen, realtime and every other caller pass no counter, so their success handler sees what it saw before""" - upstream_ws: Final = SpeakingUpstreamWebSocket() - - success_handler: Final = await _relay_deepgram_speak(upstream_ws, _speak_client(), count_client_text=None) - - assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME] - assert success_handler.call_args.kwargs["request_body"] == {} - assert success_handler.call_args.kwargs["response_body"] == [json.loads(FLUSHED_FRAME)] - - -@pytest.mark.asyncio +@pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize( ("close_code", "close_reason"), [ @@ -7946,26 +7973,19 @@ async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_ ], ) async def test_websocket_passthrough_counts_client_text_when_the_upstream_ends_with_an_error_close( - close_code: int, close_reason: str + monkeypatch: pytest.MonkeyPatch, close_code: int, close_reason: str ) -> None: upstream_ws: Final = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason) websocket: Final = _speak_client() - success_handler: Final = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len) + session: Final = await _relay_and_capture_spend( + monkeypatch, upstream_ws, websocket, DEEPGRAM_SPEAK_TARGET, "/deepgram/v1/speak", count_client_text=len + ) websocket.close.assert_awaited_once_with(code=close_code, reason=close_reason) - assert success_handler.call_args.kwargs["request_body"] == { - WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH - } - - -class _CapturingSpendLogger(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.payloads: list[StandardLoggingPayload] = [] # mutable-ok: the logger appends each payload it receives - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.payloads.append(kwargs["standard_logging_object"]) + assert [log["response_cost"] for log in session.spend_logs] == [ + pytest.approx(SPEAK_SESSION_CLIENT_LENGTH * _aura_2_usd_per_character()) + ] @pytest.mark.asyncio @@ -7973,52 +7993,20 @@ class _CapturingSpendLogger(CustomLogger): async def test_deepgram_speak_session_logs_character_spend_once_and_never_the_spoken_text( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Runs the relay with the speak counter and the real success handler, so the spend logger sees what production - writes for one Speak message""" - capturing_logger: Final = _CapturingSpendLogger() - monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger]) - monkeypatch.setattr(litellm, "success_callback", []) - monkeypatch.setattr(litellm, "callbacks", []) - queued: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: the fake worker keeps each coroutine - upstream_ws: Final = SpeakingUpstreamWebSocket() + session: Final = await _relay_and_capture_spend( + monkeypatch, + SpeakingUpstreamWebSocket(), + _speak_client(), + DEEPGRAM_SPEAK_TARGET, + "/deepgram/v1/speak", + count_client_text=deepgram_speak_characters, + ) - with ( - patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging, - patch( - "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", - return_value=FakeUpstreamConnect(upstream_ws), - ), - patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, - ): - proxy_logging.pre_call_hook = AsyncMock(return_value={}) - proxy_logging.post_call_success_hook = AsyncMock() - proxy_logging.post_call_failure_hook = AsyncMock() - worker.ensure_initialized_and_enqueue = MagicMock( - side_effect=lambda async_coroutine: queued.append(async_coroutine) - ) - await asyncio.wait_for( - websocket_passthrough_request( - websocket=_speak_client(), - target=DEEPGRAM_SPEAK_TARGET, - custom_headers={"Authorization": "Token dg-provider-key"}, - user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", team_id="team-tts"), - forward_headers=False, - endpoint="/deepgram/v1/speak", - accept_websocket=False, - count_client_text=deepgram_speak_characters, - ), - timeout=RELAY_HANG_GUARD_SECONDS, - ) - for coroutine in queued: - await coroutine - - assert len(capturing_logger.payloads) == 1 - payload: Final = capturing_logger.payloads[0] - assert payload["model"] == "aura-2-thalia-en" - assert payload["custom_llm_provider"] == "deepgram" - # Aura-2 pay as you go rate, $0.030 per 1,000 characters, from https://deepgram.com/pricing, checked 2026-10-26 - assert payload["response_cost"] == pytest.approx(len(SPOKEN_TEXT) * 0.030 / 1000) - assert SPOKEN_TEXT not in json.dumps(payload, default=str) + assert [(log["model"], log["custom_llm_provider"]) for log in session.spend_logs] == [ + ("aura-2-thalia-en", "deepgram") + ] + assert session.spend_logs[0]["response_cost"] == pytest.approx(len(SPOKEN_TEXT) * _aura_2_usd_per_character()) + assert SPOKEN_TEXT not in json.dumps(session.spend_logs[0], default=str) @pytest.mark.asyncio