From 192eed8b4a298a0a5a6d02e2f73b33e992090a6b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:59:13 -0700 Subject: [PATCH] test: type the new parametrized tests and keep the recorder from changing its argument --- .../datadog/test_datadog_cost_management.py | 4 ++- .../test_prometheus_spend_logs_metadata.py | 4 ++- tests/unit/test_router_silent_experiment.py | 31 +++++++++++++------ 3 files changed, 28 insertions(+), 11 deletions(-) diff --git a/tests/unit/integrations/datadog/test_datadog_cost_management.py b/tests/unit/integrations/datadog/test_datadog_cost_management.py index 0be299e96db..5dea4987087 100644 --- a/tests/unit/integrations/datadog/test_datadog_cost_management.py +++ b/tests/unit/integrations/datadog/test_datadog_cost_management.py @@ -330,7 +330,9 @@ async def test_extract_tags_nested_metadata_allowlisted(clean_env): ids=["mirrored call", "direct call with a client flag"], ) @pytest.mark.asyncio -async def test_extract_tags_takes_the_silent_experiment_flag_only_from_the_router(clean_env, metadata, expected_tag): +async def test_extract_tags_takes_the_silent_experiment_flag_only_from_the_router( + clean_env: None, metadata: dict[str, object], expected_tag: str | None +): logger: Final = DatadogCostManagementLogger(cost_tag_keys=["is_silent_experiment"]) log: Final = StandardLoggingPayload( custom_llm_provider="openai", diff --git a/tests/unit/integrations/test_prometheus_spend_logs_metadata.py b/tests/unit/integrations/test_prometheus_spend_logs_metadata.py index 2b37da0eacb..608195a9df8 100644 --- a/tests/unit/integrations/test_prometheus_spend_logs_metadata.py +++ b/tests/unit/integrations/test_prometheus_spend_logs_metadata.py @@ -229,7 +229,9 @@ def test_project_alias_accessible_without_prefix(monkeypatch): ids=["mirrored call", "direct call with a client flag"], ) def test_silent_experiment_label_comes_only_from_the_router_flag( - monkeypatch, standard_logging_metadata, expected_labels + monkeypatch: pytest.MonkeyPatch, + standard_logging_metadata: dict[str, object], + expected_labels: dict[str, str], ): monkeypatch.setattr("litellm.custom_prometheus_metadata_labels", ["metadata.is_silent_experiment"]) diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index 19dab0a4a3b..bf175012090 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -5,7 +5,7 @@ import threading import time from collections.abc import Callable, Generator, Mapping from types import SimpleNamespace -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -27,9 +27,13 @@ class _RecordingLogger(CustomLogger): self.failure_kwargs: list[dict[str, object]] = [] self._logged: Final = threading.Condition() - def _record(self, calls: list[dict[str, object]], kwargs: dict[str, object]) -> None: + def _record(self, outcome: Literal["success", "failure"], kwargs: dict[str, object]) -> None: with self._logged: - calls.append(kwargs) + match outcome: + case "success": + self.success_kwargs.append(kwargs) + case "failure": + self.failure_kwargs.append(kwargs) self._logged.notify_all() async def async_log_success_event( @@ -39,7 +43,7 @@ class _RecordingLogger(CustomLogger): start_time: datetime.datetime, end_time: datetime.datetime, ) -> None: - self._record(self.success_kwargs, kwargs) + self._record("success", kwargs) async def async_log_failure_event( self, @@ -48,7 +52,7 @@ class _RecordingLogger(CustomLogger): start_time: datetime.datetime, end_time: datetime.datetime, ) -> None: - self._record(self.failure_kwargs, kwargs) + self._record("failure", kwargs) def wait_until(self, logged_enough: Callable[[], bool], timeout: float = 30.0) -> None: with self._logged: @@ -70,7 +74,7 @@ class _SyncAndAsyncRecordingLogger(_RecordingLogger): start_time: datetime.datetime, end_time: datetime.datetime, ) -> None: - self._record(self.success_kwargs, kwargs) + self._record("success", kwargs) async def _settle_shared_logging_worker() -> None: @@ -396,7 +400,9 @@ def test_standard_logging_payload_marks_only_the_mirrored_calls_on_the_sync_path ) assert response.choices[0].message.content == "pong" sync_and_async_recording_logger.wait_until( - lambda: _silent_flag_by_group(sync_and_async_recording_logger).keys() >= {"primary-model", "shadow-a", "shadow-b"} + lambda: ( + _silent_flag_by_group(sync_and_async_recording_logger).keys() >= {"primary-model", "shadow-a", "shadow-b"} + ) ) assert _silent_flag_by_group(sync_and_async_recording_logger) == { @@ -435,7 +441,10 @@ async def test_a_failed_mirrored_call_reaches_failure_callbacks_marked_as_mirror await asyncio.to_thread(recording_logger.wait_until, lambda: len(recording_logger.failure_kwargs) >= 1) assert { - (call["standard_logging_object"]["model_group"], call["standard_logging_object"]["metadata"].get("is_silent_experiment")) + ( + call["standard_logging_object"]["model_group"], + call["standard_logging_object"]["metadata"].get("is_silent_experiment"), + ) for call in recording_logger.failure_kwargs } == {("shadow-a", True)} @@ -995,7 +1004,11 @@ async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_mar }, { "model_name": "shadow-a", - "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-key", "mock_response": shadow_outcome}, + "litellm_params": { + "model": "openai/gpt-5.4-nano", + "api_key": "fake-key", + "mock_response": shadow_outcome, + }, }, ] )