From a75b082f22e6027eed9e54df9c5ac739b4abf9d7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 01:53:40 -0800 Subject: [PATCH] test(google_genai): cover the limiter's Google config precedence Covers an empty config next to a filled generationConfig and both config spellings in one request, for the cap check and the injected cap, and types the new test params as Mapping. Cherry-picked from 97ddba26946 (litellm_genai_snake_case_generation_config). The line's test files gain the collections.abc.Mapping and typing.Final imports the tests use. --- .../google_genai/test_google_genai_adapter.py | 3 +- .../google_genai/test_google_genai_main.py | 8 +-- tests/unit/proxy/hooks/test_tpm_concurrent.py | 59 ++++++++++++++++++- 3 files changed, 64 insertions(+), 6 deletions(-) diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 42f193b0425..a89c945b914 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -4,6 +4,7 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json import unittest +from collections.abc import Mapping import pytest @@ -94,7 +95,7 @@ def test_list_contents_transformation(): ], ids=["camelCase", "snake_case"], ) -def test_config_parameter_mapping(config: dict[str, object]): +def test_config_parameter_mapping(config: Mapping[str, object]): """Test that config parameters are correctly mapped""" from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 8c9166be4a0..7893a343339 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -4,7 +4,7 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json - +from collections.abc import Mapping from typing import Final import httpx @@ -324,7 +324,7 @@ def _httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) -async def _call_entrypoint(entrypoint: str, output_cap: dict[str, dict[str, int]]) -> None: +async def _call_entrypoint(entrypoint: str, output_cap: Mapping[str, Mapping[str, int]]) -> None: from litellm.google_genai.main import ( agenerate_content, agenerate_content_stream, @@ -357,7 +357,7 @@ async def _call_entrypoint(entrypoint: str, output_cap: dict[str, dict[str, int] "entrypoint", ["generate_content", "agenerate_content", "generate_content_stream", "agenerate_content_stream"] ) async def test_adapter_path_sends_the_output_cap_in_every_spelling( - _httpx_transport: None, entrypoint: str, output_cap: dict[str, dict[str, int]] + _httpx_transport: None, entrypoint: str, output_cap: Mapping[str, Mapping[str, int]] ) -> None: with respx.mock() as upstream: route: Final = upstream.post(f"{_CHAT_UPSTREAM}/chat/completions").mock(side_effect=_chat_reply) @@ -369,7 +369,7 @@ async def test_adapter_path_sends_the_output_cap_in_every_spelling( @pytest.mark.asyncio @pytest.mark.parametrize("output_cap", _OUTPUT_CAP_SPELLINGS) async def test_native_gemini_request_carries_the_output_cap_in_every_spelling( - _httpx_transport: None, output_cap: dict[str, dict[str, int]] + _httpx_transport: None, output_cap: Mapping[str, Mapping[str, int]] ) -> None: from litellm.google_genai.main import agenerate_content diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 87552924a9a..dc7659d946b 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -16,7 +16,7 @@ Redis. import asyncio from datetime import datetime, timedelta -from typing import Any, Dict +from typing import Any, Dict, Final import pytest @@ -1475,6 +1475,63 @@ async def test_project_otpm_injects_google_genai_native_output_cap( assert "max_tokens" not in data +@pytest.mark.asyncio +@pytest.mark.parametrize( + "configs", + [ + {"config": {}, "generationConfig": {"maxOutputTokens": 100}}, + {"generationConfig": {"maxOutputTokens": 100}, "generation_config": {"max_output_tokens": 10}}, + {"config": {"maxOutputTokens": 100}, "generationConfig": {"maxOutputTokens": 10}}, + ], +) +async def test_project_otpm_checks_the_first_non_empty_google_genai_config(rate_limiter, configs): + handler, cache = rate_limiter + model = "gemini/gemini-3-flash-preview" + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-google-genai-config-precedence"), + project_id="project-google-genai-config-precedence", + project_metadata={"model_otpm_limit": {model: 50}}, + ) + + with pytest.raises(Exception, match="Rate limit exceeded for model_per_project_otpm") as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={"model": model, "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], **configs}, + call_type="agenerate_content", + ) + + assert getattr(exc_info.value, "status_code", None) == 429 + + +@pytest.mark.asyncio +async def test_project_otpm_injects_cap_into_the_first_non_empty_google_genai_config(rate_limiter): + handler, cache = rate_limiter + model = "gemini/gemini-3-flash-preview" + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-google-genai-config-inject-precedence"), + project_id="project-google-genai-config-inject-precedence", + project_metadata={"model_otpm_limit": {model: 40}}, + ) + data: Final = { + "model": model, + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "config": {}, + "generationConfig": {"temperature": 0}, + "generation_config": {"max_output_tokens": 5000}, + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="agenerate_content", + ) + + assert data["generationConfig"] == {"temperature": 0, "maxOutputTokens": 10} + assert data["config"] == {} + + @pytest.mark.asyncio async def test_project_otpm_over_limit_rolls_back_itpm_reservation(rate_limiter): """