diff --git a/test-quality-budget.json b/test-quality-budget.json index 55c14ec2680..143378efec7 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -9,7 +9,7 @@ "limit": 1078 }, "TQ004": { - "limit": 757 + "limit": 557 }, "TQ005": { "limit": 2810 diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 382b41807d4..858ca482eb7 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1508,7 +1508,7 @@ def test_multiple_tool_calls_in_single_choice(): print("✓ Multiple tool calls are correctly grouped in a single choice") -def test_map_reasoning_effort_adds_summary_detailed(): +def test_map_reasoning_effort_adds_summary_detailed(monkeypatch): """ Test that _map_reasoning_effort behavior with reasoning_auto_summary flag. @@ -1571,7 +1571,7 @@ def test_map_reasoning_effort_adds_summary_detailed(): # Test 3: With env var enabled (flag disabled) - summary IS added litellm.reasoning_auto_summary = False - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") result = handler._map_reasoning_effort("high") assert ( @@ -1603,7 +1603,7 @@ def test_map_reasoning_effort_adds_summary_detailed(): # Restore original values litellm.reasoning_auto_summary = original_flag if original_env is not None: - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = original_env + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", original_env) elif "LITELLM_REASONING_AUTO_SUMMARY" in os.environ: del os.environ["LITELLM_REASONING_AUTO_SUMMARY"] diff --git a/tests/test_litellm/containers/test_container_transformation.py b/tests/test_litellm/containers/test_container_transformation.py index 555fe7773f0..f0432816fce 100644 --- a/tests/test_litellm/containers/test_container_transformation.py +++ b/tests/test_litellm/containers/test_container_transformation.py @@ -341,10 +341,10 @@ class TestOpenAIContainerTransformation: assert data["expires_after"] is None assert data["file_ids"] is None - def test_container_create_response_includes_cost(self): + def test_container_create_response_includes_cost(self, monkeypatch): """Test that container create response includes code interpreter cost calculation.""" # 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="") from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py index 88cc2275ae2..fbfd609cca6 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py @@ -88,49 +88,44 @@ async def test_send_email_success(mock_env_vars): @pytest.mark.asyncio -async def test_send_email_missing_api_key(): +async def test_send_email_missing_api_key(monkeypatch): # Remove the API key from environment before initializing logger - original_key = os.environ.pop("RESEND_API_KEY", None) + monkeypatch.delenv("RESEND_API_KEY", raising=False) - try: - # Initialize the logger after removing the API key - logger = ResendEmailLogger() + # Initialize the logger after removing the API key + logger = ResendEmailLogger() - # Test data - from_email = "test@example.com" - to_email = ["recipient@example.com"] - subject = "Test Subject" - html_body = "

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 5fe4b217e4f..b7fcce8dbf3 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() + 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 + 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/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index a4e16500aee..8d662311da1 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 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase 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/test_openmeter.py b/tests/test_litellm/integrations/test_openmeter.py index b9da99b6fa9..2d09e1572db 100644 --- a/tests/test_litellm/integrations/test_openmeter.py +++ b/tests/test_litellm/integrations/test_openmeter.py @@ -236,9 +236,9 @@ class TestOpenMeterIntegration: 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 @@ class TestOpenMeterIntegration: 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 @@ class TestOpenMeterIntegration: 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/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 cf36a2b9b25..052c08a86b5 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_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 2f32145580d..a5128228742 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 @@ -377,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 a10dc46eb42..3a7e06d085a 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 @@ -2844,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 @@ -2867,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 = { @@ -2927,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. @@ -2944,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 = [ @@ -2980,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_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 82de634b488..a3dcdaf1737 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/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 15ac73ed352..1ce683d76fc 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -528,7 +528,7 @@ class TestThinkingSummaryPreservation: finally: litellm.reasoning_auto_summary = original - def test_summary_added_when_env_var_set(self): + def test_summary_added_when_env_var_set(self, monkeypatch): """When LITELLM_REASONING_AUTO_SUMMARY env var is true, summary is added.""" import litellm from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( @@ -538,7 +538,7 @@ class TestThinkingSummaryPreservation: original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = False - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") completion_kwargs = { "model": "responses/gpt-5.2", "custom_llm_provider": "openai", diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 03cbfbb8609..17bab9bf6a5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -845,14 +845,14 @@ class TestTranslateThinkingToReasoning: finally: litellm.reasoning_auto_summary = original - def test_summary_added_when_env_var_set(self): + def test_summary_added_when_env_var_set(self, monkeypatch): """When LITELLM_REASONING_AUTO_SUMMARY env var is true, summary is included.""" import litellm original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = False - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") result = _ADAPTER.translate_thinking_to_reasoning( { "type": "enabled", diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py index bc26268ee92..c925bd7de45 100644 --- a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py +++ b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py @@ -239,8 +239,8 @@ class TestAPISerpentSearchIntegration: return mock_response @pytest.mark.asyncio - async def test_asearch_quick_default(self): - os.environ["APISERPENT_API_KEY"] = "test-api-key" + async def test_asearch_quick_default(self, monkeypatch): + monkeypatch.setenv("APISERPENT_API_KEY", "test-api-key") with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", new_callable=AsyncMock, @@ -269,8 +269,8 @@ class TestAPISerpentSearchIntegration: assert response.results[0].title == "Test Result" @pytest.mark.asyncio - async def test_asearch_deep(self): - os.environ["APISERPENT_API_KEY"] = "test-api-key" + async def test_asearch_deep(self, monkeypatch): + monkeypatch.setenv("APISERPENT_API_KEY", "test-api-key") with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", new_callable=AsyncMock, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index f7ad333293c..30f479bd7ff 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -40,8 +40,8 @@ class TestAzureMAIImageGeneration: assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") - def test_mai_flash_and_2e_model_pricing_in_cost_map(self): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") flash_info = litellm.get_model_info( @@ -328,8 +328,8 @@ class TestAzureMAIImageGeneration: assert image_response.usage.total_tokens == 1046 assert image_response.size == "1792x1024" - def test_mai_image_cost_calculator_token_based(self): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + def test_mai_image_cost_calculator_token_based(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") @@ -360,8 +360,8 @@ class TestAzureMAIImageGeneration: ) assert round(cost, 10) == round(expected_cost, 10) - def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index a2f88138eaa..2d6e938ea1f 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -678,10 +678,10 @@ def test_transform_request_helper_includes_anthropic_beta_and_tools(): assert fields["tools"][0]["type"] == "computer_20250124" -def test_parallel_tool_calls_config_kept_for_sonnet_5(): +def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): 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: config = AmazonConverseConfig() @@ -708,7 +708,7 @@ def test_parallel_tool_calls_config_kept_for_sonnet_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_parallel_tool_calls_config_dropped_for_ttl_only_model( @@ -3575,7 +3575,7 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): litellm.modify_params = original_modify_params -def test_supports_native_structured_outputs(): +def test_supports_native_structured_outputs(monkeypatch): """Test model detection for native structured outputs support. Support is driven by the ``supports_native_structured_output`` flag in the @@ -3583,7 +3583,7 @@ def test_supports_native_structured_outputs(): """ 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: config = AmazonConverseConfig() @@ -3645,7 +3645,7 @@ def test_supports_native_structured_outputs(): 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_create_output_config_for_response_format(): @@ -3683,11 +3683,11 @@ def test_create_output_config_for_response_format(): assert parsed_schema == expected -def test_translate_response_format_native_output_config(): +def test_translate_response_format_native_output_config(monkeypatch): """For supported models, _translate_response_format_param should produce outputConfig.""" 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: config = AmazonConverseConfig() @@ -3743,7 +3743,7 @@ def test_translate_response_format_native_output_config(): 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_translate_response_format_fallback_tool_call(): @@ -3778,11 +3778,11 @@ def test_translate_response_format_fallback_tool_call(): assert result["json_mode"] is True -def test_native_structured_output_no_fake_stream(): +def test_native_structured_output_no_fake_stream(monkeypatch): """When using native structured outputs with streaming, fake_stream should NOT be set.""" 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: config = AmazonConverseConfig() @@ -3828,7 +3828,7 @@ def test_native_structured_output_no_fake_stream(): 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_transform_request_with_output_config(): @@ -4116,7 +4116,7 @@ def test_add_additional_properties_definitions(): ) -def test_json_object_no_schema_skips_tool_injection(): +def test_json_object_no_schema_skips_tool_injection(monkeypatch): """response_format: {type: json_object} with no schema should NOT inject the synthetic json_tool_call tool. @@ -4126,7 +4126,7 @@ def test_json_object_no_schema_skips_tool_injection(): the model respond naturally with the JSON the caller asked for.""" 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: config = AmazonConverseConfig() @@ -4152,7 +4152,7 @@ def test_json_object_no_schema_skips_tool_injection(): 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_output_config_applies_additional_properties(): @@ -4805,7 +4805,7 @@ def test_cache_control_injection_tool_config_not_added_without_injection_point() assert all("cachePoint" not in tool for tool in tools) -def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(): +def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(monkeypatch): """ Regression test: cache_control_injection_points with location=tool_config must honor the requested `control.ttl`, mirroring the message/system @@ -4819,7 +4819,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(): """ 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: config = AmazonConverseConfig() @@ -4858,10 +4858,10 @@ def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(): 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_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacking_own_pricing(): +def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacking_own_pricing(monkeypatch): """ Regression test: a regional pricing entry that omits `cache_creation_input_token_cost_above_1hr` (e.g. `jp.anthropic.claude-opus-4-7`) @@ -4870,7 +4870,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacki """ 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: assert "cache_creation_input_token_cost_above_1hr" not in litellm.model_cost["jp.anthropic.claude-opus-4-7"] @@ -4911,7 +4911,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacki 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_cache_control_injection_tool_config_drops_ttl_for_unsupported_model(): diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 950336c7ad0..20bf65ee385 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -60,9 +60,9 @@ class TestAgentCoreSearch: """ @pytest.mark.asyncio - async def test_agentcore_search_request_payload(self): + async def test_agentcore_search_request_payload(self, monkeypatch): """Validates the MCP tools/call payload and SigV4 signing without real AWS calls.""" - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) mock_response = _make_mock_response(_mcp_response_body()) @@ -321,11 +321,11 @@ class TestAgentCoreSearch: assert headers["Authorization"] == "Bearer test-jwt-token" assert signed_body == json.dumps(request_data).encode() - def test_sign_request_uses_bearer_token_from_env(self): + def test_sign_request_uses_bearer_token_from_env(self, monkeypatch): """Server token is attached when the request targets the configured gateway host.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) try: headers, _ = config.sign_request( headers={}, @@ -338,11 +338,11 @@ class TestAgentCoreSearch: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) os.environ.pop("AGENTCORE_GATEWAY_URL", None) - def test_sign_request_refuses_server_token_to_untrusted_host(self): + def test_sign_request_refuses_server_token_to_untrusted_host(self, monkeypatch): """Server-managed token must not be sent to a caller-chosen api_base.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) try: with pytest.raises(ValueError, match="Refusing to send"): config.sign_request( @@ -355,11 +355,11 @@ class TestAgentCoreSearch: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) os.environ.pop("AGENTCORE_GATEWAY_URL", None) - def test_sign_request_uses_env_token_for_gateway_api_base_without_gateway_url(self): + def test_sign_request_uses_env_token_for_gateway_api_base_without_gateway_url(self, monkeypatch): """api_base pointing at a real gateway is a trusted destination for the env token, so operators configuring api_base in yaml don't also need AGENTCORE_GATEWAY_URL.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") os.environ.pop("AGENTCORE_GATEWAY_URL", None) try: headers, _ = config.sign_request( @@ -380,12 +380,12 @@ class TestAgentCoreSearch: "https://attacker.example.com/gw.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", ], ) - def test_sign_request_refuses_sigv4_to_untrusted_host(self, untrusted_api_base): + def test_sign_request_refuses_sigv4_to_untrusted_host(self, untrusted_api_base, monkeypatch): """A SigV4 signature carries the proxy's credential scope and session token, so it must never be sent to a host that is not the operator's gateway.""" config = AgentCoreSearchConfig() os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) try: with patch.object( AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM @@ -410,12 +410,12 @@ class TestAgentCoreSearch: "http://internal-gateway.corp/mcp", ], ) - def test_sign_request_refuses_server_token_over_plaintext_http(self, plaintext_api_base): + def test_sign_request_refuses_server_token_over_plaintext_http(self, plaintext_api_base, monkeypatch): """A trusted hostname over plain http would expose the bearer token to network observers, so credentials only ride https (or localhost).""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = plaintext_api_base + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", plaintext_api_base) try: with pytest.raises(ValueError, match="plaintext"): config.sign_request( @@ -446,11 +446,11 @@ class TestAgentCoreSearch: ) mock_base_sign.assert_not_called() - def test_sign_request_allows_plain_http_for_localhost(self): + def test_sign_request_allows_plain_http_for_localhost(self, monkeypatch): """Local development against an MCP stub on 127.0.0.1 keeps working.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = "http://127.0.0.1:8931/mcp" + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", "http://127.0.0.1:8931/mcp") try: headers, _ = config.sign_request( headers={}, @@ -483,11 +483,11 @@ class TestAgentCoreSearch: # AWS_BEARER_TOKEN_BEDROCK env fallback. assert mock_base_sign.call_args.kwargs["api_key"] == "" - def test_sign_request_custom_hostname_requires_region(self): + def test_sign_request_custom_hostname_requires_region(self, monkeypatch): """Custom hostname + empty AWS config chain → clear error, no guessed region.""" config = AgentCoreSearchConfig() custom_url = "https://gateway.internal.example.com/mcp" - os.environ["AGENTCORE_GATEWAY_URL"] = custom_url + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", custom_url) mock_session = MagicMock() mock_session.region_name = None # nothing configured anywhere @@ -503,11 +503,11 @@ class TestAgentCoreSearch: finally: os.environ.pop("AGENTCORE_GATEWAY_URL", None) - def test_sign_request_custom_hostname_uses_shared_config_region(self): + def test_sign_request_custom_hostname_uses_shared_config_region(self, monkeypatch): """Custom hostname + region from AWS shared config (profile) must be honored.""" config = AgentCoreSearchConfig() custom_url = "https://gateway.internal.example.com/mcp" - os.environ["AGENTCORE_GATEWAY_URL"] = custom_url + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", custom_url) mock_session = MagicMock() mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py b/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py index daedbe5052c..962933aba28 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py @@ -40,12 +40,12 @@ class TestBedrockSSLVerify: ssl_verify = base_aws._get_ssl_verify() assert ssl_verify is True - def test_base_aws_llm_get_ssl_verify_false(self): + def test_base_aws_llm_get_ssl_verify_false(self, monkeypatch): """Test that _get_ssl_verify returns False when SSL verification is disabled.""" base_aws = BaseAWSLLM() # Set SSL_VERIFY to False via environment - os.environ["SSL_VERIFY"] = "False" + monkeypatch.setenv("SSL_VERIFY", "False") ssl_verify = base_aws._get_ssl_verify() assert ssl_verify is False @@ -53,7 +53,7 @@ class TestBedrockSSLVerify: # Clean up os.environ.pop("SSL_VERIFY", None) - def test_base_aws_llm_get_ssl_verify_custom_ca_bundle(self): + def test_base_aws_llm_get_ssl_verify_custom_ca_bundle(self, monkeypatch): """Test that _get_ssl_verify returns custom CA bundle path when SSL_CERT_FILE is set.""" base_aws = BaseAWSLLM() @@ -66,7 +66,7 @@ class TestBedrockSSLVerify: try: # Set SSL_CERT_FILE environment variable - os.environ["SSL_CERT_FILE"] = ca_bundle_path + monkeypatch.setenv("SSL_CERT_FILE", ca_bundle_path) os.environ.pop("SSL_VERIFY", None) litellm.ssl_verify = True @@ -327,7 +327,7 @@ class TestBedrockSSLVerify: os.environ.pop("SSL_CERT_FILE", None) os.unlink(ca_bundle_path) - def test_ssl_verify_priority_env_over_litellm_config(self): + def test_ssl_verify_priority_env_over_litellm_config(self, monkeypatch): """Test that SSL_VERIFY environment variable takes priority over litellm.ssl_verify.""" base_aws = BaseAWSLLM() @@ -335,7 +335,7 @@ class TestBedrockSSLVerify: litellm.ssl_verify = True # Set SSL_VERIFY environment variable to False - os.environ["SSL_VERIFY"] = "False" + monkeypatch.setenv("SSL_VERIFY", "False") try: ssl_verify = base_aws._get_ssl_verify() @@ -345,7 +345,7 @@ class TestBedrockSSLVerify: os.environ.pop("SSL_VERIFY", None) litellm.ssl_verify = True - def test_ssl_cert_file_priority_over_default(self): + def test_ssl_cert_file_priority_over_default(self, monkeypatch): """Test that SSL_CERT_FILE takes priority when ssl_verify is True.""" base_aws = BaseAWSLLM() @@ -358,7 +358,7 @@ class TestBedrockSSLVerify: try: # Set SSL_CERT_FILE environment variable - os.environ["SSL_CERT_FILE"] = ca_bundle_path + monkeypatch.setenv("SSL_CERT_FILE", ca_bundle_path) os.environ.pop("SSL_VERIFY", None) litellm.ssl_verify = True diff --git a/tests/test_litellm/llms/crusoe/test_crusoe.py b/tests/test_litellm/llms/crusoe/test_crusoe.py index 0a05126919a..34a6d37663b 100644 --- a/tests/test_litellm/llms/crusoe/test_crusoe.py +++ b/tests/test_litellm/llms/crusoe/test_crusoe.py @@ -105,14 +105,14 @@ def test_crusoe_provider_detection_by_prefix(): assert model == "meta-llama/Llama-3.3-70B-Instruct" -def test_crusoe_model_list_populated(): +def test_crusoe_model_list_populated(monkeypatch): """Test Crusoe models are present in model_prices_and_context_window.json""" import litellm original_model_cost = litellm.model_cost original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") try: - 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="") expected = [ @@ -132,4 +132,4 @@ def test_crusoe_model_list_populated(): if original_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", original_env) diff --git a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py index 3f772b263fd..153d37d549c 100644 --- a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py +++ b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py @@ -83,8 +83,8 @@ class TestDataRobotConfig: == api_base ) - def test_resolve_api_base_with_environment_variable(self, handler): - os.environ["DATAROBOT_ENDPOINT"] = "https://env.datarobot.com" + def test_resolve_api_base_with_environment_variable(self, handler, monkeypatch): + monkeypatch.setenv("DATAROBOT_ENDPOINT", "https://env.datarobot.com") assert ( handler._resolve_api_base(None) == "https://env.datarobot.com/api/v2/genai/llmgw/chat/completions/" @@ -101,7 +101,7 @@ class TestDataRobotConfig: def test_resolve_api_key(self, api_key, expected_api_key, handler): assert handler._resolve_api_key(api_key) == expected_api_key - def test_resolve_api_key_with_environment_variable(self, handler): - os.environ["DATAROBOT_API_TOKEN"] = "env_key" + def test_resolve_api_key_with_environment_variable(self, handler, monkeypatch): + monkeypatch.setenv("DATAROBOT_API_TOKEN", "env_key") assert handler._resolve_api_key(None) == "env_key" del os.environ["DATAROBOT_API_TOKEN"] diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py index a5eb836e71d..ff309bc44ed 100644 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py @@ -11,14 +11,14 @@ sys.path.insert(0, os.path.abspath("../../../..")) import litellm -def test_deepseek_supported_openai_params(): +def test_deepseek_supported_openai_params(monkeypatch): """ Test "reasoning_effort" is an openai param supported for the DeepSeek model on deepinfra """ from litellm.llms.deepinfra.chat.transformation import DeepInfraConfig # Ensure we're using the local model cost map - 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="") supported_openai_params = DeepInfraConfig().get_supported_openai_params( diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 6917092966b..fc8d71afaa9 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -81,8 +81,8 @@ def test_no_usage_details(): assert cost == 0.0 -def test_gemini_image_edit_cost_prefers_token_usage_metadata(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_edit_cost_prefers_token_usage_metadata(monkeypatch): + monkeypatch.setenv("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") @@ -120,8 +120,8 @@ def test_gemini_image_edit_cost_prefers_token_usage_metadata(): assert cost != flat_image_cost -def test_gemini_image_edit_cost_uses_output_token_details(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_edit_cost_uses_output_token_details(monkeypatch): + monkeypatch.setenv("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") @@ -176,8 +176,8 @@ def test_gemini_image_edit_cost_uses_output_token_details(): assert cost != all_output_as_image_cost -def test_gemini_image_generation_cost_uses_output_token_details(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_generation_cost_uses_output_token_details(monkeypatch): + monkeypatch.setenv("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") @@ -232,8 +232,8 @@ def test_gemini_image_generation_cost_uses_output_token_details(): assert cost != all_output_as_image_cost -def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(monkeypatch): + monkeypatch.setenv("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") @@ -264,8 +264,8 @@ def _image_response_with_web_search(web_search_requests): return ImageResponse(data=[ImageObject(b64_json="img1")], usage=usage) -def test_gemini_image_generation_cost_adds_web_search_grounding(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_generation_cost_adds_web_search_grounding(monkeypatch): + monkeypatch.setenv("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") @@ -286,8 +286,8 @@ def test_gemini_image_generation_cost_adds_web_search_grounding(): assert round(grounded - ungrounded, 10) == round(expected_web_search_cost, 10) -def test_gemini_image_generation_cost_no_web_search_when_absent(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_generation_cost_no_web_search_when_absent(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py index 0750fb9e405..cff3c6be940 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py @@ -231,10 +231,10 @@ def test_inception_in_provider_lists(): assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints -def test_inception_model_configuration(): +def test_inception_model_configuration(monkeypatch): from litellm import get_model_info - 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="") litellm.inception_models = set() litellm.add_known_models() @@ -251,8 +251,8 @@ def test_inception_model_configuration(): assert info.get("supports_response_schema") is True -def test_inception_model_list_populated(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_inception_model_list_populated(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") litellm.inception_models = set() litellm.add_known_models() diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py index 9b7c8dd3742..62688a13c35 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py @@ -143,10 +143,10 @@ async def test_inception_fim_async(): assert r.choices[0].text == "a + b" -def test_inception_fim_model_configuration(): +def test_inception_fim_model_configuration(monkeypatch): from litellm import get_model_info - 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="") litellm.text_completion_inception_models = set() litellm.add_known_models() diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py index cd866187166..e54e25cbd18 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py @@ -30,8 +30,8 @@ def _image_response_with_web_search(web_search_requests): return ImageResponse(data=[ImageObject(b64_json="img1")], usage=usage) -def test_vertex_image_generation_cost_adds_web_search_grounding(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_vertex_image_generation_cost_adds_web_search_grounding(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") @@ -55,8 +55,8 @@ def test_vertex_image_generation_cost_adds_web_search_grounding(): assert round(grounded - ungrounded, 10) == round(expected_web_search_cost, 10) -def test_vertex_image_generation_cost_no_web_search_when_absent(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_vertex_image_generation_cost_no_web_search_when_absent(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3-pro-image-preview" diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index e8374f92a19..61e1121257c 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -51,11 +51,11 @@ def test_zai_in_provider_lists(): assert "zai" in litellm.provider_list -def test_zai_models_in_model_cost(): +def test_zai_models_in_model_cost(monkeypatch): """Test that ZAI models are in the model cost map""" import os - 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="") zai_models = [ @@ -75,11 +75,11 @@ def test_zai_models_in_model_cost(): assert litellm.model_cost[model]["litellm_provider"] == "zai" -def test_zai_glm46_cost_calculation(): +def test_zai_glm46_cost_calculation(monkeypatch): """Test the cost calculation for glm-4.6""" import os - 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="") key = "zai/glm-4.6" @@ -96,11 +96,11 @@ def test_zai_glm46_cost_calculation(): assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) -def test_zai_flash_model_is_free(): +def test_zai_flash_model_is_free(monkeypatch): """Test that glm-4.5-flash has zero cost""" import os - 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="") key = "zai/glm-4.5-flash" @@ -110,11 +110,11 @@ def test_zai_flash_model_is_free(): assert info["output_cost_per_token"] == 0 -def test_glm47_supports_reasoning(): +def test_glm47_supports_reasoning(monkeypatch): """Test that GLM-4.7 supports reasoning""" import os - 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="") key = "zai/glm-4.7" @@ -124,11 +124,11 @@ def test_glm47_supports_reasoning(): assert info["supports_reasoning"] is True -def test_glm47_cost_calculation(): +def test_glm47_cost_calculation(monkeypatch): """Test cost calculation for GLM-4.7""" import os - 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="") prompt_cost, completion_cost = cost_per_token( diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index c589014f276..1c66acf8678 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -109,7 +109,7 @@ async def test_authenticate_user_admin_login_with_ui_credentials(): @pytest.mark.asyncio -async def test_authenticate_user_admin_login_with_master_key_as_password(): +async def test_authenticate_user_admin_login_with_master_key_as_password(monkeypatch): """Test admin login when UI_PASSWORD is not set, should use master_key""" master_key = "sk-1234" ui_username = "admin" @@ -131,39 +131,35 @@ async def test_authenticate_user_admin_login_with_master_key_as_password(): with patch.dict(os.environ, env_vars, clear=False): # Explicitly remove UI_PASSWORD if it exists - original_ui_password = os.environ.pop("UI_PASSWORD", None) - try: + monkeypatch.delenv("UI_PASSWORD", raising=False) + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = { + "token": "test-token-123", + "user_id": LITELLM_PROXY_ADMIN_NAME, + } + with patch( - "litellm.proxy.auth.login_utils.generate_key_helper_fn", + "litellm.proxy.auth.login_utils.user_update", new_callable=AsyncMock, - ) as mock_generate_key: - mock_generate_key.return_value = { - "token": "test-token-123", - "user_id": LITELLM_PROXY_ADMIN_NAME, - } - + return_value=None, + ) as mock_user_update: with patch( - "litellm.proxy.auth.login_utils.user_update", - new_callable=AsyncMock, - return_value=None, - ) as mock_user_update: - with patch( - "litellm.proxy.auth.login_utils.get_secret_bool", - return_value=False, - ): - result = await authenticate_user( - username=ui_username, - password=master_key, - master_key=master_key, - prisma_client=mock_prisma_client, - ) + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ): + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + ) - assert isinstance(result, LoginResult) - assert result.user_id == LITELLM_PROXY_ADMIN_NAME - assert result.user_role == LitellmUserRoles.PROXY_ADMIN - finally: - if original_ui_password: - os.environ["UI_PASSWORD"] = original_ui_password + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + assert result.user_role == LitellmUserRoles.PROXY_ADMIN @pytest.mark.asyncio @@ -319,7 +315,7 @@ async def test_authenticate_user_email_case_insensitive_login(): @pytest.mark.asyncio -async def test_authenticate_user_database_required_for_admin(): +async def test_authenticate_user_database_required_for_admin(monkeypatch): """Test that database is required for admin login""" master_key = "sk-1234" ui_username = "admin" @@ -353,7 +349,7 @@ async def test_authenticate_user_database_required_for_admin(): assert "No Database connected" in exc_info.value.message finally: if original_db_url: - os.environ["DATABASE_URL"] = original_db_url + monkeypatch.setenv("DATABASE_URL", original_db_url) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py index a2b8894910c..af0686fcc59 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py @@ -17,14 +17,14 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.exceptions import GuardrailRaisedException -def test_deepkeep_guard_config(): +def test_deepkeep_guard_config(monkeypatch): """Test DeepKeep guard configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} - os.environ["DEEPKEEP_API_KEY"] = "test-key" - os.environ["DEEPKEEP_API_BASE"] = "https://test.deepkeep.ai" - os.environ["DEEPKEEP_FIREWALL_ID"] = "fw-123" + monkeypatch.setenv("DEEPKEEP_API_KEY", "test-key") + monkeypatch.setenv("DEEPKEEP_API_BASE", "https://test.deepkeep.ai") + monkeypatch.setenv("DEEPKEEP_FIREWALL_ID", "fw-123") init_guardrails_v2( all_guardrails=[ @@ -108,11 +108,11 @@ class TestDeepKeepGuardrail: == "https://test.deepkeep.ai/v3/openai/beta/litellm_basic_guardrail_api" ) - def test_initialization_with_env_vars(self): + def test_initialization_with_env_vars(self, monkeypatch): """should initialize successfully using environment variables.""" - os.environ["DEEPKEEP_API_KEY"] = "env-key" - os.environ["DEEPKEEP_API_BASE"] = "https://env.deepkeep.ai" - os.environ["DEEPKEEP_FIREWALL_ID"] = "fw-env-456" + monkeypatch.setenv("DEEPKEEP_API_KEY", "env-key") + monkeypatch.setenv("DEEPKEEP_API_BASE", "https://env.deepkeep.ai") + monkeypatch.setenv("DEEPKEEP_FIREWALL_ID", "fw-env-456") guardrail = DeepKeepGuardrail( guardrail_name="deepkeep-env-test", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index c5b182a00ab..57adf85b3d9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -26,13 +26,13 @@ from litellm.types.utils import ( ) -def test_hiddenlayer_config_saas(): +def test_hiddenlayer_config_saas(monkeypatch): """Test Hiddenlayer SaaS configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variables for testing - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") init_guardrails_v2( all_guardrails=[ @@ -71,9 +71,9 @@ class TestHiddenlayerGuardrail: if key in os.environ: del os.environ[key] - def test_initialization(self): + def test_initialization(self, monkeypatch): """Test successful initialization with default values.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -94,9 +94,9 @@ class TestHiddenlayerGuardrail: HiddenlayerGuardrail(guardrail_name="hiddenlayer", event_hook="pre_call") @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self): + async def test_apply_guardrail_request_no_violations(self, monkeypatch): """Test apply_guardrail for request with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -151,9 +151,9 @@ class TestHiddenlayerGuardrail: assert call_args.args[0] == f"{guardrail.api_base}/detection/v1/interactions" @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self): + async def test_apply_guardrail_request_with_violations(self, monkeypatch): """Test apply_guardrail for request with violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -209,9 +209,9 @@ class TestHiddenlayerGuardrail: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self): + async def test_apply_guardrail_response_no_violations(self, monkeypatch): """Test apply_guardrail for response with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -279,10 +279,10 @@ class TestHiddenlayerGuardrail: mock_post.assert_called_once() @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self): + async def test_apply_guardrail_response_with_violations(self, monkeypatch): """Test apply_guardrail for response with violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -348,10 +348,10 @@ class TestHiddenlayerGuardrail: assert exc_info.value.status_code == 400 @pytest.mark.asyncio - async def test_apply_guardrail_api_error_handling(self): + async def test_apply_guardrail_api_error_handling(self, monkeypatch): """Test handling of API errors in apply_guardrail.""" # Set required API key - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -391,10 +391,10 @@ class TestHiddenlayerGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_validate_with_call_hiddenlayer_method(self): + async def test_validate_with_call_hiddenlayer_method(self, monkeypatch): """Test the _validate_with_guard_server internal method.""" # Set required API key - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -433,9 +433,9 @@ class TestHiddenlayerGuardrail: ) @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image(self): + async def test_apply_guardrail_request_with_image(self, monkeypatch): """Test apply_guardrail sends multimodal content (image) to HiddenLayer v1.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -498,9 +498,9 @@ class TestHiddenlayerGuardrail: assert result is not None @pytest.mark.asyncio - async def test_apply_guardrail_redact_with_image_content(self): + async def test_apply_guardrail_redact_with_image_content(self, monkeypatch): """Test that REDACT action with multimodal content extracts text properly into inputs['texts'].""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -570,12 +570,12 @@ class TestHiddenlayerGuardrail: assert config_model.__name__ == "HiddenlayerGuardrailConfigModel" -def test_hiddenlayer_config_v2(): +def test_hiddenlayer_config_v2(monkeypatch): """Test HiddenLayer V2 configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") init_guardrails_v2( all_guardrails=[ @@ -612,9 +612,9 @@ class TestHiddenlayerGuardrailV2: if key in os.environ: del os.environ[key] - def test_initialization(self): + def test_initialization(self, monkeypatch): """Test successful initialization with default values.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -633,9 +633,9 @@ class TestHiddenlayerGuardrailV2: HiddenlayerGuardrailV2(guardrail_name="hiddenlayer", event_hook="pre_call") @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self): + async def test_apply_guardrail_request_no_violations(self, monkeypatch): """Test apply_guardrail for request with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -691,9 +691,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/request-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self): + async def test_apply_guardrail_request_with_violations(self, monkeypatch): """Test apply_guardrail for request with violations detected (block via header).""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -751,9 +751,9 @@ class TestHiddenlayerGuardrailV2: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self): + async def test_apply_guardrail_response_no_violations(self, monkeypatch): """Test apply_guardrail for response with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="post_call", default_on=True @@ -816,9 +816,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self): + async def test_apply_guardrail_response_with_violations(self, monkeypatch): """Test apply_guardrail for response with violations detected (block via header).""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="post_call", default_on=True @@ -863,9 +863,9 @@ class TestHiddenlayerGuardrailV2: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_with_tool_calls(self): + async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch): """Test apply_guardrail for response containing tool calls.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="post_call", default_on=True @@ -924,9 +924,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_call_hiddenlayer_uses_correct_endpoints(self): + async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch): """Test that _call_hiddenlayer uses the v2 request/response evaluation endpoints.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -959,9 +959,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in mock_post.call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image(self): + async def test_apply_guardrail_request_with_image(self, monkeypatch): """Test apply_guardrail sends multimodal content (image) to HiddenLayer v2.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -1030,9 +1030,9 @@ class TestHiddenlayerGuardrailV2: assert texts == ["how much is on this receipt?"] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image_multimodal_response(self): + async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch): """Test that new_texts extraction handles multimodal content (list) returned by HiddenLayer v2.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py index 16185cadbdf..dcb004e5422 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -19,13 +19,13 @@ from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import ( from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -def test_lasso_guard_config(): +def test_lasso_guard_config(monkeypatch): """Test Lasso guard configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variable for testing - os.environ["LASSO_API_KEY"] = "test-key" + monkeypatch.setenv("LASSO_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py index c7a6df1361e..fa4624eac99 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py @@ -18,14 +18,14 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message -def test_onyx_guard_config(): +def test_onyx_guard_config(monkeypatch): """Test Onyx guard configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") init_guardrails_v2( all_guardrails=[ @@ -48,11 +48,11 @@ def test_onyx_guard_config(): del os.environ["ONYX_API_KEY"] -def test_onyx_guard_with_custom_timeout_from_kwargs(): +def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch): """Test Onyx guard instantiation with custom timeout passed via kwargs.""" # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -81,16 +81,16 @@ def test_onyx_guard_with_custom_timeout_from_kwargs(): del os.environ["ONYX_API_KEY"] -def test_onyx_guard_with_timeout_none_uses_env_var(): +def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch): """Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var. When timeout=None is passed (as it would be from config model with default None), the ONYX_TIMEOUT environment variable should be used. """ # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" - os.environ["ONYX_TIMEOUT"] = "60" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") + monkeypatch.setenv("ONYX_TIMEOUT", "60") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -121,11 +121,11 @@ def test_onyx_guard_with_timeout_none_uses_env_var(): del os.environ["ONYX_TIMEOUT"] -def test_onyx_guard_with_timeout_none_defaults_to_10(): +def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch): """Test Onyx guard with timeout=None and no env var defaults to 10 seconds.""" # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Ensure ONYX_TIMEOUT is not set if "ONYX_TIMEOUT" in os.environ: del os.environ["ONYX_TIMEOUT"] @@ -174,10 +174,10 @@ class TestOnyxGuardrail: if key in os.environ: del os.environ[key] - def test_initialization_with_defaults(self): + def test_initialization_with_defaults(self, monkeypatch): """Test successful initialization with default values.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -189,10 +189,10 @@ class TestOnyxGuardrail: assert guardrail.guardrail_name == "test-guard" assert guardrail.event_hook == "pre_call" - def test_initialization_with_env_vars(self): + def test_initialization_with_env_vars(self, monkeypatch): """Test initialization with environment variables.""" - os.environ["ONYX_API_BASE"] = "https://custom.onyx.security" - os.environ["ONYX_API_KEY"] = "custom-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://custom.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "custom-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -213,9 +213,9 @@ class TestOnyxGuardrail: ): OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call") - def test_initialization_with_default_timeout(self): + def test_initialization_with_default_timeout(self, monkeypatch): """Test that default timeout is 10.0 seconds.""" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -232,9 +232,9 @@ class TestOnyxGuardrail: assert timeout_param.read == 10.0 assert timeout_param.connect == 5.0 - def test_initialization_with_custom_timeout_parameter(self): + def test_initialization_with_custom_timeout_parameter(self, monkeypatch): """Test initialization with custom timeout parameter.""" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -254,14 +254,14 @@ class TestOnyxGuardrail: assert timeout_param.read == 30.0 assert timeout_param.connect == 5.0 - def test_initialization_with_timeout_from_env_var(self): + def test_initialization_with_timeout_from_env_var(self, monkeypatch): """Test initialization with timeout from ONYX_TIMEOUT environment variable. Note: The env var is only used when timeout=None is explicitly passed, since the default parameter value is 10.0 (not None). """ - os.environ["ONYX_API_KEY"] = "test-api-key" - os.environ["ONYX_TIMEOUT"] = "25" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") + monkeypatch.setenv("ONYX_TIMEOUT", "25") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -282,10 +282,10 @@ class TestOnyxGuardrail: assert timeout_param.read == 25.0 assert timeout_param.connect == 5.0 - def test_initialization_timeout_parameter_overrides_env_var(self): + def test_initialization_timeout_parameter_overrides_env_var(self, monkeypatch): """Test that timeout parameter overrides ONYX_TIMEOUT environment variable.""" - os.environ["ONYX_API_KEY"] = "test-api-key" - os.environ["ONYX_TIMEOUT"] = "25" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") + monkeypatch.setenv("ONYX_TIMEOUT", "25") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -306,10 +306,10 @@ class TestOnyxGuardrail: assert timeout_param.connect == 5.0 @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self): + async def test_apply_guardrail_request_no_violations(self, monkeypatch): """Test apply_guardrail for request with no violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -372,10 +372,10 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-call-id" @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self): + async def test_apply_guardrail_request_with_violations(self, monkeypatch): """Test apply_guardrail for request with violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -423,10 +423,10 @@ class TestOnyxGuardrail: assert "prompt_injection" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self): + async def test_apply_guardrail_response_no_violations(self, monkeypatch): """Test apply_guardrail for response with no violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -497,10 +497,10 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-call-id-2" @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self): + async def test_apply_guardrail_response_with_violations(self, monkeypatch): """Test apply_guardrail for response with violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -558,10 +558,10 @@ class TestOnyxGuardrail: assert "illegal_instructions" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_api_error_handling(self): + async def test_apply_guardrail_api_error_handling(self, monkeypatch): """Test handling of API errors in apply_guardrail.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -591,10 +591,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_timeout_error_handling(self): + async def test_apply_guardrail_timeout_error_handling(self, monkeypatch): """Test handling of timeout errors in apply_guardrail (graceful degradation).""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", @@ -629,10 +629,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_read_timeout_error_handling(self): + async def test_apply_guardrail_read_timeout_error_handling(self, monkeypatch): """Test handling of read timeout errors in apply_guardrail.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", @@ -667,10 +667,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_connect_timeout_error_handling(self): + async def test_apply_guardrail_connect_timeout_error_handling(self, monkeypatch): """Test handling of connect timeout errors in apply_guardrail.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", @@ -705,10 +705,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_no_logging_obj(self): + async def test_apply_guardrail_no_logging_obj(self, monkeypatch): """Test apply_guardrail without logging object (uses UUID).""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -747,10 +747,10 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-uuid" @pytest.mark.asyncio - async def test_validate_with_guard_server_method(self): + async def test_validate_with_guard_server_method(self, monkeypatch): """Test the _validate_with_guard_server internal method.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -788,10 +788,10 @@ class TestOnyxGuardrail: ) @pytest.mark.asyncio - async def test_validate_with_guard_server_blocked(self): + async def test_validate_with_guard_server_blocked(self, monkeypatch): """Test _validate_with_guard_server when request is blocked.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -825,10 +825,10 @@ class TestOnyxGuardrail: assert config_model.__name__ == "OnyxGuardrailConfigModel" @pytest.mark.asyncio - async def test_apply_guardrail_with_modelresponse(self): + async def test_apply_guardrail_with_modelresponse(self, monkeypatch): """Test apply_guardrail with ModelResponse object for response type.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -880,10 +880,10 @@ class TestOnyxGuardrail: assert "payload" in call_args.kwargs["json"] @pytest.mark.asyncio - async def test_apply_guardrail_response_error_handling(self): + async def test_apply_guardrail_response_error_handling(self, monkeypatch): """Test error handling when processing response data.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -925,11 +925,11 @@ class TestOnyxIntegration: """Test integration scenarios.""" @pytest.mark.asyncio - async def test_full_guardrail_flow(self): + async def test_full_guardrail_flow(self, monkeypatch): """Test full guardrail flow with multiple hooks.""" # Set environment variables - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ @@ -973,10 +973,10 @@ class TestOnyxIntegration: del os.environ["ONYX_API_KEY"] @pytest.mark.asyncio - async def test_apply_guardrail_empty_request_data(self): + async def test_apply_guardrail_empty_request_data(self, monkeypatch): """Test apply_guardrail with empty request data.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py index 55f01ebddfd..1322d93ce70 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py @@ -93,24 +93,24 @@ class TestRepelloAIInitialization: with pytest.raises(ValueError, match="asset_id"): RepelloAIGuardrail(api_key="test-api-key", guardrail_name="t") - def test_api_key_from_env(self): - os.environ["REPELLOAI_API_KEY"] = "env-key" + def test_api_key_from_env(self, monkeypatch): + monkeypatch.setenv("REPELLOAI_API_KEY", "env-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "env-key" - def test_api_key_from_argus_env(self): - os.environ["ARGUS_API_KEY"] = "argus-key" + def test_api_key_from_argus_env(self, monkeypatch): + monkeypatch.setenv("ARGUS_API_KEY", "argus-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "argus-key" - def test_argus_env_preferred_over_legacy(self): - os.environ["ARGUS_API_KEY"] = "argus-key" - os.environ["REPELLOAI_API_KEY"] = "legacy-key" + def test_argus_env_preferred_over_legacy(self, monkeypatch): + monkeypatch.setenv("ARGUS_API_KEY", "argus-key") + monkeypatch.setenv("REPELLOAI_API_KEY", "legacy-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "argus-key" - def test_explicit_api_key_preferred_over_env(self): - os.environ["ARGUS_API_KEY"] = "argus-key" + def test_explicit_api_key_preferred_over_env(self, monkeypatch): + monkeypatch.setenv("ARGUS_API_KEY", "argus-key") guardrail = RepelloAIGuardrail( api_key="explicit-key", asset_id="asset-123", guardrail_name="t" ) @@ -145,10 +145,10 @@ class TestRepelloAIInitialization: assert guardrail.api_base == DEFAULT_REPELLOAI_API_BASE assert guardrail.unreachable_fallback == "fail_closed" - def test_init_guardrails_v2_wiring(self): + def test_init_guardrails_v2_wiring(self, monkeypatch): """The guardrail registers and constructs via the config.yaml path.""" litellm.guardrail_name_config_map = {} - os.environ["REPELLOAI_API_KEY"] = "test-key" + monkeypatch.setenv("REPELLOAI_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ { diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index c8f22e6c15e..996a3ff0824 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -19,14 +19,14 @@ import litellm from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -def test_prompt_security_guard_config(): +def test_prompt_security_guard_config(monkeypatch): """Test guardrail initialization with proper configuration""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variables for testing - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") init_guardrails_v2( all_guardrails=[ @@ -78,10 +78,10 @@ def test_prompt_security_guard_config_no_api_key(): @pytest.mark.asyncio -async def test_apply_guardrail_block_request(): +async def test_apply_guardrail_block_request(monkeypatch): """Test that apply_guardrail blocks malicious prompts""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -132,10 +132,10 @@ async def test_apply_guardrail_block_request(): @pytest.mark.asyncio -async def test_apply_guardrail_modify_request(): +async def test_apply_guardrail_modify_request(monkeypatch): """Test that apply_guardrail modifies prompts when needed""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -183,10 +183,10 @@ async def test_apply_guardrail_modify_request(): @pytest.mark.asyncio -async def test_apply_guardrail_allow_request(): +async def test_apply_guardrail_allow_request(monkeypatch): """Test that apply_guardrail allows safe prompts""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -226,10 +226,10 @@ async def test_apply_guardrail_allow_request(): @pytest.mark.asyncio -async def test_apply_guardrail_block_response(): +async def test_apply_guardrail_block_response(monkeypatch): """Test that apply_guardrail blocks malicious responses""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -273,10 +273,10 @@ async def test_apply_guardrail_block_response(): @pytest.mark.asyncio -async def test_apply_guardrail_modify_response(): +async def test_apply_guardrail_modify_response(monkeypatch): """Test that apply_guardrail modifies responses when needed""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -317,10 +317,10 @@ async def test_apply_guardrail_modify_response(): @pytest.mark.asyncio -async def test_file_sanitization(): +async def test_file_sanitization(monkeypatch): """Test file sanitization for images""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -407,10 +407,10 @@ async def test_file_sanitization(): @pytest.mark.asyncio -async def test_file_sanitization_block(): +async def test_file_sanitization_block(monkeypatch): """Test that file sanitization blocks malicious files""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -491,10 +491,10 @@ async def test_file_sanitization_block(): @pytest.mark.asyncio -async def test_user_api_key_alias_forwarding(): +async def test_user_api_key_alias_forwarding(monkeypatch): """Test that user API key alias is properly sent via headers and payload""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -535,10 +535,10 @@ async def test_user_api_key_alias_forwarding(): @pytest.mark.asyncio -async def test_role_filtering(): +async def test_role_filtering(monkeypatch): """Test that tool/function messages are filtered out by default""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -600,11 +600,11 @@ async def test_role_filtering(): @pytest.mark.asyncio -async def test_check_tool_results_enabled(): +async def test_check_tool_results_enabled(monkeypatch): """Test with check_tool_results=True: transforms tool/function to 'other' role""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" - os.environ["PROMPT_SECURITY_CHECK_TOOL_RESULTS"] = "true" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 6c717d6f71c..13997fc4cd1 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -42,7 +42,7 @@ def time_controller(monkeypatch): @pytest.mark.asyncio -async def test_priority_weight_allocation(): +async def test_priority_weight_allocation(monkeypatch): """ Test that priority weights are correctly applied instead of equal splitting. @@ -53,7 +53,7 @@ async def test_priority_weight_allocation(): This validates the core fix where before it would split 50/50. """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up priority reservations litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -128,7 +128,7 @@ async def test_priority_weight_allocation(): @pytest.mark.asyncio -async def test_concurrent_priority_requests(): +async def test_concurrent_priority_requests(monkeypatch): """ Test the core issue: 5 concurrent requests with different priorities should get proper allocation based on priority weights, not equal splitting. @@ -136,7 +136,7 @@ async def test_concurrent_priority_requests(): This tests the exact scenario mentioned: priorities 0.9 and 0.1 should be 0.9/0.1, not 0.5/0.5. """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up the exact scenario from the issue litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -214,7 +214,7 @@ async def test_concurrent_priority_requests(): @pytest.mark.asyncio -async def test_100_concurrent_priority_requests(time_controller): +async def test_100_concurrent_priority_requests(time_controller, monkeypatch): """ Stress test: 100 concurrent requests with mixed priorities over 10 seconds. @@ -224,7 +224,7 @@ async def test_100_concurrent_priority_requests(time_controller): - Spread across 10 seconds to simulate real-world load """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up priority reservations litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -384,7 +384,7 @@ async def test_100_concurrent_priority_requests(time_controller): @pytest.mark.asyncio -async def test_concurrent_pre_call_hooks_stress(): +async def test_concurrent_pre_call_hooks_stress(monkeypatch): """ Stress test: 50 concurrent pre-call hooks with saturation-aware priority enforcement. @@ -394,7 +394,7 @@ async def test_concurrent_pre_call_hooks_stress(): Standard users (20% allocation) should have ~70% success rate with 30% random limiting. """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"premium": 0.8, "standard": 0.2} @@ -634,7 +634,7 @@ async def test_concurrent_pre_call_hooks_stress(): @pytest.mark.asyncio -async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): +async def test_fake_calls_case_1_no_rate_limiting_at_capacity(monkeypatch): """ Test Case 1: Saturation-Aware Rate Limiting at 50% Threshold @@ -650,7 +650,7 @@ async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): Once saturation hits 50%, strict mode enforces priority-based limits. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up priority reservations litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} @@ -759,7 +759,7 @@ async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): @pytest.mark.asyncio -async def test_fake_calls_case_2_priority_queue_during_saturation(): +async def test_fake_calls_case_2_priority_queue_during_saturation(monkeypatch): """ Test Case 2: Priority Queue Behavior During Saturation @@ -773,7 +773,7 @@ async def test_fake_calls_case_2_priority_queue_during_saturation(): When total traffic exceeds capacity, rate limiting enforces priority reservations. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} @@ -886,7 +886,7 @@ async def test_fake_calls_case_2_priority_queue_during_saturation(): @pytest.mark.asyncio -async def test_fake_calls_case_3_spillover_capacity_default_keys(): +async def test_fake_calls_case_3_spillover_capacity_default_keys(monkeypatch): """ Test Case 3: Spillover Capacity for Default Keys @@ -906,7 +906,7 @@ async def test_fake_calls_case_3_spillover_capacity_default_keys(): Tests spillover behavior where default keys share remaining capacity. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.75} litellm.priority_reservation_settings.default_priority = 0.25 @@ -1025,7 +1025,7 @@ async def test_fake_calls_case_3_spillover_capacity_default_keys(): @pytest.mark.asyncio -async def test_fake_calls_case_4_over_allocated_with_normalization(): +async def test_fake_calls_case_4_over_allocated_with_normalization(monkeypatch): """ Test Case 4: Over-Allocated Priority reservations with Normalization @@ -1042,7 +1042,7 @@ async def test_fake_calls_case_4_over_allocated_with_normalization(): - Due to concurrent burst, total successful may exceed 100 RPM in the test window - This test verifies normalization works and total capacity is reasonably bounded """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.60, "key_b": 0.80} @@ -1156,7 +1156,7 @@ async def test_fake_calls_case_4_over_allocated_with_normalization(): @pytest.mark.asyncio -async def test_fake_calls_case_5_default_value_priority_reservation(): +async def test_fake_calls_case_5_default_value_priority_reservation(monkeypatch): """ Test Case 5: Default value for priority reservation @@ -1176,7 +1176,7 @@ async def test_fake_calls_case_5_default_value_priority_reservation(): Tests complex scenario with explicit priorities and default priority. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.50, "key_b": 0.20, "key_c": 0.05} litellm.priority_reservation_settings.default_priority = 0.05 @@ -1296,7 +1296,7 @@ async def test_fake_calls_case_5_default_value_priority_reservation(): @pytest.mark.asyncio -async def test_default_priority_shared_pool(): +async def test_default_priority_shared_pool(monkeypatch): """ Test that keys without explicit priority share ONE default pool, not get individual allocations. @@ -1304,7 +1304,7 @@ async def test_default_priority_shared_pool(): - Key A, B, C (no priority) should share ONE 25 RPM pool - NOT get 25 RPM each (which would be 75 RPM total) """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"prod": 0.75} litellm.priority_reservation_settings.default_priority = 0.25 @@ -1382,7 +1382,7 @@ async def test_default_priority_shared_pool(): @pytest.mark.asyncio -async def test_async_log_success_event_increments_by_actual_tokens(): +async def test_async_log_success_event_increments_by_actual_tokens(monkeypatch): """ Test that async_log_success_event increments token counters by actual token usage. @@ -1394,7 +1394,7 @@ async def test_async_log_success_event_increments_by_actual_tokens(): from litellm.types.utils import ModelResponse, Usage - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"dev": 0.1, "prod": 0.9} dual_cache = DualCache() @@ -1483,7 +1483,7 @@ async def test_async_log_success_event_increments_by_actual_tokens(): @pytest.mark.asyncio -async def test_saturation_check_cache_ttl_configuration(): +async def test_saturation_check_cache_ttl_configuration(monkeypatch): """ Test that saturation_check_cache_ttl controls how long saturation values are cached locally. @@ -1492,7 +1492,7 @@ async def test_saturation_check_cache_ttl_configuration(): - After expiration, fresh values should be fetched from Redis - This prevents nodes from having stale saturation data in multi-node deployments """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set a short TTL for testing (5 seconds) original_ttl = litellm.priority_reservation_settings.saturation_check_cache_ttl @@ -1587,7 +1587,7 @@ async def test_saturation_check_cache_ttl_configuration(): @pytest.mark.asyncio -async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): +async def test_async_log_success_event_uses_team_priority_from_auth_metadata(monkeypatch): """ Test that async_log_success_event correctly retrieves priority from user_api_key_auth_metadata. @@ -1598,7 +1598,7 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): from litellm.types.utils import ModelResponse, Usage - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"team_priority": 0.8, "default": 0.2} dual_cache = DualCache() @@ -1680,7 +1680,7 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): @pytest.mark.asyncio -async def test_priority_429_includes_model_name_and_configured_limits(): +async def test_priority_429_includes_model_name_and_configured_limits(monkeypatch): """ The priority-based 429 should tell operators which model was hit and what the model's configured TPM/RPM are, so they can decide whether to tune the @@ -1694,7 +1694,7 @@ async def test_priority_429_includes_model_name_and_configured_limits(): """ from fastapi import HTTPException - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"prod": 0.5} dual_cache = DualCache() @@ -1774,7 +1774,7 @@ async def test_priority_429_includes_model_name_and_configured_limits(): @pytest.mark.asyncio -async def test_tpm_only_model_enforces_priority_and_model_capacity(): +async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch): """Regression: a model configured with ONLY tpm (no rpm) must still be rate limited. @@ -1789,7 +1789,7 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(): from litellm.types.utils import ModelResponse, Usage - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"dev": 0.25, "prod": 0.5} dual_cache = DualCache() diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py index 1c1e8eee145..97a986d1ade 100644 --- a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -189,7 +189,7 @@ async def test_batch_limiter_uses_atomic_check_and_increment(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): +async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(monkeypatch): """ DynamicRateLimitHandler PHASE 1 (read_only check) → PHASE 3 (increment) is non-atomic: dynamic_rate_limiter_v3.py:463-548. @@ -209,7 +209,7 @@ async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): # RPM + 1 successes before the next sees counter > RPM. MAX_SEQUENTIAL_SUCCESSES = MODEL_RPM + 1 - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"high": 0.9, "low": 0.1} dual_cache = DualCache() @@ -273,7 +273,7 @@ async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): +async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(monkeypatch): """ Regression test: dynamic limiter's enforced descriptors flow through `atomic_check_and_increment_by_n`, not the legacy @@ -283,7 +283,7 @@ async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): bundled into the atomic call alongside model_saturation_check. When not enforced, priority counter is incremented for tracking only. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"high": 0.9, "low": 0.1} dual_cache = DualCache() @@ -413,7 +413,7 @@ async def test_batch_zero_token_consumes_rpm_only(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): +async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(monkeypatch): """ Fail-closed guard: when atomic_check_and_increment_by_n returns overall_code=OVER_LIMIT but with a descriptor_key the dispatcher does @@ -425,7 +425,7 @@ async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): """ from fastapi import HTTPException - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"high": 0.9, "low": 0.1} dual_cache = DualCache() diff --git a/tests/test_litellm/test_add_deployment_no_master_key.py b/tests/test_litellm/test_add_deployment_no_master_key.py index 6db20d7d422..f7a0e90dad0 100644 --- a/tests/test_litellm/test_add_deployment_no_master_key.py +++ b/tests/test_litellm/test_add_deployment_no_master_key.py @@ -62,7 +62,7 @@ async def test_add_deployment_without_master_key(): @pytest.mark.asyncio -async def test_add_deployment_without_salt_key_or_master_key(): +async def test_add_deployment_without_salt_key_or_master_key(monkeypatch): """ Test that add_deployment() works when both master_key and LITELLM_SALT_KEY are None. @@ -70,55 +70,50 @@ async def test_add_deployment_without_salt_key_or_master_key(): such as in a local/dev environment or when just saving spend logs. """ # Remove LITELLM_SALT_KEY from environment - old_salt_key = os.environ.pop("LITELLM_SALT_KEY", None) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) - try: - # Set master_key to None - with patch("litellm.proxy.proxy_server.master_key", None): - # Mock the required dependencies - mock_prisma_client = MagicMock(spec=PrismaClient) - mock_prisma_client.db = MagicMock() - mock_prisma_client.db.litellm_config = MagicMock() - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=None + # Set master_key to None + with patch("litellm.proxy.proxy_server.master_key", None): + # Mock the required dependencies + mock_prisma_client = MagicMock(spec=PrismaClient) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_config = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock( + return_value=None + ) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + + # Create ProxyConfig instance + proxy_config = ProxyConfig() + + # Mock the internal methods + proxy_config._should_load_db_object = MagicMock(return_value=False) + proxy_config._init_non_llm_objects_in_db = AsyncMock() + + # This should NOT raise an exception + try: + await proxy_config.add_deployment( + prisma_client=mock_prisma_client, + proxy_logging_obj=mock_proxy_logging, ) - - mock_proxy_logging = MagicMock(spec=ProxyLogging) - - # Create ProxyConfig instance - proxy_config = ProxyConfig() - - # Mock the internal methods - proxy_config._should_load_db_object = MagicMock(return_value=False) - proxy_config._init_non_llm_objects_in_db = AsyncMock() - - # This should NOT raise an exception - try: - await proxy_config.add_deployment( - prisma_client=mock_prisma_client, - proxy_logging_obj=mock_proxy_logging, + assert True + except ValueError as e: + if "Master key is not initialized" in str( + e + ) or "Encryption key is not initialized" in str(e): + pytest.fail( + f"add_deployment raised ValueError about encryption key: {e}" ) - assert True - except ValueError as e: - if "Master key is not initialized" in str( - e - ) or "Encryption key is not initialized" in str(e): - pytest.fail( - f"add_deployment raised ValueError about encryption key: {e}" - ) - raise - except Exception as e: - if "Master key is not initialized" in str( - e - ) or "Encryption key is not initialized" in str(e): - pytest.fail( - f"add_deployment raised exception about encryption key: {e}" - ) - raise - finally: - # Restore LITELLM_SALT_KEY if it was set - if old_salt_key: - os.environ["LITELLM_SALT_KEY"] = old_salt_key + raise + except Exception as e: + if "Master key is not initialized" in str( + e + ) or "Encryption key is not initialized" in str(e): + pytest.fail( + f"add_deployment raised exception about encryption key: {e}" + ) + raise def test_add_deployment_sync_without_master_key(): diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/test_litellm/test_count_tokens_public_api.py index 1e2cf83dec0..ebd9c0c9edb 100644 --- a/tests/test_litellm/test_count_tokens_public_api.py +++ b/tests/test_litellm/test_count_tokens_public_api.py @@ -144,20 +144,16 @@ def test_acount_tokens_api_error_falls_back(): assert result.total_tokens > 0 -def test_acount_tokens_no_api_key_falls_back(): +def test_acount_tokens_no_api_key_falls_back(monkeypatch): """Test that missing API key falls back to local counting.""" - env_backup = os.environ.pop("OPENAI_API_KEY", None) - try: - result = asyncio.run( - litellm.acount_tokens( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - ) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + result = asyncio.run( + litellm.acount_tokens( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], ) + ) - # Should fall back to local tokenizer since no API key - assert result.total_tokens > 0 - assert result.tokenizer_type == "local_tokenizer" - finally: - if env_backup: - os.environ["OPENAI_API_KEY"] = env_backup + # Should fall back to local tokenizer since no API key + assert result.total_tokens > 0 + assert result.tokenizer_type == "local_tokenizer" diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index ba82bfaadc6..dd19334724d 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -318,7 +318,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp litellm.model_cost.pop(model_key, None) -def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(): +def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypatch): """Registering a custom override under a key shape that ``get_model_info`` cannot resolve (e.g. a triple provider prefix like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6``; a double @@ -338,7 +338,7 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(): from litellm.types.utils import PromptTokensDetailsWrapper, Usage original_model_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="") builtin_key = "us.anthropic.claude-sonnet-4-6" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index bd23ca11fbe..cd8dad39ad5 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -672,8 +672,8 @@ def test_all_model_configs(): ) == {"max_output_tokens": 10} -def test_anthropic_web_search_in_model_info(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_anthropic_web_search_in_model_info(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") supported_models = [ @@ -1193,11 +1193,11 @@ def test_max_tokens_consistency(): raise AssertionError(error_msg) -def test_get_model_info_gemini(): +def test_get_model_info_gemini(monkeypatch): """ Tests if ALL gemini models have 'tpm' and 'rpm' in the model info """ - 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_map = litellm.model_cost @@ -1252,8 +1252,8 @@ def test_get_model_info_bedrock_double_provider_prefix_resolves(local_model_cost assert info["key"] == "us.anthropic.claude-sonnet-4-6" -def test_openai_models_in_model_info(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_openai_models_in_model_info(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model_map = litellm.model_cost @@ -1408,7 +1408,7 @@ for commitment in BEDROCK_COMMITMENTS: print("block_list", block_list) -def test_supports_computer_use_utility(): +def test_supports_computer_use_utility(monkeypatch): """ Tests the litellm.utils.supports_computer_use utility function. """ @@ -1420,7 +1420,7 @@ def test_supports_computer_use_utility(): original_env_var = os.getenv("LITELLM_LOCAL_MODEL_COST_MAP") original_model_cost = getattr(litellm, "model_cost", None) - 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="") # Load with local/backup try: @@ -1438,7 +1438,7 @@ def test_supports_computer_use_utility(): if original_env_var is None: del os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env_var + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", original_env_var) if original_model_cost is not None: litellm.model_cost = original_model_cost @@ -1446,13 +1446,13 @@ def test_supports_computer_use_utility(): delattr(litellm, "model_cost") -def test_get_model_info_shows_supports_computer_use(): +def test_get_model_info_shows_supports_computer_use(monkeypatch): """ Tests if 'supports_computer_use' is correctly retrieved by get_model_info. We'll use 'claude-4-sonnet-20250514' as it's configured in the backup JSON to have supports_computer_use: True. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") # Ensure litellm.model_cost is loaded, relying on the backup mechanism if primary fails # as per previous debugging. litellm.model_cost = litellm.get_model_cost_map(url="")