From 01ab5bb7a1b00124df57dd7152546cede6360094 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 17:16:34 -0700 Subject: [PATCH] test(proxy): type the per-chunk streaming hook test helpers --- .../proxy/test_common_request_processing.py | 12 +++--- .../test_callback_capabilities_class.py | 2 +- .../proxy_logging/test_streaming_hooks.py | 42 +++++++++++-------- 3 files changed, 33 insertions(+), 23 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d8c061fc6b7..f7bb4f54f0a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4973,16 +4973,18 @@ class TestAsyncStreamingDataGeneratorFastPath: hook_spy.assert_not_awaited() @pytest.mark.asyncio - async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail(self, monkeypatch): + async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")]) ProxyLogging._callback_capabilities_cache.clear() - proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) - hook_spy = AsyncMock(side_effect=lambda **kw: kw["response"]) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=MagicMock()) + hook_spy: Final = AsyncMock(side_effect=lambda **kw: kw["response"]) monkeypatch.setattr(proxy_logging_obj, "async_post_call_streaming_hook", hook_spy) - chunks = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"] - out = [ + chunks: Final = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"] + out: Final = [ c async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=self._aiter(chunks), diff --git a/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py b/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py index 24ca87be156..09cdd31598a 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py +++ b/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py @@ -40,7 +40,7 @@ class _OverridesPerChunk(CustomLogger): class _PerChunkGuardrail(CustomGuardrail): - async def async_post_call_streaming_hook(self, *args, **kwargs): + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> None: return None diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index 1b48c62f1be..b19b31413e0 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,7 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock @@ -389,24 +389,26 @@ class _RequestOnlyGuardrail(CustomGuardrail): super().__init__(guardrail_name="request-only", event_hook=GuardrailEventHooks.pre_call, default_on=True) self.asked: List[GuardrailEventHooks] = [] - def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool: + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: self.asked.append(event_type) return super().should_run_guardrail(data=data, event_type=event_type) @pytest.mark.asyncio async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardrail( - proxy_logging, make_user_api_key_auth, monkeypatch -): - guardrail = _RequestOnlyGuardrail() + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _RequestOnlyGuardrail() monkeypatch.setattr(litellm, "callbacks", [guardrail]) - chunk = ModelResponseStream( + chunk: Final = ModelResponseStream( id="rid", choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}], created=0, model="gpt-4o-mini", ) - out = await proxy_logging.async_post_call_streaming_hook( + out: Final = await proxy_logging.async_post_call_streaming_hook( data={"messages": [{"role": "user", "content": "hi"}]}, response=chunk, user_api_key_dict=make_user_api_key_auth(), @@ -416,23 +418,25 @@ async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardr class _PerChunkRewriter(CustomLogger): - async def async_post_call_streaming_hook(self, user_api_key_dict, response): # type: ignore[override] + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> str: return f"modified-{response}" @pytest.mark.asyncio async def test_async_post_call_streaming_hook_skips_a_request_only_guardrail_beside_a_per_chunk_callback( - proxy_logging, make_user_api_key_auth, monkeypatch -): - guardrail = _RequestOnlyGuardrail() + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _RequestOnlyGuardrail() monkeypatch.setattr(litellm, "callbacks", [guardrail, _PerChunkRewriter()]) - chunk = ModelResponseStream( + chunk: Final = ModelResponseStream( id="rid", choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}], created=0, model="gpt-4o-mini", ) - out = await proxy_logging.async_post_call_streaming_hook( + out: Final = await proxy_logging.async_post_call_streaming_hook( data={"messages": [{"role": "user", "content": "hi"}]}, response=chunk, user_api_key_dict=make_user_api_key_auth(), @@ -474,8 +478,10 @@ async def _say_hi_through_chunk_hooks(proxy_logging: ProxyLogging, user_api_key_ @pytest.mark.asyncio async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_guardrail( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: guardrail: Final = _RequestOnlyGuardrail() monkeypatch.setattr(litellm, "callbacks", [guardrail, _instance_rewriter()]) assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi" @@ -484,8 +490,10 @@ async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_gua @pytest.mark.asyncio async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_class_hook( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(litellm, "callbacks", [_ClassChunkLogger(), _instance_rewriter()]) assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"