mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
49da936efb
commit
693797420d
39 changed files with 462 additions and 484 deletions
|
|
@ -9,7 +9,7 @@
|
|||
"limit": 1078
|
||||
},
|
||||
"TQ004": {
|
||||
"limit": 757
|
||||
"limit": 557
|
||||
},
|
||||
"TQ005": {
|
||||
"limit": 2810
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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=[
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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="")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue