mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(tests): use canonical litellm_enterprise import path (#27699)
The enterprise package is installed as `litellm_enterprise` (per enterprise/pyproject.toml), but several tests imported it as `enterprise.litellm_enterprise.*` — a path that only resolves because the repo root happens to sit on sys.path, letting Python's implicit namespace package machinery discover `enterprise/` as a directory. This breaks any test runner that relocates source (e.g. the mutation-testing workflow, which copies tests under `mutants/`) and also caused two `patch()` strings to target a module path that does not match what production code imports — meaning those mocks were never actually patching the production module's attribute. Replace `from enterprise.litellm_enterprise.` with the canonical `from litellm_enterprise.` across 6 test files, and fix two `patch()` target strings (and one `sys.modules` patch key in the SSO test) to match.
This commit is contained in:
ryan-crabbe-berri
GitHub
parent
9c4faeabc9
commit
63a2d1ddc9
7 changed files with 406 additions and 202 deletions
|
|
@ -19,7 +19,7 @@ def client():
|
|||
def mock_user_api_key_auth():
|
||||
"""Mock the user_api_key_auth dependency"""
|
||||
with patch(
|
||||
"enterprise.litellm_enterprise.proxy.management_endpoints.internal_user_endpoints.user_api_key_auth"
|
||||
"litellm_enterprise.proxy.management_endpoints.internal_user_endpoints.user_api_key_auth"
|
||||
) as mock_auth:
|
||||
mock_auth.return_value = {"user_id": "test_user", "api_key": "test_key"}
|
||||
yield mock_auth
|
||||
|
|
@ -31,12 +31,16 @@ class TestAvailableEnterpriseUsers:
|
|||
self, client, mock_user_api_key_auth
|
||||
):
|
||||
"""Test when max_users is set and user count is within limit"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user_data",
|
||||
{"max_users": 10},
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user_data",
|
||||
{"max_users": 10},
|
||||
),
|
||||
):
|
||||
# Mock database count
|
||||
mock_prisma.db.litellm_usertable.count = AsyncMock(return_value=5)
|
||||
|
|
@ -66,12 +70,16 @@ class TestAvailableEnterpriseUsers:
|
|||
self, client, mock_user_api_key_auth
|
||||
):
|
||||
"""Test when max_users is not set (premium_user_data is None or doesn't contain max_users)"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user_data",
|
||||
None,
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user_data",
|
||||
None,
|
||||
),
|
||||
):
|
||||
# Mock database count
|
||||
mock_prisma.db.litellm_usertable.count = AsyncMock(return_value=3)
|
||||
|
|
@ -99,12 +107,16 @@ class TestAvailableEnterpriseUsers:
|
|||
self, client, mock_user_api_key_auth
|
||||
):
|
||||
"""Test the current bug where total_users_remaining can be negative"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user_data",
|
||||
{"key": "value"},
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user_data",
|
||||
{"key": "value"},
|
||||
),
|
||||
):
|
||||
# Mock database count higher than max_users to trigger the bug
|
||||
mock_prisma.db.litellm_usertable.count = AsyncMock(return_value=8)
|
||||
|
|
@ -140,12 +152,15 @@ class TestAvailableEnterpriseUsers:
|
|||
"""Test when prisma_client is None (no database connection)"""
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
None,
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
None,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.premium_user",
|
||||
True,
|
||||
),
|
||||
):
|
||||
# Override the dependency
|
||||
client.app.dependency_overrides[mock_user_api_key_auth] = lambda: {
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from unittest.mock import patch
|
|||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from enterprise.litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
|
||||
BaseEmailLogger,
|
||||
)
|
||||
|
||||
|
|
@ -30,6 +30,7 @@ def no_invitation_wait(monkeypatch):
|
|||
|
||||
monkeypatch.setattr(BaseEmailLogger, "_wait_for_invitation_creation", _noop)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_email_logger():
|
||||
return BaseEmailLogger()
|
||||
|
|
@ -283,7 +284,10 @@ async def test_send_key_created_email_without_key(
|
|||
mock_send_email.assert_called_once()
|
||||
call_args = mock_send_email.call_args[1]
|
||||
assert "sk-secret-key-456" not in call_args["html_body"]
|
||||
assert "[Key hidden for security - retrieve from dashboard]" in call_args["html_body"]
|
||||
assert (
|
||||
"[Key hidden for security - retrieve from dashboard]"
|
||||
in call_args["html_body"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -317,7 +321,10 @@ async def test_send_key_rotated_email_without_key(
|
|||
mock_send_email.assert_called_once()
|
||||
call_args = mock_send_email.call_args[1]
|
||||
assert "sk-secret-rotated-789" not in call_args["html_body"]
|
||||
assert "[Key hidden for security - retrieve from dashboard]" in call_args["html_body"]
|
||||
assert (
|
||||
"[Key hidden for security - retrieve from dashboard]"
|
||||
in call_args["html_body"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -371,52 +378,52 @@ async def test_get_invitation_link_creates_new_when_none_exist(base_email_logger
|
|||
"""Test that _get_invitation_link creates a new invitation when none exist"""
|
||||
# Mock prisma client with no existing invitation rows
|
||||
mock_prisma = mock.MagicMock()
|
||||
|
||||
|
||||
# Mock find_many to return empty list (no existing invitations)
|
||||
async def mock_find_many_empty(*args, **kwargs):
|
||||
return []
|
||||
|
||||
|
||||
mock_prisma.db.litellm_invitationlink.find_many = mock_find_many_empty
|
||||
|
||||
|
||||
# Mock the create_invitation_for_user function
|
||||
mock_created_invitation = mock.MagicMock()
|
||||
mock_created_invitation.id = "new-invitation-id"
|
||||
|
||||
|
||||
with mock.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
with mock.patch(
|
||||
"litellm.proxy.management_helpers.user_invitation.create_invitation_for_user",
|
||||
return_value=mock_created_invitation
|
||||
return_value=mock_created_invitation,
|
||||
) as mock_create_invitation:
|
||||
# Execute
|
||||
result = await base_email_logger._get_invitation_link(
|
||||
user_id="test-user", base_url="http://test.com"
|
||||
)
|
||||
|
||||
|
||||
# Verify that create_invitation_for_user was called
|
||||
mock_create_invitation.assert_called_once()
|
||||
call_args = mock_create_invitation.call_args[1]
|
||||
assert call_args["data"].user_id == "test-user"
|
||||
assert call_args["user_api_key_dict"].user_id == "test-user"
|
||||
|
||||
|
||||
# Verify the returned link uses the new invitation ID
|
||||
assert result == "http://test.com/ui?invitation_id=new-invitation-id"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_invitation_link_uses_existing_when_available(base_email_logger):
|
||||
"""Test that _get_invitation_link uses existing invitation when available"""
|
||||
# Mock prisma client with existing invitation row
|
||||
mock_invitation_row = mock.MagicMock()
|
||||
mock_invitation_row.id = "existing-invitation-id"
|
||||
|
||||
|
||||
mock_prisma = mock.MagicMock()
|
||||
|
||||
|
||||
# Mock find_many to return existing invitation
|
||||
async def mock_find_many_existing(*args, **kwargs):
|
||||
return [mock_invitation_row]
|
||||
|
||||
|
||||
mock_prisma.db.litellm_invitationlink.find_many = mock_find_many_existing
|
||||
|
||||
|
||||
with mock.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
with mock.patch(
|
||||
"litellm.proxy.management_helpers.user_invitation.create_invitation_for_user"
|
||||
|
|
@ -425,10 +432,10 @@ async def test_get_invitation_link_uses_existing_when_available(base_email_logge
|
|||
result = await base_email_logger._get_invitation_link(
|
||||
user_id="test-user", base_url="http://test.com"
|
||||
)
|
||||
|
||||
|
||||
# Verify that create_invitation_for_user was NOT called
|
||||
mock_create_invitation.assert_not_called()
|
||||
|
||||
|
||||
# Verify the returned link uses the existing invitation ID
|
||||
assert result == "http://test.com/ui?invitation_id=existing-invitation-id"
|
||||
|
||||
|
|
@ -438,33 +445,33 @@ async def test_get_invitation_link_creates_new_when_list_is_none(base_email_logg
|
|||
"""Test that _get_invitation_link creates a new invitation when invitation_rows is None"""
|
||||
# Mock prisma client to return None
|
||||
mock_prisma = mock.MagicMock()
|
||||
|
||||
|
||||
# Mock find_many to return None
|
||||
async def mock_find_many_none(*args, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
mock_prisma.db.litellm_invitationlink.find_many = mock_find_many_none
|
||||
|
||||
|
||||
# Mock the create_invitation_for_user function
|
||||
mock_created_invitation = mock.MagicMock()
|
||||
mock_created_invitation.id = "new-invitation-from-none"
|
||||
|
||||
|
||||
with mock.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
with mock.patch(
|
||||
"litellm.proxy.management_helpers.user_invitation.create_invitation_for_user",
|
||||
return_value=mock_created_invitation
|
||||
return_value=mock_created_invitation,
|
||||
) as mock_create_invitation:
|
||||
# Execute
|
||||
result = await base_email_logger._get_invitation_link(
|
||||
user_id="test-user", base_url="http://test.com"
|
||||
)
|
||||
|
||||
|
||||
# Verify that create_invitation_for_user was called
|
||||
mock_create_invitation.assert_called_once()
|
||||
call_args = mock_create_invitation.call_args[1]
|
||||
assert call_args["data"].user_id == "test-user"
|
||||
assert call_args["user_api_key_dict"].user_id == "test-user"
|
||||
|
||||
|
||||
# Verify the returned link uses the new invitation ID
|
||||
assert result == "http://test.com/ui?invitation_id=new-invitation-from-none"
|
||||
|
||||
|
|
@ -495,13 +502,15 @@ async def test_get_email_params_user_invitation(
|
|||
user_email="test@example.com",
|
||||
)
|
||||
|
||||
assert result.logo_url == "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
|
||||
assert (
|
||||
result.logo_url
|
||||
== "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
|
||||
)
|
||||
assert result.support_contact == "support@berri.ai"
|
||||
assert result.base_url == "http://test.com/ui?invitation_id=test-id"
|
||||
assert result.recipient_email == "test@example.com"
|
||||
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_env_vars(monkeypatch):
|
||||
"""Set up test environment variables"""
|
||||
|
|
@ -513,69 +522,74 @@ def mock_env_vars(monkeypatch):
|
|||
monkeypatch.setenv("PROXY_BASE_URL", "http://test.com")
|
||||
monkeypatch.setenv("PROXY_API_URL", "https://test.com")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_email_params_custom_templates_premium_user(mock_env_vars):
|
||||
"""Test that _get_email_params returns correct values with custom templates for premium users"""
|
||||
# Mock premium_user as True
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
email_logger = BaseEmailLogger()
|
||||
|
||||
|
||||
# Test invitation email params
|
||||
invitation_params = await email_logger._get_email_params(
|
||||
email_event=EmailEvent.new_user_invitation,
|
||||
user_id="testid",
|
||||
user_email="test@example.com",
|
||||
event_message="New User Invitation"
|
||||
event_message="New User Invitation",
|
||||
)
|
||||
|
||||
|
||||
assert invitation_params.subject == "Welcome to Test Company!"
|
||||
assert invitation_params.signature == "Best regards,\nTest Company Team"
|
||||
assert invitation_params.logo_url == "https://test-company.com/logo.png"
|
||||
assert invitation_params.support_contact == "support@test-company.com"
|
||||
assert invitation_params.base_url == "http://test.com"
|
||||
|
||||
|
||||
# Test key created email params
|
||||
key_params = await email_logger._get_email_params(
|
||||
email_event=EmailEvent.virtual_key_created,
|
||||
user_id="testid",
|
||||
user_email="test@example.com",
|
||||
event_message="API Key Created"
|
||||
event_message="API Key Created",
|
||||
)
|
||||
|
||||
|
||||
assert key_params.subject == "Your Test Company API Key"
|
||||
assert key_params.signature == "Best regards,\nTest Company Team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_email_params_non_premium_user(mock_env_vars):
|
||||
"""Test that non-premium users get default templates even when custom ones are provided"""
|
||||
# Mock premium_user as False
|
||||
with patch("litellm.proxy.proxy_server.premium_user", False):
|
||||
email_logger = BaseEmailLogger()
|
||||
|
||||
|
||||
# Test invitation email params
|
||||
email_params = await email_logger._get_email_params(
|
||||
email_event=EmailEvent.new_user_invitation,
|
||||
user_email="test@example.com",
|
||||
event_message="New User Invitation"
|
||||
event_message="New User Invitation",
|
||||
)
|
||||
|
||||
|
||||
# Should use default values even though custom values are set in env
|
||||
assert email_params.subject == "LiteLLM: New User Invitation"
|
||||
assert email_params.signature == EMAIL_FOOTER
|
||||
assert email_params.logo_url == "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
|
||||
assert (
|
||||
email_params.logo_url
|
||||
== "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
|
||||
)
|
||||
assert email_params.support_contact == "support@berri.ai"
|
||||
|
||||
|
||||
# Test key created email params
|
||||
key_params = await email_logger._get_email_params(
|
||||
email_event=EmailEvent.virtual_key_created,
|
||||
user_email="test@example.com",
|
||||
event_message="API Key Created"
|
||||
event_message="API Key Created",
|
||||
)
|
||||
|
||||
|
||||
assert key_params.subject == "LiteLLM: API Key Created"
|
||||
assert key_params.signature == EMAIL_FOOTER
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_email_params_default_templates(monkeypatch):
|
||||
"""Test that _get_email_params uses default templates when custom ones aren't provided"""
|
||||
|
|
@ -583,28 +597,28 @@ async def test_get_email_params_default_templates(monkeypatch):
|
|||
monkeypatch.delenv("EMAIL_SUBJECT_INVITATION", raising=False)
|
||||
monkeypatch.delenv("EMAIL_SUBJECT_KEY_CREATED", raising=False)
|
||||
monkeypatch.delenv("EMAIL_SIGNATURE", raising=False)
|
||||
|
||||
|
||||
# Mock premium_user as True (shouldn't matter since no custom values are set)
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
email_logger = BaseEmailLogger()
|
||||
|
||||
|
||||
# Test invitation email params with default template
|
||||
invitation_params = await email_logger._get_email_params(
|
||||
email_event=EmailEvent.new_user_invitation,
|
||||
user_email="test@example.com",
|
||||
event_message="New User Invitation"
|
||||
event_message="New User Invitation",
|
||||
)
|
||||
|
||||
|
||||
assert invitation_params.subject == "LiteLLM: New User Invitation"
|
||||
assert invitation_params.signature == EMAIL_FOOTER
|
||||
|
||||
|
||||
# Test key created email params with default template
|
||||
key_params = await email_logger._get_email_params(
|
||||
email_event=EmailEvent.virtual_key_created,
|
||||
user_email="test@example.com",
|
||||
event_message="API Key Created"
|
||||
event_message="API Key Created",
|
||||
)
|
||||
|
||||
|
||||
assert key_params.subject == "LiteLLM: API Key Created"
|
||||
assert key_params.signature == EMAIL_FOOTER
|
||||
|
||||
|
|
@ -639,7 +653,10 @@ async def test_send_soft_budget_alert_email(
|
|||
call_args = mock_send_email.call_args[1]
|
||||
assert call_args["from_email"] == BaseEmailLogger.DEFAULT_LITELLM_EMAIL
|
||||
assert call_args["to_email"] == ["test@example.com"]
|
||||
assert call_args["subject"] == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0"
|
||||
assert (
|
||||
call_args["subject"]
|
||||
== "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0"
|
||||
)
|
||||
assert "$100.0" in call_args["html_body"] # soft_budget
|
||||
assert "$105.0" in call_args["html_body"] # spend
|
||||
assert "$200.0" in call_args["html_body"] # max_budget
|
||||
|
|
@ -673,13 +690,13 @@ async def test_send_soft_budget_alert_email_no_max_budget(
|
|||
call_args = mock_send_email.call_args[1]
|
||||
assert "$100.0" in call_args["html_body"] # soft_budget
|
||||
assert "$105.0" in call_args["html_body"] # spend
|
||||
assert "Maximum Budget" not in call_args["html_body"] # max_budget should not be shown
|
||||
assert (
|
||||
"Maximum Budget" not in call_args["html_body"]
|
||||
) # max_budget should not be shown
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_alerts_soft_budget_crossed(
|
||||
base_email_logger, mock_send_email
|
||||
):
|
||||
async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_email):
|
||||
"""Test that budget_alerts sends email when soft budget is crossed"""
|
||||
user_info = CallInfo(
|
||||
user_id="test_user",
|
||||
|
|
@ -708,11 +725,14 @@ async def test_budget_alerts_soft_budget_crossed(
|
|||
mock_send_email.assert_called_once()
|
||||
call_args = mock_send_email.call_args[1]
|
||||
assert call_args["to_email"] == ["test@example.com"]
|
||||
|
||||
|
||||
# Verify cache was set to prevent duplicate alerts
|
||||
mock_cache.async_set_cache.assert_called_once()
|
||||
cache_call_args = mock_cache.async_set_cache.call_args[1]
|
||||
assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:test_user"
|
||||
assert (
|
||||
cache_call_args["key"]
|
||||
== "email_budget_alerts:soft_budget_crossed:test_user"
|
||||
)
|
||||
assert cache_call_args["value"] == "SENT"
|
||||
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
|
||||
|
||||
|
|
@ -766,9 +786,7 @@ async def test_budget_alerts_soft_budget_duplicate_prevention(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_alerts_no_budgets(
|
||||
base_email_logger, mock_send_email
|
||||
):
|
||||
async def test_budget_alerts_no_budgets(base_email_logger, mock_send_email):
|
||||
"""Test that budget_alerts returns early when no budgets are set"""
|
||||
user_info = CallInfo(
|
||||
user_id="test_user",
|
||||
|
|
@ -817,7 +835,10 @@ async def test_budget_alerts_uses_token_for_cache_key(
|
|||
# Verify cache key uses token instead of user_id
|
||||
mock_cache.async_set_cache.assert_called_once()
|
||||
cache_call_args = mock_cache.async_set_cache.call_args[1]
|
||||
assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:hashed_token_123"
|
||||
assert (
|
||||
cache_call_args["key"]
|
||||
== "email_budget_alerts:soft_budget_crossed:hashed_token_123"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -838,7 +859,9 @@ async def test_get_email_params_soft_budget_crossed(
|
|||
)
|
||||
|
||||
# Should use default subject template for soft_budget_crossed
|
||||
assert result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0"
|
||||
assert (
|
||||
result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0"
|
||||
)
|
||||
assert result.recipient_email == "test@example.com"
|
||||
assert result.base_url == "http://test.com"
|
||||
|
||||
|
|
@ -867,16 +890,20 @@ async def test_budget_alerts_max_budget_alert_crossed(
|
|||
"PROXY_BASE_URL": "http://test.com",
|
||||
},
|
||||
):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
mock_send_email.assert_called_once()
|
||||
call_args = mock_send_email.call_args[1]
|
||||
assert call_args["to_email"] == ["test@example.com"]
|
||||
assert "Max Budget Alert" in call_args["subject"]
|
||||
|
||||
|
||||
mock_cache.async_set_cache.assert_called_once()
|
||||
cache_call_args = mock_cache.async_set_cache.call_args[1]
|
||||
assert cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user"
|
||||
assert (
|
||||
cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user"
|
||||
)
|
||||
assert cache_call_args["value"] == "SENT"
|
||||
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
|
||||
|
||||
|
|
@ -906,15 +933,15 @@ async def test_multi_threshold_sends_crossed_thresholds(
|
|||
base_email_logger.internal_usage_cache = mock_cache
|
||||
|
||||
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
# spend=80 crosses 50% ($50) and 75% ($75), but not 100% ($100)
|
||||
assert mock_send_email.call_count == 2
|
||||
|
||||
# Check cache keys include threshold percentage
|
||||
cache_keys = [
|
||||
c[1]["key"] for c in mock_cache.async_set_cache.call_args_list
|
||||
]
|
||||
cache_keys = [c[1]["key"] for c in mock_cache.async_set_cache.call_args_list]
|
||||
assert "email_budget_alerts:max_budget_alert:50:hashed_key_1" in cache_keys
|
||||
assert "email_budget_alerts:max_budget_alert:75:hashed_key_1" in cache_keys
|
||||
|
||||
|
|
@ -949,7 +976,9 @@ async def test_multi_threshold_dedup_cache_prevents_resend(
|
|||
base_email_logger.internal_usage_cache = mock_cache
|
||||
|
||||
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
# Only 75% should fire
|
||||
assert mock_send_email.call_count == 1
|
||||
|
|
@ -980,7 +1009,9 @@ async def test_multi_threshold_owner_email_auto_included(
|
|||
base_email_logger.internal_usage_cache = mock_cache
|
||||
|
||||
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
mock_send_email.assert_called_once()
|
||||
to_emails = mock_send_email.call_args[1]["to_email"]
|
||||
|
|
@ -1002,7 +1033,7 @@ async def test_multi_threshold_malformed_keys_skipped(
|
|||
event_group=Litellm_EntityType.KEY,
|
||||
max_budget_alert_emails={
|
||||
"fifty": ["finance@co.com"], # invalid
|
||||
"50": ["finance@co.com"], # valid, crossed
|
||||
"50": ["finance@co.com"], # valid, crossed
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -1012,7 +1043,9 @@ async def test_multi_threshold_malformed_keys_skipped(
|
|||
base_email_logger.internal_usage_cache = mock_cache
|
||||
|
||||
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
# Only the valid "50" threshold should fire
|
||||
assert mock_send_email.call_count == 1
|
||||
|
|
@ -1041,7 +1074,9 @@ async def test_multi_threshold_empty_emails_only_owner(
|
|||
base_email_logger.internal_usage_cache = mock_cache
|
||||
|
||||
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
mock_send_email.assert_called_once()
|
||||
to_emails = mock_send_email.call_args[1]["to_email"]
|
||||
|
|
@ -1067,11 +1102,13 @@ async def test_no_map_preserves_old_single_threshold(
|
|||
base_email_logger.internal_usage_cache = mock_cache
|
||||
|
||||
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
|
||||
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
|
||||
await base_email_logger.budget_alerts(
|
||||
type="max_budget_alert", user_info=user_info
|
||||
)
|
||||
|
||||
mock_send_email.assert_called_once()
|
||||
call_args = mock_send_email.call_args[1]
|
||||
assert call_args["to_email"] == ["test@example.com"]
|
||||
# Old path cache key has no threshold percentage
|
||||
cache_key = mock_cache.async_set_cache.call_args[1]["key"]
|
||||
assert cache_key == "email_budget_alerts:max_budget_alert:test_user"
|
||||
assert cache_key == "email_budget_alerts:max_budget_alert:test_user"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from enterprise.litellm_enterprise.enterprise_callbacks.callback_controls import (
|
||||
from litellm_enterprise.enterprise_callbacks.callback_controls import (
|
||||
EnterpriseCallbackControls,
|
||||
)
|
||||
from litellm.constants import X_LITELLM_DISABLE_CALLBACKS
|
||||
|
|
@ -18,168 +18,282 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
|
||||
|
||||
class TestEnterpriseCallbackControls:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_premium_user(self):
|
||||
"""Fixture to mock premium user check as True"""
|
||||
with patch.object(EnterpriseCallbackControls, '_should_allow_dynamic_callback_disabling', return_value=True):
|
||||
with patch.object(
|
||||
EnterpriseCallbackControls,
|
||||
"_should_allow_dynamic_callback_disabling",
|
||||
return_value=True,
|
||||
):
|
||||
yield
|
||||
|
||||
@pytest.fixture
|
||||
|
||||
@pytest.fixture
|
||||
def mock_non_premium_user(self):
|
||||
"""Fixture to mock premium user check as False"""
|
||||
with patch.object(EnterpriseCallbackControls, '_should_allow_dynamic_callback_disabling', return_value=False):
|
||||
with patch.object(
|
||||
EnterpriseCallbackControls,
|
||||
"_should_allow_dynamic_callback_disabling",
|
||||
return_value=False,
|
||||
):
|
||||
yield
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request_headers(self):
|
||||
"""Fixture to mock get_proxy_server_request_headers"""
|
||||
with patch('enterprise.litellm_enterprise.enterprise_callbacks.callback_controls.get_proxy_server_request_headers') as mock_headers:
|
||||
with patch(
|
||||
"litellm_enterprise.enterprise_callbacks.callback_controls.get_proxy_server_request_headers"
|
||||
) as mock_headers:
|
||||
yield mock_headers
|
||||
|
||||
def test_callback_disabled_langfuse_string(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_langfuse_string(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that 'langfuse' string callback is disabled when specified in headers"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_callback_disabled_langfuse_customlogger(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_langfuse_customlogger(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that LangfusePromptManagement CustomLogger instance is disabled when 'langfuse' specified in headers"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
langfuse_logger = LangfusePromptManagement()
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(langfuse_logger, litellm_params, standard_callback_dynamic_params)
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
langfuse_logger, litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_callback_disabled_s3_v2_string(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_s3_v2_string(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that 's3_v2' string callback is disabled when specified in headers"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "s3_v2"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("s3_v2", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"s3_v2", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_callback_disabled_s3_v2_customlogger(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_s3_v2_customlogger(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that S3Logger CustomLogger instance is disabled when 's3_v2' specified in headers"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "s3_v2"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# Mock S3Logger to avoid async initialization issues
|
||||
with patch('litellm.integrations.s3_v2.S3Logger.__init__', return_value=None):
|
||||
with patch("litellm.integrations.s3_v2.S3Logger.__init__", return_value=None):
|
||||
s3_logger = S3Logger()
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(s3_logger, litellm_params, standard_callback_dynamic_params)
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
s3_logger, litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_callback_disabled_datadog_string(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_datadog_string(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that 'datadog' string callback is disabled when specified in headers"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "datadog"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("datadog", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"datadog", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_callback_disabled_datadog_customlogger(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_datadog_customlogger(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that DataDogLogger CustomLogger instance is disabled when 'datadog' specified in headers"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "datadog"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# Mock DataDogLogger to avoid async initialization issues
|
||||
with patch('litellm.integrations.datadog.datadog.DataDogLogger.__init__', return_value=None):
|
||||
with patch(
|
||||
"litellm.integrations.datadog.datadog.DataDogLogger.__init__",
|
||||
return_value=None,
|
||||
):
|
||||
datadog_logger = DataDogLogger()
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(datadog_logger, litellm_params, standard_callback_dynamic_params)
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
datadog_logger, litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_multiple_callbacks_disabled(self, mock_premium_user, mock_request_headers):
|
||||
"""Test that multiple callbacks can be disabled with comma-separated list"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse,datadog,s3_v2"}
|
||||
mock_request_headers.return_value = {
|
||||
X_LITELLM_DISABLE_CALLBACKS: "langfuse,datadog,s3_v2"
|
||||
}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
# Test each callback is disabled
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("datadog", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("s3_v2", litellm_params, standard_callback_dynamic_params) is True
|
||||
|
||||
# Test non-disabled callback is not disabled
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("prometheus", litellm_params, standard_callback_dynamic_params) is False
|
||||
|
||||
def test_callback_not_disabled_when_not_in_list(self, mock_premium_user, mock_request_headers):
|
||||
# Test each callback is disabled
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"datadog", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"s3_v2", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
# Test non-disabled callback is not disabled
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"prometheus", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_callback_not_disabled_when_not_in_list(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that callbacks not in the disabled list are not disabled"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("datadog", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"datadog", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_callback_not_disabled_when_no_header(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_not_disabled_when_no_header(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that callbacks are not disabled when the header is not present"""
|
||||
mock_request_headers.return_value = {}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_callback_not_disabled_when_header_none(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_not_disabled_when_header_none(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that callbacks are not disabled when the header value is None"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: None}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_non_premium_user_cannot_disable_callbacks(self, mock_non_premium_user, mock_request_headers):
|
||||
def test_non_premium_user_cannot_disable_callbacks(
|
||||
self, mock_non_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that non-premium users cannot disable callbacks even with the header"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_case_insensitive_callback_matching(self, mock_premium_user, mock_request_headers):
|
||||
def test_case_insensitive_callback_matching(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that callback matching is case insensitive"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "LANGFUSE,DataDog"}
|
||||
mock_request_headers.return_value = {
|
||||
X_LITELLM_DISABLE_CALLBACKS: "LANGFUSE,DataDog"
|
||||
}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# Test lowercase callbacks are disabled
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("datadog", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"datadog", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_whitespace_handling_in_disabled_callbacks(self, mock_premium_user, mock_request_headers):
|
||||
def test_whitespace_handling_in_disabled_callbacks(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that whitespace around callback names is handled correctly"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: " langfuse , datadog , s3_v2 "}
|
||||
mock_request_headers.return_value = {
|
||||
X_LITELLM_DISABLE_CALLBACKS: " langfuse , datadog , s3_v2 "
|
||||
}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("datadog", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("s3_v2", litellm_params, standard_callback_dynamic_params) is True
|
||||
|
||||
def test_custom_logger_not_in_registry(self, mock_premium_user, mock_request_headers):
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"datadog", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"s3_v2", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_custom_logger_not_in_registry(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that CustomLogger not in registry is not disabled"""
|
||||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "unknown_logger"}
|
||||
mock_request_headers.return_value = {
|
||||
X_LITELLM_DISABLE_CALLBACKS: "unknown_logger"
|
||||
}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# Create a mock CustomLogger that's not in the registry
|
||||
class UnknownLogger(CustomLogger):
|
||||
pass
|
||||
|
||||
|
||||
unknown_logger = UnknownLogger()
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(unknown_logger, litellm_params, standard_callback_dynamic_params)
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
unknown_logger, litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_exception_handling(self, mock_premium_user, mock_request_headers):
|
||||
|
|
@ -188,32 +302,64 @@ class TestEnterpriseCallbackControls:
|
|||
mock_request_headers.side_effect = Exception("Test exception")
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_callback_disabled_via_request_body_langfuse(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_via_request_body_langfuse(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that callbacks can be disabled via request body litellm_disabled_callbacks"""
|
||||
mock_request_headers.return_value = {} # No headers
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams(litellm_disabled_callbacks=["langfuse"])
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams(
|
||||
litellm_disabled_callbacks=["langfuse"]
|
||||
)
|
||||
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_callback_disabled_via_request_body_multiple(self, mock_premium_user, mock_request_headers):
|
||||
def test_callback_disabled_via_request_body_multiple(
|
||||
self, mock_premium_user, mock_request_headers
|
||||
):
|
||||
"""Test that multiple callbacks can be disabled via request body"""
|
||||
mock_request_headers.return_value = {} # No headers
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams(litellm_disabled_callbacks=["langfuse", "datadog", "s3_v2"])
|
||||
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams(
|
||||
litellm_disabled_callbacks=["langfuse", "datadog", "s3_v2"]
|
||||
)
|
||||
|
||||
# Test each callback is disabled
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("datadog", litellm_params, standard_callback_dynamic_params) is True
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("s3_v2", litellm_params, standard_callback_dynamic_params) is True
|
||||
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"datadog", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"s3_v2", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
# Test non-disabled callback is not disabled
|
||||
assert EnterpriseCallbackControls.is_callback_disabled_dynamically("prometheus", litellm_params, standard_callback_dynamic_params) is False
|
||||
assert (
|
||||
EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"prometheus", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_admin_can_disable_dynamic_callback_disabling(self, mock_request_headers):
|
||||
"""
|
||||
|
|
@ -223,11 +369,13 @@ class TestEnterpriseCallbackControls:
|
|||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# Mock litellm.allow_dynamic_callback_disabling set to False
|
||||
with patch('litellm.allow_dynamic_callback_disabling', False):
|
||||
with patch('litellm.proxy.proxy_server.premium_user', True):
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
with patch("litellm.allow_dynamic_callback_disabling", False):
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is False
|
||||
|
||||
def test_admin_can_enable_dynamic_callback_disabling(self, mock_request_headers):
|
||||
|
|
@ -238,14 +386,18 @@ class TestEnterpriseCallbackControls:
|
|||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# Mock litellm.allow_dynamic_callback_disabling set to True
|
||||
with patch('litellm.allow_dynamic_callback_disabling', True):
|
||||
with patch('litellm.proxy.proxy_server.premium_user', True):
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
with patch("litellm.allow_dynamic_callback_disabling", True):
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_default_admin_setting_allows_dynamic_callback_disabling(self, mock_request_headers):
|
||||
def test_default_admin_setting_allows_dynamic_callback_disabling(
|
||||
self, mock_request_headers
|
||||
):
|
||||
"""
|
||||
Test that when allow_dynamic_callback_disabling is not set,
|
||||
it defaults to True and allows dynamic callback disabling for premium users
|
||||
|
|
@ -253,8 +405,10 @@ class TestEnterpriseCallbackControls:
|
|||
mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"}
|
||||
litellm_params = {"proxy_server_request": {"url": "test"}}
|
||||
standard_callback_dynamic_params = StandardCallbackDynamicParams()
|
||||
|
||||
|
||||
# litellm.allow_dynamic_callback_disabling should default to True
|
||||
with patch('litellm.proxy.proxy_server.premium_user', True):
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically("langfuse", litellm_params, standard_callback_dynamic_params)
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
result = EnterpriseCallbackControls.is_callback_disabled_dynamically(
|
||||
"langfuse", litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
assert result is True
|
||||
|
|
|
|||
|
|
@ -327,7 +327,7 @@ class TestFileSearchGuardInResponsesMain:
|
|||
class TestManagedFilesVectorStoreAccess:
|
||||
def _make_hook(self):
|
||||
"""Return a ManagedFiles instance with prisma_client mocked."""
|
||||
from enterprise.litellm_enterprise.proxy.hooks.managed_files import (
|
||||
from litellm_enterprise.proxy.hooks.managed_files import (
|
||||
_PROXY_LiteLLMManagedFiles as ManagedFiles,
|
||||
)
|
||||
|
||||
|
|
@ -471,7 +471,7 @@ class TestManagedFilesVectorStoreAccess:
|
|||
@pytest.mark.asyncio
|
||||
async def test_F6_non_responses_call_type_skipped(self):
|
||||
"""Access check only runs for aresponses/responses call types."""
|
||||
from enterprise.litellm_enterprise.proxy.hooks.managed_files import (
|
||||
from litellm_enterprise.proxy.hooks.managed_files import (
|
||||
_PROXY_LiteLLMManagedFiles as ManagedFiles,
|
||||
)
|
||||
from litellm.proxy._types import CallTypes
|
||||
|
|
@ -499,7 +499,7 @@ class TestManagedFilesVectorStoreAccess:
|
|||
|
||||
class TestGetVectorStoreIdsFromFileSearchTools:
|
||||
def _make_hook(self):
|
||||
from enterprise.litellm_enterprise.proxy.hooks.managed_files import (
|
||||
from litellm_enterprise.proxy.hooks.managed_files import (
|
||||
_PROXY_LiteLLMManagedFiles as ManagedFiles,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -660,7 +660,7 @@ async def test_azure_content_safety_post_call_checks_all_choices(user_api_key):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_secret_detection_redacts_multimodal_text_parts(user_api_key):
|
||||
from enterprise.litellm_enterprise.enterprise_callbacks.secret_detection import (
|
||||
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
||||
_ENTERPRISE_SecretDetection,
|
||||
)
|
||||
|
||||
|
|
@ -696,7 +696,7 @@ async def test_secret_detection_redacts_multimodal_text_parts(user_api_key):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_secret_detection_redacts_responses_api_input(user_api_key):
|
||||
from enterprise.litellm_enterprise.enterprise_callbacks.secret_detection import (
|
||||
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
||||
_ENTERPRISE_SecretDetection,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ async def test_project_perm_check_uses_current_team_not_caller_supplied():
|
|||
"""The permission check must look at the project's existing team. Even
|
||||
if the caller is admin of an unrelated team, they must not pass when no
|
||||
explicit team_object is forced through."""
|
||||
from enterprise.litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_check_user_permission_for_project,
|
||||
)
|
||||
|
||||
|
|
@ -56,7 +56,7 @@ async def test_project_perm_check_uses_current_team_not_caller_supplied():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_perm_check_allows_team_admin_of_existing_team():
|
||||
from enterprise.litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_check_user_permission_for_project,
|
||||
)
|
||||
|
||||
|
|
@ -76,7 +76,7 @@ async def test_project_perm_check_allows_team_admin_of_existing_team():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_perm_check_proxy_admin_always_allowed():
|
||||
from enterprise.litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_check_user_permission_for_project,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1797,15 +1797,13 @@ class TestCustomUISSO:
|
|||
):
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"enterprise.litellm_enterprise.proxy.auth.custom_sso_handler": None
|
||||
},
|
||||
{"litellm_enterprise.proxy.auth.custom_sso_handler": None},
|
||||
):
|
||||
# Temporarily mock the google_login function call to test the import error path
|
||||
async def mock_google_login():
|
||||
# This mimics the relevant part of google_login that would trigger the import error
|
||||
try:
|
||||
from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import ( # noqa: F401
|
||||
from litellm_enterprise.proxy.auth.custom_sso_handler import ( # noqa: F401
|
||||
EnterpriseCustomSSOHandler,
|
||||
)
|
||||
|
||||
|
|
@ -1828,7 +1826,7 @@ class TestCustomUISSO:
|
|||
"""Test successful custom UI SSO sign-in with valid headers"""
|
||||
from fastapi_sso.sso.base import OpenID
|
||||
|
||||
from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import (
|
||||
from litellm_enterprise.proxy.auth.custom_sso_handler import (
|
||||
EnterpriseCustomSSOHandler,
|
||||
)
|
||||
from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler
|
||||
|
|
@ -1903,7 +1901,7 @@ class TestCustomUISSO:
|
|||
@pytest.mark.asyncio
|
||||
async def test_handle_custom_ui_sso_sign_in_rejects_untrusted_proxy(self):
|
||||
"""Custom UI SSO rejects spoofed identity headers from direct clients."""
|
||||
from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import (
|
||||
from litellm_enterprise.proxy.auth.custom_sso_handler import (
|
||||
EnterpriseCustomSSOHandler,
|
||||
)
|
||||
from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler
|
||||
|
|
@ -1943,7 +1941,7 @@ class TestCustomUISSO:
|
|||
"""
|
||||
from fastapi_sso.sso.base import OpenID
|
||||
|
||||
from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import (
|
||||
from litellm_enterprise.proxy.auth.custom_sso_handler import (
|
||||
EnterpriseCustomSSOHandler,
|
||||
)
|
||||
from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue