test(passthrough): type the Deepgram speak tests

Every new test and helper now annotates its parameters and returns, and every new binding is Final, per AGENTS.md. The two CustomLogger overrides keep the base class's untyped signature.
This commit is contained in:
shrey kharbanda 2026-10-26 18:57:56 -07:00
parent 64c6955c96
commit 71d23db938
4 changed files with 84 additions and 64 deletions

View file

@ -344,7 +344,7 @@ def test_deepgram_listen_addon_pricing_models(upstream_url: str, expected: tuple
)
def test_aura_rows_bill_speech_per_character_at_the_published_rate(
pricing_model: str, usd_per_thousand_characters: float
):
) -> None:
"""Pay as you go rates from https://deepgram.com/pricing, checked 2026-10-26"""
prompt_cost, completion_cost = litellm.cost_per_token(
model=pricing_model, custom_llm_provider="deepgram", call_type="speech", prompt_characters=1000
@ -384,7 +384,7 @@ SPEAK_BASE: Final = "wss://api.deepgram.com/v1/speak"
),
],
)
def test_deepgram_speak_websocket_target(api_base: str | None, query_string: str, expected: str):
def test_deepgram_speak_websocket_target(api_base: str | None, query_string: str, expected: str) -> None:
assert deepgram_speak_websocket_target(api_base=api_base, query_string=query_string) == expected
@ -392,7 +392,7 @@ def test_deepgram_speak_websocket_target(api_base: str | None, query_string: str
"query_string",
["", "model=aura-2-thalia-en", "model=aura-2-thalia-en&model=aura-2-zeus-en", "encoding=mulaw&model=aura-luna-en"],
)
def test_the_authorized_speak_model_is_the_only_model_the_upstream_target_carries(query_string: str):
def test_the_authorized_speak_model_is_the_only_model_the_upstream_target_carries(query_string: str) -> None:
"""The key is checked against deepgram_speak_requested_model, so the upstream must use exactly that voice"""
target: Final = deepgram_speak_websocket_target(api_base=None, query_string=query_string)
@ -421,7 +421,9 @@ def test_the_authorized_speak_model_is_the_only_model_the_upstream_target_carrie
pytest.param("nova-3", "nova-3", id="listen model"),
],
)
def test_deepgram_speak_registry_key_names_the_aura_generation_row_else_the_voice(model: str, expected_family: str):
def test_deepgram_speak_registry_key_names_the_aura_generation_row_else_the_voice(
model: str, expected_family: str
) -> None:
target: Final = deepgram_speak_websocket_target(api_base=None, query_string=f"model={model}")
assert deepgram_speak_registry_key(target) == f"deepgram/{expected_family}"
@ -438,19 +440,19 @@ def test_deepgram_speak_registry_key_names_the_aura_generation_row_else_the_voic
pytest.param("nova-3", False, id="a listen model has no speech row"),
],
)
def test_deepgram_speak_is_priced(model: str, expected: bool):
def test_deepgram_speak_is_priced(model: str, expected: bool) -> None:
assert deepgram_speak_is_priced(f"{SPEAK_BASE}?model={model}") is expected
@pytest.mark.usefixtures("local_model_cost_map")
def test_deepgram_speak_is_not_priced_once_its_family_row_is_gone(monkeypatch: pytest.MonkeyPatch):
def test_deepgram_speak_is_not_priced_once_its_family_row_is_gone(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delitem(litellm.model_cost, "deepgram/aura-2")
assert deepgram_speak_is_priced(f"{SPEAK_BASE}?model=aura-2-thalia-en") is False
@pytest.mark.usefixtures("local_model_cost_map")
def test_a_cost_map_row_named_after_a_new_voice_prices_it_at_that_rate(monkeypatch: pytest.MonkeyPatch):
def test_a_cost_map_row_named_after_a_new_voice_prices_it_at_that_rate(monkeypatch: pytest.MonkeyPatch) -> None:
"""The 1008 refusal tells the operator to add the voice's row, so adding it must be enough"""
url: Final = f"{SPEAK_BASE}?model=aura-3-foo-en"
monkeypatch.setitem(
@ -488,5 +490,5 @@ def test_a_cost_map_row_named_after_a_new_voice_prices_it_at_that_rate(monkeypat
pytest.param("1" * 5_000, 0, id="integer longer than the parser allows"),
],
)
def test_deepgram_speak_characters_counts_speak_text_and_never_raises(frame: str | bytes, expected: int):
def test_deepgram_speak_characters_counts_speak_text_and_never_raises(frame: str | bytes, expected: int) -> None:
assert deepgram_speak_characters(frame) == expected

View file

@ -43,7 +43,7 @@ pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
pytest.param("https://gateway.internal/deepgram/v1/speak", "POST", False, id="http passthrough to a speak url"),
],
)
def test_is_deepgram_speak_route(url_route: str, request_method: str, expected: bool):
def test_is_deepgram_speak_route(url_route: str, request_method: str, expected: bool) -> None:
assert DeepgramSpeakPassthroughLoggingHandler.is_deepgram_speak_route(url_route, request_method) is expected
@ -70,7 +70,7 @@ def _handler_cost(upstream_url: str, characters: object) -> float | None:
return response_cost
def test_handler_bills_spoken_characters_at_the_aura_2_rate_and_names_the_voice():
def test_handler_bills_spoken_characters_at_the_aura_2_rate_and_names_the_voice() -> None:
logging_obj: Final = _logging_obj()
handler_result: Final = DeepgramSpeakPassthroughLoggingHandler.deepgram_speak_passthrough_handler(
@ -91,7 +91,7 @@ def test_handler_bills_spoken_characters_at_the_aura_2_rate_and_names_the_voice(
)
def test_handler_bills_an_aura_1_voice_at_the_aura_1_rate():
def test_handler_bills_an_aura_1_voice_at_the_aura_1_rate() -> None:
assert _handler_cost(AURA_1_URL, 1000) == pytest.approx(1000 * AURA_1_USD_PER_CHARACTER)
@ -99,12 +99,12 @@ def test_handler_bills_an_aura_1_voice_at_the_aura_1_rate():
"characters",
[pytest.param(0, id="nothing spoken"), pytest.param(None, id="no count")],
)
def test_handler_bills_zero_for_a_session_without_speak_text(characters: object):
def test_handler_bills_zero_for_a_session_without_speak_text(characters: object) -> None:
assert _handler_cost(AURA_2_URL, characters) == 0.0
@pytest.mark.parametrize("model", ["flux-tts-en", "aura-3-x-en", "aura-x", "nova-3"])
def test_handler_reports_no_cost_rather_than_a_wrong_one_for_an_unpriced_voice(model: str):
def test_handler_reports_no_cost_rather_than_a_wrong_one_for_an_unpriced_voice(model: str) -> None:
assert _handler_cost(f"wss://api.deepgram.com/v1/speak?model={model}", 1000) is None
@ -126,7 +126,7 @@ class _CapturingLogger(CustomLogger):
)
async def test_success_handler_dispatches_deepgram_speak_and_logs_character_based_spend(
monkeypatch: pytest.MonkeyPatch, characters: int, expected_cost: float
):
) -> None:
"""Drives the shared passthrough success handler the way the relay does at socket close and reads what a spend
logger receives"""
capturing_logger: Final = _CapturingLogger()

View file

@ -108,13 +108,13 @@ async def _serve(websocket: _FakeWebSocket, user_api_key_dict: UserAPIKeyAuth |
return relay
def test_deepgram_websocket_routes_registered():
def test_deepgram_websocket_routes_registered() -> None:
ws_paths = {route.path for route in router.routes if isinstance(route, WebSocketRoute)}
assert {*LISTEN_PATHS, *SPEAK_PATHS} <= ws_paths
@pytest.mark.parametrize("path", (*LISTEN_PATHS, *SPEAK_PATHS))
def test_deepgram_websocket_routes_are_lazily_loaded_mapped_pass_through_routes(path):
def test_deepgram_websocket_routes_are_lazily_loaded_mapped_pass_through_routes(path: str) -> None:
"""The route must be reachable before the passthrough module is imported and must be authed and
billed as a mapped pass-through route like the other provider prefixes."""
feature = next(feature for feature in LAZY_FEATURES if feature.name == "llm_passthrough")
@ -492,7 +492,9 @@ def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm
),
],
)
def test_deepgram_routes_accept_the_litellm_key_in_a_deepgram_sdk_token_header(url, expected_target, monkeypatch):
def test_deepgram_routes_accept_the_litellm_key_in_a_deepgram_sdk_token_header(
url: str, expected_target: str, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Deepgram SDKs send ``Authorization: Token <key>``, so pointing one at the proxy with a LiteLLM key as its API
key must authenticate through the real key auth path and still send only the server's key upstream"""
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
@ -520,7 +522,7 @@ def test_deepgram_routes_accept_the_litellm_key_in_a_deepgram_sdk_token_header(u
]
def test_a_non_deepgram_websocket_route_still_refuses_a_token_header():
def test_a_non_deepgram_websocket_route_still_refuses_a_token_header() -> None:
relay: Final = _FakeRelay()
client: Final = TestClient(_app_with_relay(relay))
@ -545,7 +547,9 @@ async def _serve_speak(websocket: _FakeWebSocket, user_api_key_dict: UserAPIKeyA
@pytest.mark.parametrize("path", SPEAK_PATHS)
async def test_deepgram_speak_forwards_query_injects_only_provider_auth_and_counts_characters(path, monkeypatch):
async def test_deepgram_speak_forwards_query_injects_only_provider_auth_and_counts_characters(
path: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
websocket: Final = _FakeWebSocket(path, "model=aura-2-thalia-en&encoding=linear16&sample_rate=24000")
caller: Final = UserAPIKeyAuth(api_key="sk-litellm-virtual", team_id="team-tts")
@ -569,7 +573,7 @@ async def test_deepgram_speak_forwards_query_injects_only_provider_auth_and_coun
assert websocket.closed is None
async def test_deepgram_listen_does_not_count_client_text(monkeypatch):
async def test_deepgram_listen_does_not_count_client_text(monkeypatch: pytest.MonkeyPatch) -> None:
"""Listen bills audio seconds from upstream frames, so it must keep handing the relay no counter"""
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
@ -596,7 +600,9 @@ async def test_deepgram_listen_does_not_count_client_text(monkeypatch):
),
],
)
async def test_deepgram_speak_names_exactly_one_model_upstream(query, expected_target, monkeypatch):
async def test_deepgram_speak_names_exactly_one_model_upstream(
query: str, expected_target: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
@ -605,7 +611,9 @@ async def test_deepgram_speak_names_exactly_one_model_upstream(query, expected_t
assert [call.target for call in relay.calls] == [expected_target]
async def test_deepgram_speak_honours_server_configured_api_base_and_ignores_the_callers(monkeypatch):
async def test_deepgram_speak_honours_server_configured_api_base_and_ignores_the_callers(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The server's Deepgram key must only ever go to the server-configured host"""
monkeypatch.setenv("DEEPGRAM_API_BASE", "https://api.eu.deepgram.com/v1/")
websocket: Final = _FakeWebSocket(
@ -620,7 +628,7 @@ async def test_deepgram_speak_honours_server_configured_api_base_and_ignores_the
]
async def test_deepgram_speak_closes_before_accepting_when_provider_credentials_missing():
async def test_deepgram_speak_closes_before_accepting_when_provider_credentials_missing() -> None:
websocket: Final = _FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en")
with patch(GET_CREDENTIALS, return_value=None):
@ -643,7 +651,9 @@ async def test_deepgram_speak_closes_before_accepting_when_provider_credentials_
pytest.param("model=aura-2-thalia-en", "deepgram/aura-2", "deepgram/aura-2", id="family row removed"),
],
)
async def test_deepgram_speak_refuses_sessions_it_cannot_price(query, missing_key, removed_row, monkeypatch):
async def test_deepgram_speak_refuses_sessions_it_cannot_price(
query: str, missing_key: str, removed_row: str | None, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A session without a per-character price would be logged at no cost, so the proxy closes it before Deepgram is
contacted and names the registry row to add"""
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
@ -662,7 +672,7 @@ async def test_deepgram_speak_refuses_sessions_it_cannot_price(query, missing_ke
assert "dg-provider-key" not in websocket.closed[1]
def test_deepgram_speak_rejects_connections_without_a_litellm_key():
def test_deepgram_speak_rejects_connections_without_a_litellm_key() -> None:
relay: Final = _FakeRelay()
client: Final = TestClient(_app_with_relay(relay))
@ -694,7 +704,9 @@ def test_deepgram_speak_rejects_connections_without_a_litellm_key():
pytest.param("model=&encoding=mulaw", None, id="voice blank, default denied"),
],
)
def test_deepgram_speak_authorizes_the_voice_it_will_actually_send_upstream(query, expected_target, monkeypatch):
def test_deepgram_speak_authorizes_the_voice_it_will_actually_send_upstream(
query: str, expected_target: str | None, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A key allowed only one voice must not reach another by leaving ``model`` out or naming a second one, because
the real key auth path must see the same model the upstream target carries"""
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)

View file

@ -7824,24 +7824,24 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock:
return increment_spend_counters
DEEPGRAM_SPEAK_TARGET = "wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en"
SPEAK_AUDIO_CHUNK = bytes(range(256)) * 8
SPOKEN_TEXT = "The quick brown fox jumps over the lazy dog."
SPEAK_FRAME = json.dumps({"type": "Speak", "text": SPOKEN_TEXT})
FLUSH_FRAME = json.dumps({"type": "Flush"})
FLUSHED_FRAME = json.dumps({"type": "Flushed", "sequence_id": 0})
SPEAK_CLOSE_FRAME = json.dumps({"type": "Close"})
CLIENT_BINARY_FRAME = b"\x00\xffclient-binary"
SPEAK_SESSION_FRAMES = ({"text": SPEAK_FRAME}, {"bytes": CLIENT_BINARY_FRAME}, {"text": FLUSH_FRAME})
SPEAK_SESSION_CLIENT_LENGTH = sum(map(len, (SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME)))
RELAY_HANG_GUARD_SECONDS = 5.0
DEEPGRAM_SPEAK_TARGET: Final = "wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en"
SPEAK_AUDIO_CHUNK: Final = bytes(range(256)) * 8
SPOKEN_TEXT: Final = "The quick brown fox jumps over the lazy dog."
SPEAK_FRAME: Final = json.dumps({"type": "Speak", "text": SPOKEN_TEXT})
FLUSH_FRAME: Final = json.dumps({"type": "Flush"})
FLUSHED_FRAME: Final = json.dumps({"type": "Flushed", "sequence_id": 0})
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:
"""Answers each Speak with audio as soon as it arrives and ends the session on Close, the way Deepgram speak
does, so a relay that held audio back until the end would never see the Close it waits for"""
def __init__(self, close_code: int = 1000, close_reason: str = ""):
def __init__(self, close_code: int = 1000, close_reason: str = "") -> None:
self.received: list[str | bytes] = []
self.close = AsyncMock()
self._close_code = close_code
@ -7862,36 +7862,40 @@ class SpeakingUpstreamWebSocket:
from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK
from websockets.frames import Close
reply = await self._replies.get()
reply: Final = await self._replies.get()
if reply is not None:
return reply
close = Close(self._close_code, self._close_reason)
closed = ConnectionClosedOK if self._close_code == 1000 else ConnectionClosedError
close: Final = Close(self._close_code, self._close_reason)
closed: Final = ConnectionClosedOK if self._close_code == 1000 else ConnectionClosedError
raise closed(rcvd=close, sent=close, rcvd_then_sent=True)
def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRAMES):
def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRAMES) -> MagicMock:
"""A client socket that sends ``frames``, then Close once the first audio has reached it"""
first_audio = asyncio.Event()
first_audio: Final = asyncio.Event()
async def script():
async def script() -> AsyncIterator[dict[str, str | bytes]]:
for frame in frames:
yield {"type": "websocket.receive", **frame}
await first_audio.wait()
yield {"type": "websocket.receive", "text": SPEAK_CLOSE_FRAME}
await asyncio.Event().wait()
steps = script()
steps: Final = script()
async def receive():
async def receive() -> dict[str, str | bytes]:
return await anext(steps)
websocket = _client_websocket(receive)
websocket: Final = _client_websocket(receive)
websocket.send_bytes = AsyncMock(side_effect=lambda data: first_audio.set())
return websocket
async def _relay_deepgram_speak(upstream_ws, websocket, count_client_text):
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),
@ -7918,13 +7922,13 @@ async def _relay_deepgram_speak(upstream_ws, websocket, count_client_text):
@pytest.mark.asyncio
async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed():
async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed() -> 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"""
upstream_ws = SpeakingUpstreamWebSocket()
websocket = _speak_client()
upstream_ws: Final = SpeakingUpstreamWebSocket()
websocket: Final = _speak_client()
success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
success_handler: Final = await _relay_deepgram_speak(upstream_ws, websocket, 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]
@ -7936,11 +7940,11 @@ async def test_websocket_passthrough_counts_client_text_without_changing_what_is
@pytest.mark.asyncio
async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_body():
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 = SpeakingUpstreamWebSocket()
upstream_ws: Final = SpeakingUpstreamWebSocket()
success_handler = await _relay_deepgram_speak(upstream_ws, _speak_client(), count_client_text=None)
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"] == {}
@ -7957,11 +7961,11 @@ 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
):
upstream_ws = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason)
websocket = _speak_client()
) -> None:
upstream_ws: Final = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason)
websocket: Final = _speak_client()
success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
success_handler: Final = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len)
websocket.close.assert_awaited_once_with(code=close_code, reason=close_reason)
assert success_handler.call_args.kwargs["request_body"] == {
@ -7980,15 +7984,17 @@ class _CapturingSpendLogger(CustomLogger):
@pytest.mark.asyncio
@pytest.mark.usefixtures("local_model_cost_map")
async def test_deepgram_speak_session_logs_character_spend_once_and_never_the_spoken_text(monkeypatch):
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 = _CapturingSpendLogger()
capturing_logger: Final = _CapturingSpendLogger()
monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
queued: list[Coroutine[object, object, None]] = [] # mutable-ok: the fake logging worker keeps each coroutine
upstream_ws = SpeakingUpstreamWebSocket()
queued: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: the fake worker keeps each coroutine
upstream_ws: Final = SpeakingUpstreamWebSocket()
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as proxy_logging,
@ -8021,7 +8027,7 @@ async def test_deepgram_speak_session_logs_character_spend_once_and_never_the_sp
await coroutine
assert len(capturing_logger.payloads) == 1
payload = capturing_logger.payloads[0]
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