diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 15ec440b0bb..7fcdc8473d7 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -4,6 +4,7 @@ from datetime import datetime import contextlib import copy import json +import logging import os from collections.abc import Mapping from dataclasses import dataclass @@ -190,7 +191,9 @@ async def test_url_with_format_param(model, sync_mode, monkeypatch): # URL->image conversion helpers so suite-level network/client state from # earlier tests cannot prevent the mocked provider client from being hit. fake_base64_image = "data:image/png;base64,ZmFrZS1pbWFnZQ==" - monkeypatch.setattr(prompt_factory, "convert_url_to_base64", lambda url: fake_base64_image) + monkeypatch.setattr( + prompt_factory, "convert_url_to_base64", lambda url: fake_base64_image + ) monkeypatch.setattr( prompt_factory.BedrockImageProcessor, "get_image_details", @@ -305,7 +308,9 @@ async def test_url_with_format_param_openai(model, sync_mode): } ], } - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object( + client.chat.completions.with_raw_response, "create" + ) as mock_client: try: if sync_mode: response = completion(**args, client=client) @@ -357,7 +362,9 @@ def test_strip_input_examples_for_non_anthropic_providers(): } ] - assert not litellm_main._should_allow_input_examples(custom_llm_provider="openai", model="gpt-4o-mini") + assert not litellm_main._should_allow_input_examples( + custom_llm_provider="openai", model="gpt-4o-mini" + ) cleaned = litellm_main._drop_input_examples_from_tools(tools=tools) @@ -369,7 +376,9 @@ def test_strip_input_examples_for_non_anthropic_providers(): def test_custom_provider_with_extra_headers(): from litellm.llms.custom_httpx.http_handler import HTTPHandler - with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post: + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: response = litellm.completion( model="custom/custom", messages=[{"role": "user", "content": "Hello, how are you?"}], @@ -384,7 +393,9 @@ def test_custom_provider_with_extra_headers(): def test_custom_provider_with_extra_body(): from litellm.llms.custom_httpx.http_handler import HTTPHandler - with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post: + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: response = litellm.completion( model="custom/custom", messages=[{"role": "user", "content": "Hello, how are you?"}], @@ -411,7 +422,9 @@ def test_custom_provider_with_extra_body(): } # test that extra_body is not passed if not provided - with patch.object(litellm.llms.custom_httpx.http_handler.HTTPHandler, "post") as mock_post: + with patch.object( + litellm.llms.custom_httpx.http_handler.HTTPHandler, "post" + ) as mock_post: response = litellm.completion( model="custom/custom", messages=[{"role": "user", "content": "Hello, how are you?"}], @@ -442,7 +455,9 @@ def set_openrouter_api_key(): @pytest.mark.asyncio -async def test_extra_body_with_fallback(respx_mock: respx.MockRouter, set_openrouter_api_key, monkeypatch): +async def test_extra_body_with_fallback( + respx_mock: respx.MockRouter, set_openrouter_api_key, monkeypatch +): """ test regression for https://github.com/BerriAI/litellm/issues/8425. @@ -510,7 +525,9 @@ async def test_extra_body_with_fallback(respx_mock: respx.MockRouter, set_openro # Verify the response assert response is not None - assert len(respx_mock.calls) > 0, "Mock was not called - check if aiohttp transport is properly disabled" + assert ( + len(respx_mock.calls) > 0 + ), "Mock was not called - check if aiohttp transport is properly disabled" # Get the request from the mock request: httpx.Request = respx_mock.calls[0].request @@ -534,7 +551,9 @@ async def test_extra_body_with_fallback(respx_mock: respx.MockRouter, set_openro @pytest.mark.parametrize("env_base", ["OPENAI_BASE_URL", "OPENAI_API_BASE"]) @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) -async def test_openai_env_base(respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch): +async def test_openai_env_base( + respx_mock: respx.MockRouter, env_base, openai_api_response, monkeypatch +): "This tests OpenAI env variables are honored, including legacy OPENAI_API_BASE" # Ensure aiohttp transport is disabled to use httpx which respx can mock litellm.disable_aiohttp_transport = True @@ -549,7 +568,9 @@ async def test_openai_env_base(respx_mock: respx.MockRouter, env_base, openai_ap messages = [{"role": "user", "content": "Hello, how are you?"}] # Configure respx mock to intercept the request - mock_route = respx_mock.post(url__regex=r"http://localhost:12345/v1/chat/completions.*").mock( + mock_route = respx_mock.post( + url__regex=r"http://localhost:12345/v1/chat/completions.*" + ).mock( return_value=httpx.Response( status_code=200, json={ @@ -583,7 +604,9 @@ async def test_openai_env_base(respx_mock: respx.MockRouter, env_base, openai_ap assert response.choices[0].message.content == "Hello from mocked response!" # Verify the mock was called - assert mock_route.called, "Mock route was not called - request may have bypassed respx" + assert ( + mock_route.called + ), "Mock route was not called - request may have bypassed respx" finally: # Clean up to avoid affecting other tests litellm.disable_aiohttp_transport = False @@ -659,7 +682,9 @@ def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRout model = "gpt-5.2" messages = [{"role": "user", "content": "hi"}] - respx_mock.post("https://api.openai.com/v1/chat/completions").mock(return_value=_mocked_openai_chat_response(model)) + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_mocked_openai_chat_response(model) + ) request = return_raw_request( endpoint=CallTypes.completion, @@ -676,7 +701,9 @@ def test_completion_forwards_verbosity_in_raw_request(respx_mock: respx.MockRout @pytest.mark.asyncio -async def test_acompletion_forwards_verbosity_to_provider_request(respx_mock: respx.MockRouter, monkeypatch): +async def test_acompletion_forwards_verbosity_to_provider_request( + respx_mock: respx.MockRouter, monkeypatch +): """Regression test: acompletion() must forward the verbosity param to the provider request body.""" original_disable_aiohttp = litellm.disable_aiohttp_transport try: @@ -737,9 +764,9 @@ def test_responses_api_bridge_check_gpt_5_4_pro(): model=model_name, custom_llm_provider="openai", ) - assert model_info.get("mode") == "responses", ( - f"{model_name} should have mode='responses', got '{model_info.get('mode')}'" - ) + assert ( + model_info.get("mode") == "responses" + ), f"{model_name} should have mode='responses', got '{model_info.get('mode')}'" def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_responses(): @@ -1156,7 +1183,7 @@ def test_responses_api_bridge_check_openai_backed_custom_api_base_with_unset_eff tools=[{"type": "function", "function": {"name": "get_capital"}}], reasoning_effort=None, api_base=api_base, - ) + ) assert model == "gpt-5.6" assert model_info.get("mode") == "responses" @@ -1181,7 +1208,7 @@ def test_responses_api_bridge_check_lookalike_custom_api_base_with_unset_effort_ tools=[{"type": "function", "function": {"name": "get_capital"}}], reasoning_effort=None, api_base=api_base, - ) + ) assert model == "gpt-5.6" assert model_info.get("mode") != "responses" @@ -1201,7 +1228,7 @@ def test_responses_api_bridge_check_privatelink_api_base_via_env_with_unset_effo tools=[{"type": "function", "function": {"name": "get_capital"}}], reasoning_effort=None, api_base=None, - ) + ) assert model == "gpt-5.6" assert model_info.get("mode") == "responses" @@ -1562,7 +1589,9 @@ def test_responses_api_bridge_check_handles_exception(): with patch("litellm.main._get_model_info_helper") as mock_get_model_info: mock_get_model_info.side_effect = Exception("Model not found") - model_info, model = responses_api_bridge_check(model="responses/custom-model", custom_llm_provider="custom") + model_info, model = responses_api_bridge_check( + model="responses/custom-model", custom_llm_provider="custom" + ) assert model == "custom-model" assert model_info["mode"] == "responses" @@ -2343,7 +2372,9 @@ def test_image_edit_merges_headers_and_extra_headers(): mock_image_edit_config = MagicMock() mock_image_edit_config.get_supported_openai_params.return_value = set() - mock_image_edit_config.map_openai_params.side_effect = lambda **kwargs: dict(kwargs["image_edit_optional_params"]) + mock_image_edit_config.map_openai_params.side_effect = lambda **kwargs: dict( + kwargs["image_edit_optional_params"] + ) with ( patch( @@ -2699,7 +2730,10 @@ def test_mock_completion_stream_with_model_response(): # Verify the content is streamed correctly accumulated_content = "" for chunk in chunks: - if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): accumulated_content += chunk.choices[0].delta.content assert "This is a test response" in accumulated_content or len(chunks) > 0 @@ -2757,7 +2791,10 @@ async def test_async_mock_completion_stream_with_model_response(): # Verify the content is streamed correctly accumulated_content = "" for chunk in chunks: - if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): accumulated_content += chunk.choices[0].delta.content assert "This is an async test response" in accumulated_content or len(chunks) > 0 @@ -2824,7 +2861,9 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): ), ] - response = stream_chunk_builder_text_completion(chunks=chunks, messages=[{"role": "user", "content": "say hello"}]) + response = stream_chunk_builder_text_completion( + chunks=chunks, messages=[{"role": "user", "content": "say hello"}] + ) assert response.choices[0].text == "Hello world" assert response.choices[0].finish_reason == "stop" @@ -3272,7 +3311,10 @@ def _text_chunk(content, finish_reason=None, usage=None): def _priced_at(prompt_tokens, completion_tokens): prices = litellm.model_cost[STREAM_COST_MODEL] - return prompt_tokens * prices["input_cost_per_token"] + completion_tokens * prices["output_cost_per_token"] + return ( + prompt_tokens * prices["input_cost_per_token"] + + completion_tokens * prices["output_cost_per_token"] + ) @pytest.fixture @@ -3339,9 +3381,9 @@ def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map usage=STREAMED_USAGE, ) - assert litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) == pytest.approx( - litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL) - ) + assert litellm.completion_cost( + completion_response=rebuilt, model=STREAM_COST_MODEL + ) == pytest.approx(litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL)) def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): @@ -3360,7 +3402,9 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) assert cost > 0 - assert cost == pytest.approx(_priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens)) + assert cost == pytest.approx( + _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) + ) @pytest.mark.asyncio @@ -3813,8 +3857,6 @@ def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_ def test_aiohttp_openai_warns_only_when_http2_enabled( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, http2_on: bool ): - import logging - from litellm.main import base_llm_aiohttp_handler monkeypatch.setattr(litellm, "http2", http2_on)