mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-15 18:24:07 +00:00
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
This commit is contained in:
+8
-4
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user