From 909ab23b89589375d8037319ff32fea048d710fe Mon Sep 17 00:00:00 2001 From: Bisma Nawaz Date: Fri, 21 Aug 2026 03:42:34 +0500 Subject: [PATCH] test: annotate parametrized traffic-type test inputs --- tests/test_litellm/llms/gemini/test_cost_calculator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index c44f29ba168..1f7bfa69527 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -318,7 +318,9 @@ def test_gemini_image_generation_cost_no_web_search_when_absent(): ("SOMETHING_UNKNOWN", None), ], ) -def test_map_traffic_type_to_service_tier(traffic_type, expected_service_tier): +def test_map_traffic_type_to_service_tier( + traffic_type: str | None, expected_service_tier: str | None +): """ Gemini/Vertex usageMetadata.trafficType maps to the LiteLLM service_tier that selects flex/priority cost keys. ON_DEMAND_FLEX (Vertex's flex opt-in