From b63d64dcf531b4ff623c481ed23e52fc9939d060 Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Tue, 1 Jul 2025 15:15:40 -0700 Subject: [PATCH] add subject and signature tests --- .../types/enterprise_callbacks/send_emails.py | 23 ++---- .../test_email_customization.py | 77 +++++++++++++++++++ 2 files changed, 83 insertions(+), 17 deletions(-) create mode 100644 tests/proxy_admin_ui_tests/test_email_customization.py diff --git a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py index 3adeefa62b..2d3c8adf2c 100644 --- a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py +++ b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py @@ -1,10 +1,10 @@ -from enum import Enum -from typing import Dict, List, Optional -from pydantic import BaseModel +import enum +from typing import Dict, List + +from pydantic import BaseModel, Field from litellm.proxy._types import WebhookEvent - class EmailParams(BaseModel): logo_url: str support_contact: str @@ -18,44 +18,33 @@ class SendKeyCreatedEmailEvent(WebhookEvent): virtual_key: str """ The virtual key that was created - this will be sk-123xxx, since we will be emailing this to the user to start using the key """ -class EmailEvent(str, Enum): +class EmailEvent(str, enum.Enum): virtual_key_created = "Virtual Key Created" new_user_invitation = "New User Invitation" - class EmailEventSettings(BaseModel): event: EmailEvent enabled: bool - - class EmailEventSettingsUpdateRequest(BaseModel): settings: List[EmailEventSettings] - - class EmailEventSettingsResponse(BaseModel): settings: List[EmailEventSettings] - - class DefaultEmailSettings(BaseModel): """Default settings for email events""" - settings: Dict[EmailEvent, bool] = Field( default_factory=lambda: { EmailEvent.virtual_key_created: False, # Off by default EmailEvent.new_user_invitation: True, # On by default } ) - def to_dict(self) -> Dict[str, bool]: """Convert to dictionary with string keys for storage""" return {event.value: enabled for event, enabled in self.settings.items()} - @classmethod def get_defaults(cls) -> Dict[str, bool]: """Get the default settings as a dictionary with string keys""" - return cls().to_dict() + return cls().to_dict() \ No newline at end of file diff --git a/tests/proxy_admin_ui_tests/test_email_customization.py b/tests/proxy_admin_ui_tests/test_email_customization.py new file mode 100644 index 0000000000..5d3798e926 --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_email_customization.py @@ -0,0 +1,77 @@ +import pytest +from unittest.mock import patch, MagicMock +import os +import sys +sys.path.insert(0, os.path.abspath("../..")) + +from enterprise.litellm_enterprise.enterprise_callbacks.send_emails.base_email import BaseEmailLogger +from enterprise.litellm_enterprise.types.enterprise_callbacks.send_emails import EmailEvent +from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER + +@pytest.fixture +def mock_env_vars(monkeypatch): + """Set up test environment variables""" + monkeypatch.setenv("EMAIL_LOGO_URL", "https://test-company.com/logo.png") + monkeypatch.setenv("EMAIL_SUPPORT_CONTACT", "support@test-company.com") + monkeypatch.setenv("EMAIL_SIGNATURE", "Best regards,\nTest Company Team") + monkeypatch.setenv("EMAIL_SUBJECT_INVITATION", "Welcome to Test Company!") + monkeypatch.setenv("EMAIL_SUBJECT_KEY_CREATED", "Your Test Company API Key") + monkeypatch.setenv("PROXY_BASE_URL", "http://test.com") + +@pytest.mark.asyncio +async def test_get_email_params_custom_templates(mock_env_vars): + """Test that _get_email_params returns correct values with custom templates""" + email_logger = BaseEmailLogger() + + # Test invitation email params + invitation_params = await email_logger._get_email_params( + email_event=EmailEvent.new_user_invitation, + user_email="test@example.com", + 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_email="test@example.com", + 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_default_templates(monkeypatch): + """Test that _get_email_params uses default templates when custom ones aren't provided""" + # Clear any existing environment variables + monkeypatch.delenv("EMAIL_SUBJECT_INVITATION", raising=False) + monkeypatch.delenv("EMAIL_SUBJECT_KEY_CREATED", raising=False) + monkeypatch.delenv("EMAIL_SIGNATURE", raising=False) + + 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" + ) + + 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" + ) + + assert key_params.subject == "LiteLLM: API Key Created" + assert key_params.signature == EMAIL_FOOTER \ No newline at end of file