test: unwind environment writes in tests/test_litellm with monkeypatch (#37806)

* test: use monkeypatch.setenv for env writes in tests/test_litellm

`os.environ["X"] = v` inside a test leaks the value into every test that runs
after it in the same worker, so ordering decides the result. 262 of those
writes across 40 files now go through pytest's `monkeypatch` fixture, which
restores the previous value at teardown.

The rewrite skips any test that a mock.patch-family decorator wraps, any test
with defaulted positional parameters, any test whose own name is called
directly elsewhere, and rebinds nothing inside nested defs, because in each of
those cases appending a fixture parameter changes what pytest or mock binds.

Ratchets the TQ004 ceiling from 768 to 506.

* fix(test): delete the key through monkeypatch instead of popping it first

Five tests popped a key straight out of `os.environ`, ran, then restored it with
`monkeypatch.setenv`. By the time monkeypatch saw the name it was already gone,
so it recorded "absent" as the value to go back to and deleted the key at
teardown. On a worker that inherited a real `RESEND_API_KEY`, `SENDGRID_API_KEY`,
`UI_PASSWORD`, `LITELLM_SALT_KEY` or `OPENAI_API_KEY`, every test after the first
one ran without it.

`monkeypatch.delenv(..., raising=False)` removes the key and restores whatever
was there, so the try/finally the manual restore needed goes with it.

* chore(test): leave the two cost-calc files to the PR that rewrites them fully

Both files are also in #37815, which converts the module-global writes as well
as the env writes and folds them into one fixture. Two PRs rewriting the same
lines differently is a conflict nobody benefits from resolving, so this one
drops back to staging on those two and keeps the other 39.

TQ004 clears 200 here instead of 275; the rest moves with #37815.
This commit is contained in:
yuneng-jiang 2026-08-21 20:28:37 -07:00 committed by GitHub
parent 49da936efb
commit 693797420d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 462 additions and 484 deletions

View file

@ -9,7 +9,7 @@
"limit": 1078
},
"TQ004": {
"limit": 757
"limit": 557
},
"TQ005": {
"limit": 2810

View file

@ -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"]

View file

@ -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 (

View file

@ -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 = "<p>Test email body</p>"
# Test data
from_email = "test@example.com"
to_email = ["recipient@example.com"]
subject = "Test Subject"
html_body = "<p>Test email body</p>"
# 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

View file

@ -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="<p>Test email body</p>",
)
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="<p>Test email body</p>",
)
@pytest.mark.asyncio

View file

@ -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

View file

@ -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 = {

View file

@ -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,

View file

@ -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"

View file

@ -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():

View file

@ -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"]

View file

@ -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",

View file

@ -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",

View file

@ -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,

View file

@ -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")

View file

@ -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():

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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"]

View file

@ -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(

View file

@ -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"

View file

@ -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()

View file

@ -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()

View file

@ -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"

View file

@ -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(

View file

@ -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

View file

@ -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",

View file

@ -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

View file

@ -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=[

View file

@ -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

View file

@ -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=[
{

View file

@ -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

View file

@ -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()

View file

@ -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()

View file

@ -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():

View file

@ -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"

View file

@ -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"

View file

@ -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="")