Test email body
" + # Test data + from_email = "test@example.com" + to_email = ["recipient@example.com"] + subject = "Test Subject" + html_body = "Test email body
" - # Create mock HTTP client and inject it directly into the logger - # This ensures the mock is used regardless of any caching issues - mock_response = mock.Mock(spec=Response) - mock_response.raise_for_status.return_value = None - mock_response.status_code = 200 - mock_response.json.return_value = {"id": "test_email_id"} + # Create mock HTTP client and inject it directly into the logger + # This ensures the mock is used regardless of any caching issues + mock_response = mock.Mock(spec=Response) + mock_response.raise_for_status.return_value = None + mock_response.status_code = 200 + mock_response.json.return_value = {"id": "test_email_id"} - mock_async_client = mock.AsyncMock() - mock_async_client.post.return_value = mock_response + mock_async_client = mock.AsyncMock() + mock_async_client.post.return_value = mock_response - # Directly inject the mock client to bypass any caching - logger.async_httpx_client = mock_async_client + # Directly inject the mock client to bypass any caching + logger.async_httpx_client = mock_async_client - # Send email - await logger.send_email( - from_email=from_email, - to_email=to_email, - subject=subject, - html_body=html_body, - ) + # Send email + await logger.send_email( + from_email=from_email, + to_email=to_email, + subject=subject, + html_body=html_body, + ) - # Verify the HTTP client was called with None as the API key - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - assert call_args[1]["headers"] == {"Authorization": "Bearer None"} - finally: - # Restore the original key if it existed - if original_key is not None: - os.environ["RESEND_API_KEY"] = original_key + # Verify the HTTP client was called with None as the API key + mock_async_client.post.assert_called_once() + call_args = mock_async_client.post.call_args + assert call_args[1]["headers"] == {"Authorization": "Bearer None"} @pytest.mark.asyncio diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index 8aceee814dcf..7903e9323d3a 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -98,22 +98,18 @@ async def test_send_email_success(mock_env_vars, mock_async_client): @pytest.mark.asyncio -async def test_send_email_missing_api_key(): - original_key = os.environ.pop("SENDGRID_API_KEY", None) +async def test_send_email_missing_api_key(monkeypatch): + monkeypatch.delenv("SENDGRID_API_KEY", raising=False) - try: - logger = SendGridEmailLogger() - - with pytest.raises(ValueError, match='SENDGRID_API_KEY is not set'): - await logger.send_email( - from_email="test@example.com", - to_email=["recipient@example.com"], - subject="Test Subject", - html_body="Test email body
", - ) - finally: - if original_key is not None: - os.environ["SENDGRID_API_KEY"] = original_key + logger = SendGridEmailLogger() + + with pytest.raises(ValueError, match='SENDGRID_API_KEY is not set'): + await logger.send_email( + from_email="test@example.com", + to_email=["recipient@example.com"], + subject="Test Subject", + html_body="Test email body
", + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index 90ef89699eaa..2ad42db8132e 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -16,11 +16,9 @@ 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -import json import os import sys -import pytest import litellm diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py b/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py index 8fc3eca9adf2..8b4e6ab1bab9 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py @@ -220,17 +220,13 @@ def test_stream_transformation_error_handling(): # Create a wrapper mock_wrapper = GoogleGenAIStreamWrapper(completion_stream=iter([])) - # Try to transform - this should handle errors gracefully - try: - streaming_chunk = adapter.translate_streaming_completion_to_generate_content( + # An empty `choices` leaves nothing to emit, so the adapter drops the chunk + assert ( + adapter.translate_streaming_completion_to_generate_content( mock_response, mock_wrapper ) - # If no exception is raised, that's fine - we just want to ensure no crash - assert True - except Exception as e: - # If an exception is raised, it should be a ValueError with appropriate message - assert isinstance(e, ValueError) - # We won't check the exact message as it might vary + is None + ) def test_non_stream_response_when_stream_requested(): diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/test_litellm/google_genai/test_google_genai_main.py index 8f56b4e4bc03..8441b62e559b 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -13,11 +13,9 @@ 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -import json import os import sys -import pytest import litellm diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py index dd97de24df34..4a15da87c89f 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -1,5 +1,6 @@ import json import os +import re import sys from unittest.mock import MagicMock, patch @@ -158,7 +159,7 @@ def test_bitbucket_client_get_file_content_access_denied(mock_get): client = BitBucketClient(config) - with pytest.raises(Exception, match="Access denied to file 'test.prompt'"): + with pytest.raises(Exception, match=re.escape("Access denied to file 'test.prompt'")): client.get_file_content("test.prompt") diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py index cb786d9c2922..1a50a6991da4 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py @@ -1,4 +1,3 @@ -import os import time from unittest.mock import AsyncMock @@ -12,34 +11,13 @@ @pytest.fixture -def clean_env(): - # Save original env - original_api_key = os.environ.get("DD_API_KEY") - original_app_key = os.environ.get("DD_APP_KEY") - original_site = os.environ.get("DD_SITE") - - # Set test env - os.environ["DD_API_KEY"] = "test_api_key" - os.environ["DD_APP_KEY"] = "test_app_key" - os.environ["DD_SITE"] = "test.datadoghq.com" - - yield - - # Restore original env - if original_api_key: - os.environ["DD_API_KEY"] = original_api_key - else: - del os.environ["DD_API_KEY"] - - if original_app_key: - os.environ["DD_APP_KEY"] = original_app_key - else: - del os.environ["DD_APP_KEY"] - - if original_site: - os.environ["DD_SITE"] = original_site - else: - del os.environ["DD_SITE"] +def clean_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key, value in ( + ("DD_API_KEY", "test_api_key"), + ("DD_APP_KEY", "test_app_key"), + ("DD_SITE", "test.datadoghq.com"), + ): + monkeypatch.setenv(key, value) @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py index a4a4ca334b0e..eade92d66727 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py @@ -1,4 +1,3 @@ -import os import time from datetime import datetime, timedelta from unittest.mock import AsyncMock @@ -11,25 +10,16 @@ @pytest.fixture -def clean_env(): - """Set test env vars and restore originals after test.""" - keys = ["DD_API_KEY", "DD_APP_KEY", "DD_SITE", "DD_ENV", "DD_SERVICE", "DD_VERSION"] - originals = {k: os.environ.get(k) for k in keys} - - os.environ["DD_API_KEY"] = "test_api_key" - os.environ["DD_APP_KEY"] = "test_app_key" - os.environ["DD_SITE"] = "test.datadoghq.com" - os.environ["DD_ENV"] = "test-env" - os.environ["DD_SERVICE"] = "test-service" - os.environ["DD_VERSION"] = "1.0.0" - - yield - - for k, v in originals.items(): - if v is not None: - os.environ[k] = v - elif k in os.environ: - del os.environ[k] +def clean_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key, value in ( + ("DD_API_KEY", "test_api_key"), + ("DD_APP_KEY", "test_app_key"), + ("DD_SITE", "test.datadoghq.com"), + ("DD_ENV", "test-env"), + ("DD_SERVICE", "test-service"), + ("DD_VERSION", "1.0.0"), + ): + monkeypatch.setenv(key, value) @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index a4e16500aee5..8d662311da16 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -8,10 +8,10 @@ class TestGCSBucketBase: - def test_construct_request_headers_with_project_id(self): + def test_construct_request_headers_with_project_id(self, monkeypatch): """Test that construct_request_headers correctly uses project_id if passed from env""" test_project_id = "test-project" - os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = test_project_id + monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", test_project_id) try: # Create handler diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py b/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py index 1f7706882f6d..adccd94141f5 100644 --- a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py @@ -1,4 +1,5 @@ import os +import re import sys from unittest.mock import MagicMock, patch @@ -172,7 +173,7 @@ def test_gitlab_client_get_file_content_access_denied(mock_get): mock_get.side_effect = err client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) - with pytest.raises(Exception, match="Access denied to file 'test.prompt'"): + with pytest.raises(Exception, match=re.escape("Access denied to file 'test.prompt'")): client.get_file_content("test.prompt") diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/test_litellm/integrations/test_galileo.py index 0533b7ca7d18..8905795bbc61 100644 --- a/tests/test_litellm/integrations/test_galileo.py +++ b/tests/test_litellm/integrations/test_galileo.py @@ -112,7 +112,6 @@ def test_galileo_input_text_from_messages(): def test_galileo_get_output_str_responses_api(galileo_v2_env): - from litellm.types.llms.openai import ResponsesAPIResponse logger = GalileoObserve() resp_dict = { diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 3c7dd51bff8a..73a62e5594d4 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -14,7 +14,6 @@ from litellm.integrations.langfuse.langfuse import LangFuseLogger sys.path.insert(0, os.path.abspath("../..")) -from litellm.integrations.langfuse.langfuse import LangFuseLogger # Import LangfuseUsageDetails directly from the module where it's defined from litellm.types.integrations.langfuse import * diff --git a/tests/test_litellm/integrations/test_openmeter.py b/tests/test_litellm/integrations/test_openmeter.py index 539e3f99cdc2..2d09e1572db1 100644 --- a/tests/test_litellm/integrations/test_openmeter.py +++ b/tests/test_litellm/integrations/test_openmeter.py @@ -33,7 +33,7 @@ def test_openmeter_logger_initialization(self): def test_openmeter_logger_missing_api_key(self): """Test that OpenMeterLogger raises exception when API key is missing""" os.environ.pop("OPENMETER_API_KEY", None) - with pytest.raises(Exception, match="Missing keys.*OPENMETER_API_KEY"): + with pytest.raises(Exception, match=r"Missing keys.*OPENMETER_API_KEY"): OpenMeterLogger() def test_common_logic_with_string_user(self): @@ -236,9 +236,9 @@ def test_cloudevents_structure(self): assert result["data"]["completion_tokens"] == 8 assert result["data"]["total_tokens"] == 23 - def test_custom_event_type(self): + def test_custom_event_type(self, monkeypatch): """Test that custom event type is used when set""" - os.environ["OPENMETER_EVENT_TYPE"] = "custom_event_type" + monkeypatch.setenv("OPENMETER_EVENT_TYPE", "custom_event_type") logger = OpenMeterLogger() @@ -374,10 +374,10 @@ def test_common_logic_integer_token_user_id(self): assert isinstance(result["subject"], str) assert result["subject"] == "12345" - def test_common_logic_trust_request_user_false_ignores_request_user(self): + def test_common_logic_trust_request_user_false_ignores_request_user(self, monkeypatch): """OPENMETER_TRUST_REQUEST_USER=false makes the key-bound user_id win over a request-supplied `user` (forge-attribution mitigation).""" - os.environ["OPENMETER_TRUST_REQUEST_USER"] = "false" + monkeypatch.setenv("OPENMETER_TRUST_REQUEST_USER", "false") logger = OpenMeterLogger() kwargs = { @@ -400,11 +400,11 @@ def test_common_logic_trust_request_user_false_ignores_request_user(self): assert result["subject"] == "real-tenant-id" assert result["subject"] != "forged-by-client" - def test_common_logic_trust_request_user_false_still_raises_without_key_user(self): + def test_common_logic_trust_request_user_false_still_raises_without_key_user(self, monkeypatch): """OPENMETER_TRUST_REQUEST_USER=false still raises when no user_api_key_user_id is available — the request `user` is not a fallback in this mode.""" - os.environ["OPENMETER_TRUST_REQUEST_USER"] = "false" + monkeypatch.setenv("OPENMETER_TRUST_REQUEST_USER", "false") logger = OpenMeterLogger() kwargs = { diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 8cccfd937e7d..933e41d17a06 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -751,7 +751,7 @@ async def test_strip_base64_mixed_nested_objects(): @pytest.mark.asyncio -async def test_s3_verify_false_handling(): +async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): """ Test that s3_verify=False is properly handled and not treated as None. @@ -763,15 +763,19 @@ async def test_s3_verify_false_handling(): import litellm # Set up s3_callback_params with s3_verify=False - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, # This should NOT be ignored - "s3_use_ssl": False, # This should also NOT be ignored - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, # This should NOT be ignored + "s3_use_ssl": False, # This should also NOT be ignored + }, + ) with patch("asyncio.create_task"): with patch( @@ -801,12 +805,9 @@ async def test_s3_verify_false_handling(): "ssl_verify": False }, f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio -async def test_s3_verify_none_handling(): +async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): """ Test that s3_verify=None uses default behavior. """ @@ -815,12 +816,16 @@ async def test_s3_verify_none_handling(): import litellm # Set up s3_callback_params without s3_verify - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_aws_access_key_id": "test-key", - "s3_aws_secret_access_key": "test-secret", - "s3_region_name": "us-east-1", - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_aws_access_key_id": "test-key", + "s3_aws_secret_access_key": "test-secret", + "s3_region_name": "us-east-1", + }, + ) with patch("asyncio.create_task"): with patch( @@ -846,12 +851,9 @@ async def test_s3_verify_none_handling(): assert call_kwargs["params"].get("ssl_verify") is None # Either params is None or params={'ssl_verify': None} is acceptable - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio -async def test_s3_verify_false_creates_httpx_client_with_verify_false(): +async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatch: pytest.MonkeyPatch): """ Test that when s3_verify=False, the actual httpx client has verify=False. @@ -862,14 +864,18 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(): import litellm # Set up s3_callback_params with s3_verify=False - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, + ) with patch("asyncio.create_task"): # Create logger - this creates the httpx client @@ -888,12 +894,9 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(): httpx_client._verify is False ), f"Expected httpx client _verify=False, got {httpx_client._verify}" - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio -async def test_s3_verify_false_async_client(): +async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): """ Test that the async httpx client respects s3_verify=False. """ @@ -903,14 +906,18 @@ async def test_s3_verify_false_async_client(): from litellm.types.integrations.s3_v2 import s3BatchLoggingElement # Set up s3_callback_params with s3_verify=False - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, + ) with patch("asyncio.create_task"): logger = S3Logger() @@ -945,9 +952,6 @@ async def test_s3_verify_false_async_client(): httpx_client._verify is False ), f"Expected async httpx client _verify=False, got {httpx_client._verify}" - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio async def test_strip_base64_recursive_redaction(): @@ -1169,26 +1173,22 @@ def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # -------------------------------------------------------------- -def test_s3_callback_params_override_uses_alternate_dict(): +def test_s3_callback_params_override_uses_alternate_dict(monkeypatch): """`s3_callback_params_override` makes the logger read its config from the override dict instead of `litellm.s3_callback_params`.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"} - try: - logger = S3Logger( - s3_callback_params_override={ - "s3_bucket_name": "audit-bucket", - "s3_path": "audit-prefix", - "s3_region_name": "us-west-2", - } - ) - assert logger.s3_bucket_name == "audit-bucket" - assert logger.s3_path == "audit-prefix" - assert logger.s3_region_name == "us-west-2" - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"}) + logger = S3Logger( + s3_callback_params_override={ + "s3_bucket_name": "audit-bucket", + "s3_path": "audit-prefix", + "s3_region_name": "us-west-2", + } + ) + assert logger.s3_bucket_name == "audit-bucket" + assert logger.s3_path == "audit-prefix" + assert logger.s3_region_name == "us-west-2" def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): @@ -1198,43 +1198,31 @@ def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): monkeypatch.setenv("MY_AUDIT_BUCKET", "resolved-bucket") override = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"} - original_global = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"} - try: - logger = S3Logger(s3_callback_params_override=override) - assert logger.s3_bucket_name == "resolved-bucket" - assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - assert ( - litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - ) - finally: - litellm.s3_callback_params = original_global + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"}) + logger = S3Logger(s3_callback_params_override=override) + assert logger.s3_bucket_name == "resolved-bucket" + assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" + assert ( + litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" + ) -def test_s3_callback_params_override_none_falls_back_to_global(): +def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch): """No override → behaves exactly as today (reads `litellm.s3_callback_params`).""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "from-global"} - try: - logger = S3Logger() - assert logger.s3_bucket_name == "from-global" - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"}) + logger = S3Logger() + assert logger.s3_bucket_name == "from-global" -def test_s3_callback_params_override_empty_dict_is_opt_in(): +def test_s3_callback_params_override_empty_dict_is_opt_in(monkeypatch): """An empty override dict skips the global entirely (env/IAM-only config).""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "from-global"} - try: - logger = S3Logger(s3_callback_params_override={}) - assert logger.s3_bucket_name is None - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"}) + logger = S3Logger(s3_callback_params_override={}) + assert logger.s3_bucket_name is None def _expected_content_md5(payload: dict) -> str: @@ -1374,20 +1362,20 @@ async def test_async_upload_sets_server_side_encryption_header_when_configured() assert headers["x-amz-server-side-encryption"] == "aws:kms" -def test_s3_server_side_encryption_read_from_callback_params(): +def test_s3_server_side_encryption_read_from_callback_params(monkeypatch): """s3_server_side_encryption can be configured via s3_callback_params.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "aws:kms" - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" @pytest.mark.asyncio @@ -1505,21 +1493,21 @@ async def test_async_upload_omits_kms_key_id_header_when_not_configured(): assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers -def test_s3_sse_kms_key_id_read_from_callback_params(): +def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch): """s3_sse_kms_key_id can be configured via s3_callback_params.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - } - try: - logger = S3Logger() - assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, + ) + logger = S3Logger() + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") @pytest.mark.asyncio @@ -1561,83 +1549,79 @@ async def test_async_upload_infers_aws_kms_when_only_key_id_set(): ) -def test_s3_sse_kms_key_id_read_from_audit_override_params(): +def test_s3_sse_kms_key_id_read_from_audit_override_params(monkeypatch): """The audit-log override path must honor s3_sse_kms_key_id too.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "normal-logs-bucket"} - try: - logger = S3Logger( - s3_callback_params_override={ - "s3_bucket_name": "audit-logs-bucket", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id", - } - ) - assert logger.s3_bucket_name == "audit-logs-bucket" - assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id") - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-logs-bucket"}) + logger = S3Logger( + s3_callback_params_override={ + "s3_bucket_name": "audit-logs-bucket", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id", + } + ) + assert logger.s3_bucket_name == "audit-logs-bucket" + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id") -def test_kms_key_id_dropped_when_algorithm_is_not_kms(): +def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch): """ AES256 plus a KMS key id is an invalid S3 combination; the key id must be dropped at init so uploads keep working instead of silently 400ing. """ import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "AES256", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "AES256" - assert logger.s3_sse_kms_key_id is None - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "AES256" + assert logger.s3_sse_kms_key_id is None -def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(): +def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch): """ A YAML boolean in s3_server_side_encryption must not crash logger init and must not discard the valid key id; aws:kms is inferred from the key id. """ import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": True, - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "aws:kms" - assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") -def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(): +def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch): """A mistyped key id (unquoted YAML number) must not disable the valid algorithm.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": 12345, - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "aws:kms" - assert logger.s3_sse_kms_key_id is None - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" + assert logger.s3_sse_kms_key_id is None _ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE" diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py index 524589abf5e2..05b0bde16bb0 100644 --- a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py +++ b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py @@ -86,29 +86,21 @@ def test_preserves_existing_headers(self, config): assert headers["X-Custom"] == "value" assert headers["x-goog-api-key"] == "test-key" - def test_api_revision_new_schema_by_default(self, config): + def test_api_revision_new_schema_by_default(self, config, monkeypatch: pytest.MonkeyPatch): # Default: use_legacy_interactions_schema=False → new steps schema - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - headers = config.validate_environment( - headers={}, model="gemini-2.5-flash", litellm_params=None - ) - assert headers["Api-Revision"] == "2026-05-20" - finally: - litellm.use_legacy_interactions_schema = original + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + headers = config.validate_environment( + headers={}, model="gemini-2.5-flash", litellm_params=None + ) + assert headers["Api-Revision"] == "2026-05-20" - def test_api_revision_legacy_schema_when_flag_set(self, config): + def test_api_revision_legacy_schema_when_flag_set(self, config, monkeypatch: pytest.MonkeyPatch): # Flag on → legacy outputs schema until June 8, 2026 - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = True - headers = config.validate_environment( - headers={}, model="gemini-2.5-flash", litellm_params=None - ) - assert headers["Api-Revision"] == "2026-05-07" - finally: - litellm.use_legacy_interactions_schema = original + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", True) + headers = config.validate_environment( + headers={}, model="gemini-2.5-flash", litellm_params=None + ) + assert headers["Api-Revision"] == "2026-05-07" class TestGetCompleteUrl: @@ -561,23 +553,19 @@ def test_get_interaction_raises_without_key(self, config): class TestTransformRequestSchemaCoalescing: """Test new-schema request coalescing (Api-Revision: 2026-05-20).""" - def test_response_mime_type_folded_into_response_format(self, config): - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="summarise", - optional_params={ - "response_mime_type": "application/json", - "response_format": {"type": "object", "properties": {}}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + def test_response_mime_type_folded_into_response_format(self, config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="summarise", + optional_params={ + "response_mime_type": "application/json", + "response_format": {"type": "object", "properties": {}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) # response_mime_type must not appear as a top-level body key assert "response_mime_type" not in body @@ -586,25 +574,21 @@ def test_response_mime_type_folded_into_response_format(self, config): assert rf["mime_type"] == "application/json" assert "schema" in rf - def test_image_config_moved_to_response_format(self, config): - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="draw a sunset", - optional_params={ - "generation_config": { - "temperature": 0.7, - "image_config": {"aspect_ratio": "1:1", "image_size": "1K"}, - } - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + def test_image_config_moved_to_response_format(self, config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw a sunset", + optional_params={ + "generation_config": { + "temperature": 0.7, + "image_config": {"aspect_ratio": "1:1", "image_size": "1K"}, + } + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) # image_config removed from generation_config assert "image_config" not in body.get("generation_config", {}) @@ -613,95 +597,85 @@ def test_image_config_moved_to_response_format(self, config): assert rf["type"] == "image" assert rf["aspect_ratio"] == "1:1" - def test_response_mime_type_skipped_when_response_format_is_list(self, config): + def test_response_mime_type_skipped_when_response_format_is_list(self, config, monkeypatch: pytest.MonkeyPatch): """Lists are already polymorphic; do not wrap them into schema.""" - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - rf_list = [ - {"type": "text", "mime_type": "application/json"}, - {"type": "image", "aspect_ratio": "1:1"}, - ] - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="multimodal", - optional_params={ - "response_format": rf_list, - "response_mime_type": "application/json", - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + rf_list = [ + {"type": "text", "mime_type": "application/json"}, + {"type": "image", "aspect_ratio": "1:1"}, + ] + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="multimodal", + optional_params={ + "response_format": rf_list, + "response_mime_type": "application/json", + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert body["response_format"] == rf_list assert "response_mime_type" not in body def test_image_config_appended_to_response_format_list_without_mutating_input( - self, config + self, + config, + monkeypatch: pytest.MonkeyPatch, ): """When response_format is already a list, image_config must not mutate optional_params.""" - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - text_rf = {"type": "text", "mime_type": "application/json"} - optional_params = { - "response_format": [text_rf], - "generation_config": { - "image_config": {"aspect_ratio": "16:9", "image_size": "2K"}, - }, - } - original_rf = optional_params["response_format"] - - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="draw and summarise", - optional_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + text_rf = {"type": "text", "mime_type": "application/json"} + optional_params = { + "response_format": [text_rf], + "generation_config": { + "image_config": {"aspect_ratio": "16:9", "image_size": "2K"}, + }, + } + original_rf = optional_params["response_format"] - assert optional_params["response_format"] is original_rf - assert len(optional_params["response_format"]) == 1 - assert body["response_format"] == [ - text_rf, - {"type": "image", "aspect_ratio": "16:9", "image_size": "2K"}, - ] - - # Retry must not append a second image entry into the caller's list. - body_retry = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="draw and summarise", - optional_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - assert len(optional_params["response_format"]) == 1 - assert body_retry["response_format"] == body["response_format"] - finally: - litellm.use_legacy_interactions_schema = original + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw and summarise", + optional_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) - def test_legacy_schema_passes_fields_unchanged(self, config): - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = True - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="hello", - optional_params={ - "response_mime_type": "application/json", - "generation_config": {"image_config": {"aspect_ratio": "16:9"}}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + assert optional_params["response_format"] is original_rf + assert len(optional_params["response_format"]) == 1 + assert body["response_format"] == [ + text_rf, + {"type": "image", "aspect_ratio": "16:9", "image_size": "2K"}, + ] + + # Retry must not append a second image entry into the caller's list. + body_retry = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw and summarise", + optional_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert len(optional_params["response_format"]) == 1 + assert body_retry["response_format"] == body["response_format"] + + def test_legacy_schema_passes_fields_unchanged(self, config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", True) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="hello", + optional_params={ + "response_mime_type": "application/json", + "generation_config": {"image_config": {"aspect_ratio": "16:9"}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert body["response_mime_type"] == "application/json" assert body["generation_config"]["image_config"]["aspect_ratio"] == "16:9" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index cf36a2b9b253..052c08a86b58 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -56,8 +56,8 @@ def test_bedrock_guardrail_cost_no_pricing_entry(monkeypatch): assert bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") == 0.0 -def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") assert litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] == { "automatedReasoningPolicyUnits": 0.00017, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index f66056a54e24..c8c360327930 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1,6 +1,4 @@ import json -import os -import sys import pytest from fastapi.testclient import TestClient @@ -28,10 +26,6 @@ StandardBuiltInToolsParams, ) -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - from litellm.litellm_core_utils.llm_cost_calc.utils import ( PromptTokensDetailsResult, TokenTypeCostBreakdown, @@ -44,13 +38,17 @@ from litellm.types.utils import CacheCreationTokenDetails, Usage -def test_reasoning_tokens_no_price_set(): +@pytest.fixture +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + +def test_reasoning_tokens_no_price_set(_local_model_cost_map): # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token) model = "o1" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] usage = Usage( completion_tokens=1578, @@ -87,11 +85,9 @@ def test_reasoning_tokens_no_price_set(): ) -def test_reasoning_tokens_gemini(): +def test_reasoning_tokens_gemini(_local_model_cost_map): model = "gemini-2.5-flash" custom_llm_provider = "gemini" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( completion_tokens=1578, @@ -132,12 +128,10 @@ def test_reasoning_tokens_gemini(): ) -def test_reasoning_tokens_gemini_3_1_flash_lite(): +def test_reasoning_tokens_gemini_3_1_flash_lite(_local_model_cost_map): """Test cost calculation for gemini-3.1-flash-lite-preview with reasoning tokens""" model = "gemini-3.1-flash-lite-preview" custom_llm_provider = "gemini" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( completion_tokens=1000, @@ -270,11 +264,9 @@ def test_image_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round(expected_completion_cost, 12) -def test_video_output_tokens_gemini_omni_flash_preview(): +def test_video_output_tokens_gemini_omni_flash_preview(_local_model_cost_map): """Video output tokens are billed at output_cost_per_video_token, not the text rate and not zero.""" model = "gemini-omni-flash-preview" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") text_tokens = 100 video_tokens = 46336 @@ -310,11 +302,9 @@ def test_video_output_tokens_gemini_omni_flash_preview(): ) -def test_video_input_tokens_gemini_omni_flash_preview(): +def test_video_input_tokens_gemini_omni_flash_preview(_local_model_cost_map): """Video input tokens are billed at the standard input rate instead of being dropped.""" model = "gemini-omni-flash-preview" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( completion_tokens=10, @@ -369,12 +359,10 @@ def test_video_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12) -def test_generic_cost_per_token_above_200k_tokens(): +def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map): # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing model = "gemini-2.5-pro" custom_llm_provider = "vertex_ai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] prompt_tokens = 220 * 1e6 @@ -420,12 +408,10 @@ def test_get_token_base_cost_picks_highest_crossed_tier(): assert prompt_base_cost == 9e-6 -def test_generic_cost_per_token_gpt54_above_272k_tokens(): +def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] prompt_tokens = 273000 # Above 272K threshold @@ -450,12 +436,10 @@ def test_generic_cost_per_token_gpt54_above_272k_tokens(): assert round(completion_cost, 10) == round(expected_completion, 10) -def test_generic_cost_per_token_minimax_m3_above_512k_tokens(): +def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_map): """MiniMax-M3: prompts >512K input tokens priced at 2x input, output, and cache read.""" model = "minimax/MiniMax-M3" custom_llm_provider = "minimax" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] prompt_tokens = 600000 @@ -493,10 +477,8 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(): "bedrock_mantle/openai.gpt-5.6-luna", ], ) -def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(model): +def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model): """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] assert model_cost_map["max_input_tokens"] == 1000000 @@ -827,12 +809,10 @@ def test_generic_cost_per_token_tiered_pricing_bills_reasoning_at_tier_rate(): litellm.model_cost.pop(model, None) -def test_generic_cost_per_token_gpt55(): +def test_generic_cost_per_token_gpt55(_local_model_cost_map): """gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input.""" model = "gpt-5.5" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] @@ -867,12 +847,10 @@ def test_generic_cost_per_token_gpt55(): ) -def test_generic_cost_per_token_gpt55_pro(): +def test_generic_cost_per_token_gpt55_pro(_local_model_cost_map): """gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input.""" model = "gpt-5.5-pro" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] @@ -913,13 +891,13 @@ def test_generic_cost_per_token_gpt55_pro(): @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost,cache_write_cost", [ - ("gpt-5.6", 5e-6, 3e-5, 5e-7, 6.25e-6), - ("gpt-5.6-sol", 5e-6, 3e-5, 5e-7, 6.25e-6), + ("gpt-5.6", 4e-6, 2e-5, 4e-7, 5e-6), + ("gpt-5.6-sol", 4e-6, 2e-5, 4e-7, 5e-6), ("gpt-5.6-terra", 2e-6, 1.2e-5, 2e-7, 2.5e-6), ("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7), ], ) -def test_generic_cost_per_token_gpt56( +def test_generic_cost_per_token_gpt56(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost, cache_write_cost ): """gpt-5.6 (sol/terra/luna): base pricing + new cache-write cost. @@ -927,8 +905,6 @@ def test_generic_cost_per_token_gpt56( Cache writes are billed at 1.25x the uncached input rate for this family. """ custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] @@ -965,16 +941,31 @@ def test_generic_cost_per_token_gpt56( assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) +def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map): + """Regression: the bare gpt-5.6 alias routes to GPT-5.6 Sol, so every cost field on + the two entries has to hold the same value. They drifted once before, when Sol took + its promotional cut and gpt-5.6 was left on the pre-cut rates, overbilling callers + who used the alias.""" + alias = litellm.model_cost["gpt-5.6"] + sol = litellm.model_cost["gpt-5.6-sol"] + + cost_fields = sorted(field for field in sol if "cost" in field) + assert len(cost_fields) == 23 + + for field in cost_fields: + assert alias.get(field) == sol.get(field), field + + @pytest.mark.parametrize( "model,flex_long_input_cost,flex_long_output_cost", [ - ("gpt-5.6", 5e-6, 2.25e-5), - ("gpt-5.6-sol", 5e-6, 2.25e-5), + ("gpt-5.6", 4e-6, 1.5e-5), + ("gpt-5.6-sol", 4e-6, 1.5e-5), ("gpt-5.6-terra", 2e-6, 9e-6), ("gpt-5.6-luna", 2e-7, 9e-7), ], ) -def test_generic_cost_per_token_gpt56_flex_above_272k( +def test_generic_cost_per_token_gpt56_flex_above_272k(_local_model_cost_map, model, flex_long_input_cost, flex_long_output_cost ): """A >272K flex request bills the flex long-context rate, not the standard one. @@ -983,8 +974,6 @@ def test_generic_cost_per_token_gpt56_flex_above_272k( ``*_above_272k_tokens_flex`` keys these requests silently fell back to the standard long-context price, billing 2x what OpenAI charges. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") prompt_tokens = 300000 completion_tokens = 1000 @@ -1023,11 +1012,9 @@ def test_generic_cost_per_token_gpt56_flex_above_272k( ("flex", 300000, 2e-6, 2.5e-6, 2e-7), ], ) -def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context( +def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context(_local_model_cost_map, service_tier, prompt_tokens, input_rate, cache_write_rate, cache_read_rate ): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") cached_tokens = 50000 cache_write_tokens = 40000 @@ -1115,14 +1102,14 @@ def test_generic_cost_per_token_gpt56_cyber( ("azure/eu/gpt-5.6-luna", 2.2e-7, 1.32e-6, 2.2e-8), ], ) -def test_generic_cost_per_token_azure_gpt56( +def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost ): - """Azure gpt-5.6 (global + us/eu regional): pricing mirrors the openai - family for global deployments and carries the standard 10% regional uplift. + """Azure gpt-5.6 (global + us/eu regional): Azure prices this family on its own + schedule and carries the standard 10% regional uplift on top. It did not take the + promotional cut OpenAI applied to gpt-5.6-sol, so these rates deliberately sit + above the openai ones and must not be lowered to match them. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] assert model_cost_map["litellm_provider"] == "azure" @@ -1163,7 +1150,7 @@ def test_generic_cost_per_token_azure_gpt56( ("gpt-5.5-pro-2026-04-23", False, True, False), ], ) -def test_gpt55_reasoning_effort_flags_match_live_openai_api( +def test_gpt55_reasoning_effort_flags_match_live_openai_api(_local_model_cost_map, model, expected_none, expected_xhigh, expected_minimal ): """Pin reasoning_effort capability flags to OpenAI's actual API contract. @@ -1172,8 +1159,6 @@ def test_gpt55_reasoning_effort_flags_match_live_openai_api( ``Unsupported value: 'reasoning_effort' does not support 'minimal' with this model``. gpt-5.5-pro additionally rejects 'none' and 'low'. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") m = litellm.model_cost[model] assert ( @@ -1194,7 +1179,7 @@ def test_gpt55_reasoning_effort_flags_match_live_openai_api( ("gpt-5.5-pro", "gpt-5.5-pro-2026-04-23"), ], ) -def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( +def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(_local_model_cost_map, base_model, dated_model ): """Dated snapshots must carry the same reasoning_effort capability flags as @@ -1206,8 +1191,6 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( behavior between ``gpt-5.5`` and ``gpt-5.5-2026-04-23``. Pinning to a dated variant must never lose capabilities relative to the base alias. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") base = litellm.model_cost[base_model] dated = litellm.model_cost[dated_model] @@ -1234,7 +1217,7 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( ("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6), ], ) -def test_azure_gpt55_entries_present_with_correct_pricing( +def test_azure_gpt55_entries_present_with_correct_pricing(_local_model_cost_map, model, expected_mode, expected_input, expected_output, expected_cache_read ): """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure. @@ -1243,8 +1226,6 @@ def test_azure_gpt55_entries_present_with_correct_pricing( on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro. Cache discount is 10% of input. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") m = litellm.model_cost[model] assert m["litellm_provider"] == "azure" @@ -1269,12 +1250,10 @@ def test_azure_gpt55_entries_present_with_correct_pricing( ("azure/gpt-5.5-pro", False, False, True), ], ) -def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api( +def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api(_local_model_cost_map, model, expected_none, expected_minimal, expected_xhigh ): """Azure entries pin reasoning_effort flags to OpenAI's actual API contract.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") m = litellm.model_cost[model] assert m.get("supports_none_reasoning_effort") is expected_none @@ -1654,11 +1633,9 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): assert round(result, 6) == round(expected, 6) -def test_service_tier_flex_pricing(): +def test_service_tier_flex_pricing(_local_model_cost_map): """Test that flex service tier uses correct pricing (approximately 50% of standard).""" # Set up environment for local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -1711,11 +1688,9 @@ def test_service_tier_flex_pricing(): ), f"Flex total cost mismatch: {flex_total} vs {expected_flex_total}" -def test_service_tier_default_pricing(): +def test_service_tier_default_pricing(_local_model_cost_map): """Test that when no service tier is provided, standard pricing is used.""" # Set up environment for local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano model = "gpt-5-nano" @@ -1762,11 +1737,9 @@ def test_service_tier_default_pricing(): ), f"Standard completion cost mismatch: {default_cost[1]} vs {expected_standard_completion}" -def test_service_tier_fallback_pricing(): +def test_service_tier_fallback_pricing(_local_model_cost_map): """Test that when service tier is provided but model doesn't have those keys, it falls back to standard pricing.""" # Set up environment for local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-4 which doesn't have flex pricing keys model = "gpt-4" @@ -1874,15 +1847,13 @@ def test_service_tier_ultrafast_pricing(): assert completion_cost == pytest.approx(400 * 3e-04) -def test_service_tier_ultrafast_fallback_pricing(): +def test_service_tier_ultrafast_fallback_pricing(_local_model_cost_map): """Without *_ultrafast keys an ultrafast request bills the standard rate, not zero. Guards the suffix fallback in _get_cost_per_unit: "_fast" is a substring of "_ultrafast", so a shortest-first suffix match would strip the wrong suffix and price the request at 0. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) @@ -1909,9 +1880,10 @@ def test_service_tier_ultrafast_fallback_pricing(): [ "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview", + "gemini-3.1-flash-lite-image", ], ) -def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): +def test_gemini_image_generation_cost_with_zero_text_tokens(_local_model_cost_map, model: str): """ Test that image_tokens are correctly costed when text_tokens=0. @@ -1921,8 +1893,6 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): https://github.com/BerriAI/litellm/issues/17410 """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") custom_llm_provider = "vertex_ai" @@ -1977,13 +1947,11 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): ), f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}" -def test_vertex_image_generation_cost_prefers_token_usage_metadata(): +def test_vertex_image_generation_cost_prefers_token_usage_metadata(_local_model_cost_map): """ When usage metadata exists on image responses, Vertex image generation cost should be calculated from token pricing, not flat output_cost_per_image. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3.1-flash-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") @@ -2022,13 +1990,11 @@ def test_vertex_image_generation_cost_prefers_token_usage_metadata(): assert cost != len(image_response.data) * model_info["output_cost_per_image"] -def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(): +def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(_local_model_cost_map): """ Without usage metadata, Vertex image generation cost should fall back to output_cost_per_image * number_of_images. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3.1-flash-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") @@ -2046,13 +2012,11 @@ def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(): assert round(cost, 10) == round(expected_cost, 10) -def test_gemini_image_generation_cost_prefers_token_usage_metadata(): +def test_gemini_image_generation_cost_prefers_token_usage_metadata(_local_model_cost_map): """ When usage metadata exists on image responses, Gemini image generation cost should be calculated from token pricing, not flat output_cost_per_image. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -2091,13 +2055,11 @@ def test_gemini_image_generation_cost_prefers_token_usage_metadata(): assert cost != len(image_response.data) * model_info["output_cost_per_image"] -def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(): +def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(_local_model_cost_map): """ Without usage metadata, Gemini image generation cost should fall back to output_cost_per_image * number_of_images. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -2194,7 +2156,7 @@ def test_reasoning_tokens_without_text_tokens_gpt5_nano(): ), "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!" -def test_image_count_prevents_text_tokens_fallback(): +def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map): """ Test that the text_tokens fallback in generic_cost_per_token does not override text_tokens=0 when image_count > 0. @@ -2203,8 +2165,6 @@ def test_image_count_prevents_text_tokens_fallback(): When image_count > 0, text_tokens=0 is intentional (image-only request), not "text_tokens not set by provider." """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Simulate Nova image-only embedding: prompt_tokens estimated from # embedding dimensions (768 for 3072-dim), image_count=1 @@ -2238,20 +2198,6 @@ def test_image_count_prevents_text_tokens_fallback(): # --------------------------------------------------------------------------- -@pytest.fixture -def _local_model_cost_map(): - prev_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - prev_model_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - try: - yield - finally: - litellm.model_cost = prev_model_cost - if prev_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = prev_env @pytest.mark.parametrize("model", ["gpt-5.4", "gpt-realtime-2.1", "gpt-realtime-2.1-mini"]) @@ -2585,7 +2531,7 @@ def test_threshold_keys_exclude_service_tier_variants(): ("cerebras/qwen-3-32b", "cerebras", 250, 0), ], ) -def test_token_type_cost_breakdown_is_provider_agnostic( +def test_token_type_cost_breakdown_is_provider_agnostic(_local_model_cost_map, model, custom_llm_provider, reasoning_tokens, cached_tokens ): """ @@ -2597,8 +2543,6 @@ def test_token_type_cost_breakdown_is_provider_agnostic( there - not the top-level cache_read_input_tokens attribute the old breakdown code relied on - is what makes Vertex/OpenAI/Azure cache costs show up at all. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( prompt_tokens=1000, @@ -2629,10 +2573,8 @@ def test_token_type_cost_breakdown_is_provider_agnostic( assert breakdown.cache_read_cost == pytest.approx(cached_tokens * cache_read_rate) -def test_token_type_cost_breakdown_matches_real_gemini_numbers(): +def test_token_type_cost_breakdown_matches_real_gemini_numbers(_local_model_cost_map): """Hard-coded against the exact gemini-2.5-flash response that exposed the gap.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( prompt_tokens=209, @@ -2655,9 +2597,7 @@ def test_token_type_cost_breakdown_matches_real_gemini_numbers(): assert breakdown.cache_creation_cost == 0.0 -def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(_local_model_cost_map): usage = Usage( prompt_tokens=200_000, @@ -2679,9 +2619,7 @@ def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(): assert breakdown.cache_read_cost == pytest.approx(50_000 * 4e-07) -def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(_local_model_cost_map): usage = Usage( prompt_tokens=199_999, @@ -2703,14 +2641,12 @@ def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(): assert breakdown.cache_read_cost == pytest.approx(50_000 * 2e-07) -def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(): +def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(_local_model_cost_map): """ Bedrock/Anthropic report cache tokens as top-level usage fields; the Usage constructor maps them onto prompt_tokens_details, so the breakdown must still pick up both cache-read and cache-creation costs. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "anthropic.claude-3-5-haiku-20241022-v1:0" usage = Usage( @@ -2734,14 +2670,12 @@ def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage( ) -def test_token_type_cost_breakdown_reads_cache_write_tokens(): +def test_token_type_cost_breakdown_reads_cache_write_tokens(_local_model_cost_map): """ Some OpenAI-compatible providers (e.g. kimi-k2) report cache-write tokens under `cache_write_tokens` rather than `cache_creation_tokens`. The breakdown must read it the same way the total-cost normalization does, so the two agree. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "anthropic.claude-3-5-haiku-20241022-v1:0" usage = Usage( @@ -2762,7 +2696,7 @@ def test_token_type_cost_breakdown_reads_cache_write_tokens(): ) -def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): +def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(_local_model_cost_map): """ Regression: OpenAI gpt-5.6 reports cache-write tokens under prompt_tokens_details.cache_write_tokens (not the Anthropic cache_creation_tokens @@ -2770,8 +2704,6 @@ def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): input rate. Customer report: cache creation tokens were never counted for the GPT-5.6 series, so cost was undercounted on cache-write requests. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.6" usage = Usage( @@ -2793,14 +2725,12 @@ def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): assert prompt_cost > 1000 * info["input_cost_per_token"] -def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(): +def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(_local_model_cost_map): """ Regression for #34801: when a provider reports text_tokens covering the whole prompt alongside cache-write tokens (and no cache reads), the cache-write tokens must be backed out of the text total instead of being billed twice. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.6" usage = Usage( @@ -2819,15 +2749,13 @@ def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(): assert prompt_cost == pytest.approx(expected_prompt) -def test_token_type_cost_breakdown_reconciles_with_generic_total(): +def test_token_type_cost_breakdown_reconciles_with_generic_total(_local_model_cost_map): """ Both-ways check: the reasoning subset must sum with the remaining (text) output cost to exactly the completion total, and the cache-read subset with the remaining input cost to exactly the prompt total, as computed by generic_cost_per_token. A mismatch here would mean the breakdown misrepresents what was actually billed. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-2.5-flash" custom_llm_provider = "vertex_ai" @@ -2860,9 +2788,7 @@ def test_token_type_cost_breakdown_reconciles_with_generic_total(): assert text_input_cost + breakdown.cache_read_cost == pytest.approx(prompt_cost) -def test_token_type_cost_breakdown_zero_without_special_tokens(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) breakdown = get_token_type_cost_breakdown( @@ -2899,7 +2825,7 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(): ), ], ) -def test_token_type_cost_breakdown_openai_responses_api_cache_write_read( +def test_token_type_cost_breakdown_openai_responses_api_cache_write_read(_local_model_cost_map, raw_usage, expect_read, expect_write ): """Regression for #34309: OpenAI Responses API reports cache tokens under @@ -2908,8 +2834,6 @@ def test_token_type_cost_breakdown_openai_responses_api_cache_write_read( cache_read_cost / cache_creation_cost from the transformed usage.""" from litellm.responses.utils import ResponseAPILoggingUtils - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.6" usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage) @@ -2950,15 +2874,13 @@ def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): ) -def test_token_type_cost_breakdown_applies_regional_uplift(): +def test_token_type_cost_breakdown_applies_regional_uplift(_local_model_cost_map): """ Regional OpenAI hosts (eu./us.) apply a flat uplift to every token cost. The per-type breakdown must apply the same uplift via data_residency so it stays reconciled with the uplifted input_cost/output_cost totals, instead of being logged at the base rate. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.4" custom_llm_provider = "openai" @@ -3006,15 +2928,13 @@ def test_token_type_cost_breakdown_applies_regional_uplift(): assert text_input_cost + eu.cache_read_cost == pytest.approx(prompt_cost) -def test_token_type_cost_breakdown_applies_vertex_regional_uplift(): +def test_token_type_cost_breakdown_applies_vertex_regional_uplift(_local_model_cost_map): """ Non-global Vertex endpoints apply a flat 1.1x uplift to every token cost. The per-type breakdown must apply the same uplift via vertex_location so it stays reconciled with the uplifted input_cost/output_cost totals, instead of being logged at the global rate. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-haiku-4-5@20251001" custom_llm_provider = "vertex_ai" @@ -3057,7 +2977,7 @@ def test_token_type_cost_breakdown_applies_vertex_regional_uplift(): assert text_input_cost + regional.cache_read_cost == pytest.approx(prompt_cost) -def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch): +def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(_local_model_cost_map, monkeypatch): """ Anthropic's regional (geo) uplift lives in provider_specific_entry and is applied to every token type in the totals, so the per-type breakdown must @@ -3070,7 +2990,6 @@ def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch) ) monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-geo-breakdown-model" litellm.register_model( @@ -3191,9 +3110,7 @@ def test_image_response_input_image_tokens_priced_at_image_rate(details_as_dict) @pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_DAY0_LAUNCH_PRICING) -def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, output_cost, cache_read_cost): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_gemini_36_flash_and_35_flash_lite_launch_pricing(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost): model_cost_map = litellm.model_cost[model] assert model_cost_map["input_cost_per_token"] == input_cost @@ -3206,9 +3123,7 @@ def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, out assert model_cost_map["max_input_tokens"] == 1048576 -def test_generic_cost_per_token_gemini_36_flash(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_generic_cost_per_token_gemini_36_flash(_local_model_cost_map): usage = Usage( prompt_tokens=1000, @@ -3274,9 +3189,7 @@ def test_gemini_36_flash_batch_introductory_pricing(model, _local_model_cost_map assert model_cost_map["output_cost_per_token_batches"] == 1.875e-06 -def test_generic_cost_per_token_gemini_35_flash_lite(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map): usage = Usage( prompt_tokens=1000, @@ -3300,8 +3213,8 @@ def test_generic_cost_per_token_gemini_35_flash_lite(): @pytest.mark.parametrize( "service_tier,input_rate,cache_read_rate,cache_write_rate,output_rate", [ - ("flex", 2.5e-6, 2.5e-7, 3.125e-6, 1.5e-5), - ("priority", 1e-5, 1e-6, 1.25e-5, 6e-5), + ("flex", 2e-6, 2e-7, 2.5e-6, 1e-5), + ("priority", 8e-6, 8e-7, 1e-5, 4e-5), ], ) def test_service_tier_cache_creation_rates_for_gpt_5_6( @@ -3314,7 +3227,7 @@ def test_service_tier_cache_creation_rates_for_gpt_5_6( ): """Regression: gpt-5.6 publishes cache_creation_input_token_cost_flex/_priority, so a flex or priority request must bill cache writes at that tier's rate instead of falling - back to the standard 6.25e-6 rate.""" + back to the standard cache-write rate.""" usage = Usage( prompt_tokens=10_000, completion_tokens=500, @@ -3361,8 +3274,8 @@ def test_fast_service_tier_bills_at_the_priority_rate(_local_model_cost_map): model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast" ) - expected_prompt = 800 * 1e-05 + 200 * 1e-06 - expected_completion = 500 * 6e-05 + expected_prompt = 800 * 8e-06 + 200 * 8e-07 + expected_completion = 500 * 4e-05 assert fast == priority assert fast[0] == pytest.approx(expected_prompt, rel=1e-9) @@ -3397,8 +3310,8 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m ) assert fast == priority - assert fast[0] == pytest.approx(300_000 * 1e-05, rel=1e-9) - assert fast[1] == pytest.approx(1_000 * 4.5e-05, rel=1e-9) + assert fast[0] == pytest.approx(300_000 * 8e-06, rel=1e-9) + assert fast[1] == pytest.approx(1_000 * 3e-05, rel=1e-9) def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 0c945151a90f..a51282287425 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -15,13 +15,7 @@ ) # Adds the parent directory to the system path -@pytest.fixture -def local_model_cost_map(monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - -# Test basic web search cost calculations def test_web_search_cost_low(): web_search_options = WebSearchOptions(search_context_size="low") model_info = litellm.get_model_info("gpt-4o-search-preview") @@ -383,12 +377,12 @@ def test_get_cost_for_vertex_ai_gemini_web_search(model, custom_llm_provider): assert cost == 0.035, f"Expected $0.035 grounding cost, got ${cost}" -def test_azure_assistant_features_integrated_cost_tracking(): +def test_azure_assistant_features_integrated_cost_tracking(monkeypatch): """ Test integrated cost tracking for Azure assistant features. """ # Force use of local model cost map for CI/CD consistency - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure/gpt-4o" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 08d8c17cc2eb..3a7e06d085ab 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -332,7 +332,6 @@ def test_bedrock_get_document_format_fallback_mimes(): This tests the fallback mechanism when mimetypes.guess_all_extensions returns empty results, which can happen in Docker containers where mimetypes depends on OS-installed MIME types. """ - from unittest.mock import patch # Test DOCX fallback docx_mime = ( @@ -2845,7 +2844,7 @@ def test_anthropic_messages_pt_file_block_preserves_cache_control(): assert text_block["cache_control"]["type"] == "ephemeral" -def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(): +def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(monkeypatch): """ Tools with cache_control ttl should preserve the ttl in the cachePoint block for Claude 4.5+ models on Bedrock, matching the behavior of system @@ -2868,7 +2867,7 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: tool_with_1h = { @@ -2928,10 +2927,10 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) -def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): +def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(monkeypatch): """ End-to-end: _bedrock_tools_pt should produce cachePoint blocks with ttl for Claude 4.5+ models when tools have cache_control with ttl. @@ -2945,7 +2944,7 @@ def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: tools = [ @@ -2981,7 +2980,7 @@ def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index d5676aaf288f..38f46b26eea3 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -762,3 +762,50 @@ def test_azure_404_with_invalid_request_error_type_maps_to_not_found(): assert excinfo.value.status_code == 404 assert "Response with id 'resp_abc' not found." in excinfo.value.message + + +def test_bedrock_mantle_400_maps_to_bad_request(): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + original_exception = BaseLLMException( + status_code=400, + message=( + '{"error": {"code": "validation_error", "message": ' + "\"invalid request body: Invalid 'input': value did not match any expected variant\", " + '"type": "invalid_request_error"}}' + ), + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gpt-5.6-terra", + original_exception=original_exception, + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.status_code == 400 + assert "Invalid 'input'" in excinfo.value.message + assert type(excinfo.value) is litellm.BadRequestError + + +def test_bedrock_mantle_context_overflow_maps_to_context_window_exceeded(): + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + original_exception = BaseLLMException( + status_code=400, + message=( + '{"error":{"code":"validation_error",' + '"message":"prompt tokens (1055489) exceed model maximum (1050000) for openai.gpt-5.6-sol",' + '"param":null,"type":"invalid_request_error"}}' + ), + ) + + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model="openai.gpt-5.6-sol", + original_exception=original_exception, + custom_llm_provider="bedrock_mantle", + ) + + assert excinfo.value.status_code == 400 + assert "prompt is too long: 1055489 tokens > 1050000 maximum" in excinfo.value.message diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index c6aac4d39917..b2a4263fade3 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -366,6 +366,17 @@ def test_shipped_rules_stack_adaptive_and_mid_conversation_flags(shipped_cost_ma assert info["supports_function_calling"] is True +def test_shipped_rules_flag_unmapped_fable_as_always_on_thinking(shipped_cost_map): + """An unmapped Fable/Mythos id picks up ``thinking_always_on`` from the + claude-always-on-thinking rule, while other unmapped Claudes stay unflagged.""" + model = "claude-fable-5-1" + assert model not in litellm.model_cost + info = litellm.get_model_info(model, custom_llm_provider="anthropic") + assert info["thinking_always_on"] is True + other = litellm.get_model_info("claude-opus-4-9", custom_llm_provider="anthropic") + assert other.get("thinking_always_on") is None + + @pytest.mark.parametrize( "model,provider", [ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 82de634b4889..a3dcdaf17371 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -64,7 +64,7 @@ def test_post_call_serializes_dict_with_datetime(logging_obj): assert "2026-05-11" in serialized -def test_sentry_sample_rate(): +def test_sentry_sample_rate(monkeypatch): existing_sample_rate = os.getenv("SENTRY_API_SAMPLE_RATE") try: # test with default value by removing the environment variable @@ -76,7 +76,7 @@ def test_sentry_sample_rate(): assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "1.0" # test with custom value - os.environ["SENTRY_API_SAMPLE_RATE"] = "0.5" + monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", "0.5") set_callbacks(["sentry"]) # Check if the custom sample rate is set correctly @@ -86,13 +86,13 @@ def test_sentry_sample_rate(): finally: # Restore the original environment variable if existing_sample_rate: - os.environ["SENTRY_API_SAMPLE_RATE"] = existing_sample_rate + monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", existing_sample_rate) else: if "SENTRY_API_SAMPLE_RATE" in os.environ: del os.environ["SENTRY_API_SAMPLE_RATE"] -def test_sentry_environment(): +def test_sentry_environment(monkeypatch): """Test that SENTRY_ENVIRONMENT is properly handled during Sentry initialization""" existing_environment = os.getenv("SENTRY_ENVIRONMENT") existing_dsn = os.getenv("SENTRY_DSN") @@ -115,7 +115,7 @@ def test_sentry_environment(): try: # Set a mock DSN to allow Sentry initialization - os.environ["SENTRY_DSN"] = "https://test@sentry.io/123456" + monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") # Test with default value (no environment set) if existing_environment: @@ -129,7 +129,7 @@ def test_sentry_environment(): assert call_kwargs["environment"] == "production" # Test with custom environment value - os.environ["SENTRY_ENVIRONMENT"] = "development" + monkeypatch.setenv("SENTRY_ENVIRONMENT", "development") mock_init.reset_mock() set_callbacks(["sentry"]) @@ -139,7 +139,7 @@ def test_sentry_environment(): assert call_kwargs["environment"] == "development" # Test with staging environment - os.environ["SENTRY_ENVIRONMENT"] = "staging" + monkeypatch.setenv("SENTRY_ENVIRONMENT", "staging") mock_init.reset_mock() set_callbacks(["sentry"]) @@ -154,13 +154,13 @@ def test_sentry_environment(): finally: # Restore the original environment variables if existing_environment: - os.environ["SENTRY_ENVIRONMENT"] = existing_environment + monkeypatch.setenv("SENTRY_ENVIRONMENT", existing_environment) else: if "SENTRY_ENVIRONMENT" in os.environ: del os.environ["SENTRY_ENVIRONMENT"] if existing_dsn: - os.environ["SENTRY_DSN"] = existing_dsn + monkeypatch.setenv("SENTRY_DSN", existing_dsn) else: if "SENTRY_DSN" in os.environ: del os.environ["SENTRY_DSN"] diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py index b953bfaa5656..f5339daad208 100644 --- a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py +++ b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py @@ -1,13 +1,14 @@ """Tests for the shared PTU rules: which deployments accrue flat cost, and what that zeroes.""" import os -from datetime import datetime, timezone +from datetime import date, datetime, timezone from unittest.mock import patch import pytest from litellm.litellm_core_utils.ptu_pricing import ( ptu_config_error, + ptu_identity_error, CUSTOM_PRICING_FIELDS, PTU_EMPTIED_PRICING_FIELDS, PTU_ZEROED_PRICING_FIELDS, @@ -209,3 +210,78 @@ def test_an_inverted_window_is_caught_before_the_count_and_rate_gate(): } assert ptu_config_error(window_only) == "ptu_effective_to must be after ptu_effective_from" + + +# --- the identity a config.yaml reservation has to declare --------------------------- + + +def test_a_declared_unique_id_is_accepted(): + assert ptu_identity_error(declared_id="azure-ptu-eastus", taken=False) is None + + +@pytest.mark.parametrize("missing", [None, ""], ids=["absent", "blank"]) +def test_a_reservation_without_an_id_is_refused(missing): + error = ptu_identity_error(declared_id=missing, taken=False) + + assert error is not None + assert error.startswith("model_info.id is required when PTU fields are set") + + +def test_the_refusal_names_the_id_the_deployment_already_uses(): + """An operator who invents a fresh name starts a second identity beside the charges + already written, which is the duplicate this rule exists to prevent.""" + error = ptu_identity_error(declared_id=None, taken=False, current_id="0ba149287615") + + assert error is not None + assert "0ba149287615" in error + + +def test_the_refusal_points_at_the_model_info_route_when_the_current_id_is_unknown(): + error = ptu_identity_error(declared_id=None, taken=False) + + assert error is not None + assert "GET /model/info" in error + + +def test_an_id_declared_twice_is_refused(): + error = ptu_identity_error(declared_id="azure-ptu-eastus", taken=True) + + assert error is not None + assert "declared on more than one deployment" in error + + +def test_the_deployment_is_named_when_the_caller_supplies_one(): + error = ptu_identity_error(declared_id=None, taken=False, model_name="azure-ptu") + + assert error is not None + assert error.startswith("PTU configuration on model 'azure-ptu' is invalid:") + + +def test_a_bare_yaml_date_bound_is_read_as_that_day_opening(): + """An unquoted 2027-01-01 in config.yaml loads as a date, not a string. Discarding it + took the whole deployment out of PTU handling, so it billed per token and accrued no + flat cost while the provider invoiced the reservation hourly.""" + terms = ptu_terms({**_VALID, "ptu_effective_to": date(2027, 1, 1)}) + + assert terms is not None + assert terms.effective_to == datetime(2027, 1, 1, tzinfo=timezone.utc) + + +def test_a_bare_yaml_date_start_is_read_as_that_day_opening(): + terms = ptu_terms({**_VALID, "ptu_effective_from": date(2026, 5, 1)}) + + assert terms is not None + assert terms.effective_from == datetime(2026, 5, 1, tzinfo=timezone.utc) + + +def test_the_string_zero_is_a_declared_id(): + """0 is a perfectly stable id, and ModelInfo stores it as a string. Reading it as absent + refused a deployment whose identity was never in doubt.""" + assert ptu_identity_error(declared_id="0", taken=False) is None + + +def test_an_empty_id_is_no_id(): + error = ptu_identity_error(declared_id="", taken=False) + + assert error is not None + assert error.startswith("model_info.id is required") diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index ccf353b1b6ce..61b63e2b9176 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -1,6 +1,4 @@ import json -import os -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -8,7 +6,6 @@ import litellm -sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( @@ -1326,7 +1323,7 @@ async def test_log_messages_includes_tools_in_model_call_details(): @pytest.mark.asyncio -async def test_realtime_guardrail_blocks_prompt_injection(): +async def test_realtime_guardrail_blocks_prompt_injection(monkeypatch: pytest.MonkeyPatch): """ Test that when a transcription event containing prompt injection arrives from the backend, a registered guardrail blocks it — sending a warning to the client @@ -1350,7 +1347,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) # --- client websocket mock --- client_ws = MagicMock() @@ -1405,11 +1402,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No f"Expected guardrail_violation error type, got: {error_events[0]}" ) - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_guardrail_allows_clean_transcript(): +async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.MonkeyPatch): """ Test that a clean transcript passes through the guardrail and triggers response.create to the backend. @@ -1430,7 +1426,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1463,11 +1459,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No response_creates = [e for e in sent_to_backend if e.get("type") == "response.create"] assert len(response_creates) == 1, f"Clean transcript should trigger response.create, got: {sent_to_backend}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_text_input_guardrail_blocks_and_returns_error(): +async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ Test that when conversation.item.create arrives with text that triggers a guardrail, the proxy blocks it (doesn't forward to backend) and returns an error event directly @@ -1495,7 +1490,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1558,11 +1553,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No ] assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(): +async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ Test that a client-supplied function_call_output whose content triggers a guardrail is blocked: it is not forwarded to the backend, and an error @@ -1590,7 +1584,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1648,11 +1642,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No assert sanitized_item["call_id"] == "call_123" assert "test@example.com" not in sanitized_item["output"] - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_function_call_output_guardrail_allows_clean_output(): +async def test_realtime_function_call_output_guardrail_allows_clean_output(monkeypatch: pytest.MonkeyPatch): """ Test that a clean function_call_output passes through and reaches the backend when guardrails are configured. @@ -1670,7 +1663,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1714,11 +1707,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No ] assert len(forwarded) == 1, f"Clean function_call_output should be forwarded, got: {forwarded}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_text_input_guardrail_uses_pre_call_mode(): +async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch): """ Test that _has_realtime_guardrails returns True for a guardrail configured with pre_call mode (not just realtime_input_transcription). @@ -1736,7 +1728,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() backend_ws = MagicMock() @@ -1751,11 +1743,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No "pre_call-only guardrail must not disable server_vad auto-response" ) - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_session_created_injects_session_update_for_audio_guardrail(): +async def test_realtime_session_created_injects_session_update_for_audio_guardrail(monkeypatch: pytest.MonkeyPatch): """ Test that when an audio transcription guardrail is configured, a session.created event from the backend triggers a session.update injection (create_response: false) @@ -1775,7 +1766,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1809,11 +1800,12 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No "GA session.update must nest turn_detection under audio.input" ) - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only(): +async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only( + monkeypatch: pytest.MonkeyPatch, +): """ pre_call-only guardrails must not inject create_response:false on realtime sessions — that breaks server_vad for audio-only voice agents (e.g. Model Armor). @@ -1831,7 +1823,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1853,11 +1845,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No session_updates = [e for e in sent_to_backend if e.get("type") == "session.update"] assert len(session_updates) == 0, f"pre_call-only guardrail must not inject session.update, got: {sent_to_backend}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(): +async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monkeypatch: pytest.MonkeyPatch): """Model Armor-style pre_call + post_call must not gate audio VAD.""" import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -1867,18 +1858,22 @@ class ModelArmorStyleGuardrail(CustomGuardrail): async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): return inputs - litellm.callbacks = [ - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_pre_call", - event_hook=GuardrailEventHooks.pre_call, - default_on=False, - ), - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_post_call", - event_hook=GuardrailEventHooks.post_call, - default_on=False, - ), - ] + monkeypatch.setattr( + litellm, + "callbacks", + [ + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_pre_call", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ), + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_post_call", + event_hook=GuardrailEventHooks.post_call, + default_on=False, + ), + ], + ) client_ws = MagicMock() backend_ws = MagicMock() @@ -1900,11 +1895,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No assert streaming._has_realtime_guardrails() is True assert streaming._has_audio_transcription_guardrails() is False - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_end_session_after_n_fails_closes_connection(): +async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.MonkeyPatch): """ Test that end_session_after_n_fails=2 closes the backend websocket after the second guardrail violation in a session. @@ -1923,7 +1917,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No default_on=True, end_session_after_n_fails=2, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1948,11 +1942,10 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No assert backend_ws.close.called, "Expected backend_ws.close() to be called after 2 violations" assert streaming._violation_count == 2 - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_on_violation_end_session_closes_on_first_fail(): +async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest.MonkeyPatch): """ Test that on_violation='end_session' closes the session immediately on the first violation, regardless of end_session_after_n_fails. @@ -1971,7 +1964,7 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No default_on=True, on_violation="end_session", ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1995,7 +1988,6 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No assert backend_ws.close.called, "Expected session to close immediately with on_violation=end_session" assert streaming._violation_count == 1 - litellm.callbacks = [] # cleanup @pytest.mark.asyncio @@ -2898,53 +2890,47 @@ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=No ) -def test_setup_folds_in_auto_response_disable_when_transcription_guardrail_active(): +def test_setup_folds_in_auto_response_disable_when_transcription_guardrail_active(monkeypatch: pytest.MonkeyPatch): """Gemini rejects a second setup, so a transcription guardrail's auto-response disable must be folded into the one-and-only setup; otherwise the model auto-responds and the guardrail is bypassed.""" import litellm - litellm.callbacks = [_transcription_guardrail()] - try: - streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) - setup = json.dumps( - { - "setup": { - "model": "models/gemini-3.1-flash-live-preview", - "generationConfig": {"responseModalities": ["AUDIO"]}, - "inputAudioTranscription": {}, - } + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) + setup = json.dumps( + { + "setup": { + "model": "models/gemini-3.1-flash-live-preview", + "generationConfig": {"responseModalities": ["AUDIO"]}, + "inputAudioTranscription": {}, } - ) - out = json.loads(streaming._maybe_inject_guardrail_auto_response_disable(setup)) - aad = out["setup"]["realtimeInputConfig"]["automaticActivityDetection"] - assert aad["disabled"] is True - finally: - litellm.callbacks = [] + } + ) + out = json.loads(streaming._maybe_inject_guardrail_auto_response_disable(setup)) + aad = out["setup"]["realtimeInputConfig"]["automaticActivityDetection"] + assert aad["disabled"] is True -def test_setup_unchanged_without_transcription_guardrail(): +def test_setup_unchanged_without_transcription_guardrail(monkeypatch: pytest.MonkeyPatch): import litellm - litellm.callbacks = [] + monkeypatch.setattr(litellm, "callbacks", []) streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) setup = json.dumps({"setup": {"model": "x", "generationConfig": {"responseModalities": ["AUDIO"]}}}) out = streaming._maybe_inject_guardrail_auto_response_disable(setup) assert json.loads(out) == json.loads(setup) -def test_non_bidi_setup_left_untouched_for_followup_capable_providers(): +def test_non_bidi_setup_left_untouched_for_followup_capable_providers(monkeypatch: pytest.MonkeyPatch): """OpenAI realtime accepts a follow-up session.update, so a non-bidi message (no top-level 'setup' key) must be left untouched even with a guardrail on.""" import litellm - litellm.callbacks = [_transcription_guardrail()] - try: - streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) - msg = json.dumps({"type": "session.update", "session": {"instructions": "hi"}}) - assert streaming._maybe_inject_guardrail_auto_response_disable(msg) == msg - finally: - litellm.callbacks = [] + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) + msg = json.dumps({"type": "session.update", "session": {"instructions": "hi"}}) + assert streaming._maybe_inject_guardrail_auto_response_disable(msg) == msg @pytest.mark.asyncio diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index eec4b307c87b..ee3e7719d52f 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -13,7 +13,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens @@ -634,7 +634,6 @@ def test_token_counter(): import unittest -from unittest.mock import MagicMock, patch from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding @@ -1025,13 +1024,12 @@ def test_token_counter_with_image_url(): } ] - try: + with pytest.raises(ValueError, match="Invalid detail value") as exc_info: token_counter(model="gpt-3.5-turbo", messages=messages_invalid) - pytest.fail("Expected ValueError for invalid detail value") - except ValueError as e: - assert "Invalid detail value" in str( - e - ), f"Expected detail validation error, got: {e}" + e = exc_info.value + assert "Invalid detail value" in str( + e + ), f"Expected detail validation error, got: {e}" def test_token_counter_with_thinking_content(): diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index 751b548adcda..aaaa43a0dc48 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -100,12 +100,12 @@ def test_encodes_path_segments_without_collapsing_valid_model_paths(self): @pytest.mark.parametrize("value", ["", ".", "..", None]) def test_rejects_empty_and_dot_segments(self, value): - with pytest.raises(ValueError, match="resource_id (is required|cannot be a dot path segment)"): + with pytest.raises(ValueError, match=r"resource_id (is required|cannot be a dot path segment)"): encode_url_path_segment(value, field_name="resource_id") @pytest.mark.parametrize("value", ["../model", "model/../other", "/model"]) def test_rejects_dot_segments_in_multi_segment_paths(self, value): - with pytest.raises(ValueError, match="model (is required|cannot be a dot path segment)"): + with pytest.raises(ValueError, match=r"model (is required|cannot be a dot path segment)"): encode_url_path_segments(value, field_name="model") diff --git a/tests/test_litellm/llms/anthropic/batches/test_transformation.py b/tests/test_litellm/llms/anthropic/batches/test_transformation.py index 4a2adb01ea5f..1635abcefd82 100644 --- a/tests/test_litellm/llms/anthropic/batches/test_transformation.py +++ b/tests/test_litellm/llms/anthropic/batches/test_transformation.py @@ -619,7 +619,6 @@ def fake_transform_parsed(*, completion_response, raw_response, model_response): # automatically. See base_batches_config_test.py. # --------------------------------------------------------------------------- # -from litellm.types.utils import LlmProviders # noqa: E402 from tests.test_litellm.llms.base_llm.batches.base_batches_config_test import ( # noqa: E402 BatchesConfigContractTests, ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index c38235b510e3..dab91aa59c32 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2813,18 +2813,6 @@ def test_raw_adaptive_thinking_untouched_for_46_plus_model(): assert result["thinking"] == {"type": "adaptive"} -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - @pytest.mark.parametrize( "model, expected", @@ -6206,3 +6194,41 @@ def test_is_anthropic_usage_object_rejects_responses_api_usage(): "output_tokens_details": {"reasoning_tokens": 0}, } ) + + +@pytest.mark.parametrize( + "model, expected_dropped", + [ + # always-on-thinking models reject thinking.type=disabled with a 400 + ("claude-fable-5", True), + ("claude-mythos-5", True), + # unmapped future family member -> claude-always-on-thinking fallback rule + ("claude-fable-5-1", True), + # adaptive-capable models that ACCEPT disabled must keep it verbatim + ("claude-opus-5", False), + ("claude-sonnet-5", False), + ("claude-opus-4-8", False), + # legacy models keep it verbatim + ("claude-sonnet-4-5-20250929", False), + ], +) +def test_disabled_thinking_omitted_only_for_always_on_models( + local_model_cost_map, model, expected_dropped +): + """``thinking={"type": "disabled"}`` is omitted for always-on-thinking models + (Fable/Mythos, which 400 on it: the API remedy is to omit the param) and is + forwarded verbatim for every model that accepts it.""" + config = AnthropicConfig() + + request = config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 64, "thinking": {"type": "disabled"}}, + litellm_params={}, + headers={}, + ) + + if expected_dropped: + assert "thinking" not in request + else: + assert request["thinking"] == {"type": "disabled"} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py index 5b7f2a60f689..f48d51dbe1e0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py @@ -66,5 +66,5 @@ def test_prepare_completion_kwargs_keeps_prompt_cache_key_through_responses_rero {"custom_llm_provider": "openai"}, thinking={"type": "enabled", "budget_tokens": 1024}, ) - assert completion_kwargs["model"] == "responses/openai/gpt-5.6-luna" + assert completion_kwargs["model"] == "openai/responses/gpt-5.6-luna" assert completion_kwargs["prompt_cache_key"] == "session-abc" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 4b2103e69eac..ffc1211709a0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -1202,6 +1202,8 @@ def _fake_user_api_key_auth( model_max_budget=None, end_user_model_max_budget=None, end_user_id=None, + user_model_max_budget=None, + user_id=None, token=None, ): """Build a minimal stand-in for ``UserAPIKeyAuth`` with just the fields @@ -1220,6 +1222,8 @@ class _Auth: auth.model_max_budget = model_max_budget auth.end_user_model_max_budget = end_user_model_max_budget auth.end_user_id = end_user_id + auth.user_model_max_budget = user_model_max_budget + auth.user_id = user_id auth.token = token return auth @@ -1548,6 +1552,78 @@ async def test_summary_model_denied_when_key_over_model_budget(): assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" +async def test_summary_model_denied_when_user_over_model_budget(): + """Internal-user per-model budget is enforced for the summary subrequest too. + + This file propagates `user_api_key_user_model_max_budget` into the summary + subrequest's metadata, so its spend charges the user's counter. Enforcing + only the key and end-user scopes would let compaction increment a counter it + can never be refused by, which is the asymmetry this PR exists to remove. + """ + import litellm + + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("