From 5c20617a21f7155b313cfc5a42e91b929072ff3c Mon Sep 17 00:00:00 2001 From: joereyna Date: Wed, 11 Mar 2026 20:05:40 -0700 Subject: [PATCH] fix: mock completion_cost in routing test and restore helper consistency - Mock litellm.completion_cost in test_pass_through_success_handler_gemini_routing to decouple it from model_prices_and_context_window.json; prevents the same breakage if gemini-2.0-flash is ever removed from the pricing map - Revert _create_passthrough_logging_payload URL back to gemini-1.5-flash to eliminate inconsistency with the other tests that use gemini-1.5-flash explicitly --- .../test_gemini_passthrough_logging_handler.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py index 1a348e14ca..be38d08327 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -68,7 +68,7 @@ class TestGeminiPassthroughLoggingHandler: def _create_passthrough_logging_payload(self) -> PassthroughStandardLoggingPayload: """Create a mock passthrough logging payload for testing""" return PassthroughStandardLoggingPayload( - url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent", + url="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, request_method="POST", ) @@ -242,7 +242,11 @@ class TestGeminiPassthroughLoggingHandler: assert "kwargs" in result @pytest.mark.asyncio - async def test_pass_through_success_handler_gemini_routing(self): + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler.litellm.completion_cost", + return_value=0.000050, + ) + async def test_pass_through_success_handler_gemini_routing(self, mock_completion_cost): """Test that the success handler correctly routes Gemini requests to the Gemini handler""" handler = PassThroughEndpointLogging() @@ -277,14 +281,14 @@ class TestGeminiPassthroughLoggingHandler: assert result is None # Verify that the logging object has the cost set (from Gemini handler) - assert mock_logging_obj.model_call_details["response_cost"] is not None + assert mock_logging_obj.model_call_details["response_cost"] == 0.000050 assert mock_logging_obj.model_call_details["model"] == "gemini-2.0-flash" assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" # Verify that _handle_logging was called with the correct kwargs handler._handle_logging.assert_called_once() call_kwargs = handler._handle_logging.call_args[1] - assert call_kwargs["response_cost"] is not None + assert call_kwargs["response_cost"] == 0.000050 assert call_kwargs["model"] == "gemini-2.0-flash" assert call_kwargs["custom_llm_provider"] == "gemini"