scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

refactor(redis): move email budget alert claims into a typed store

Browse source
This commit is contained in:
Devin AI 2026-10-30 22:10:14 +00:00
parent 3a70bbc5ad
commit 30ffa81066
4 changed files with 655 additions and 457 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

275
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
View file

@ -6,7 +6,8 @@ Base class for sending emails to user after creating keys or invite links
import html
import json
import os
from typing import Final, List, Literal, Optional
from datetime import timedelta
from typing import Final, Literal, Protocol
from litellm_enterprise.types.enterprise_callbacks.send_emails import (
EmailEvent,
@ -18,6 +19,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
EMAIL_BUDGET_ALERT_TTL,
@ -46,6 +48,17 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
WebhookEvent,
)
from litellm.redis.stores.email_alert_claims import (
EmailAlertClaimsStore,
EmailBudgetAlert,
MaxBudgetAlert,
SoftBudgetAlert,
TeamMemberAlert,
ThresholdBudgetAlert,
TransportEmailAlertClaimsStore,
alert_key,
)
from litellm.redis.transport.contract import RedisFailure
from litellm.repositories.table_repositories import InvitationLinkRepository
from litellm.repositories.user_repository import UserRepository
from litellm.secret_managers.main import get_secret_bool
@ -54,13 +67,27 @@ from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
_BUDGET_ALERT_CLAIMS_TARGET: Final = "budget_alert_claims"
def _max_budget_alert_id(user_info: CallInfo) -> str:
class _LocalAlertCache(Protocol):
async def async_increment(self, key: str, value: float, *, ttl: int) -> float: ...
def delete_cache(self, key: str) -> None: ...
class _BudgetAlertCache(Protocol):
@property
def in_memory_cache(self) -> _LocalAlertCache | None: ...
@property
def redis_cache(self) -> RedisCache | None: ...
def _max_budget_alert_id(user_info: CallInfo) -> str | TeamMemberAlert:
if user_info.event_group == Litellm_EntityType.TEAM_MEMBER:
return f"team_member:{user_info.user_id}:{user_info.team_id}"
return TeamMemberAlert(user_info.user_id, user_info.team_id)
return user_info.token or user_info.user_id or "default_id"
def _parse_email_list(raw) -> List[str]:
def _parse_email_list(raw) -> list[str]:
"""Parse emails from a list or comma-separated string."""
if isinstance(raw, list):
return [e.strip() for e in raw if isinstance(e, str) and e.strip()]
@ -80,7 +107,7 @@ class BaseEmailLogger(CustomLogger):
def __init__(
self,
internal_usage_cache: Optional[DualCache] = None,
internal_usage_cache: DualCache | None = None,
**kwargs,
):
"""
@ -104,9 +131,7 @@ class BaseEmailLogger(CustomLogger):
event_message=event.event_message,
)
verbose_proxy_logger.debug(
f"send_user_invitation_email_event: {json.dumps(event, indent=4, default=str)}"
)
verbose_proxy_logger.debug(f"send_user_invitation_email_event: {json.dumps(event, indent=4, default=str)}")
email_html_content = USER_INVITATION_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
@ -124,11 +149,7 @@ class BaseEmailLogger(CustomLogger):
html_body=email_html_content,
)
pass
async def send_key_created_email(
self, send_key_created_email_event: SendKeyCreatedEmailEvent
):
async def send_key_created_email(self, send_key_created_email_event: SendKeyCreatedEmailEvent):
"""
Send email to user after creating key for the user
"""
@ -144,9 +165,7 @@ class BaseEmailLogger(CustomLogger):
)
# Check if API key should be included in email
include_api_key = get_secret_bool(
secret_name="EMAIL_INCLUDE_API_KEY", default_value=True
)
include_api_key = get_secret_bool(secret_name="EMAIL_INCLUDE_API_KEY", default_value=True)
if include_api_key is None:
include_api_key = True # Default to True if not set
key_token_display = (
@ -171,11 +190,8 @@ class BaseEmailLogger(CustomLogger):
subject=email_params.subject,
html_body=email_html_content,
)
pass
async def send_key_rotated_email(
self, send_key_rotated_email_event: SendKeyRotatedEmailEvent
):
async def send_key_rotated_email(self, send_key_rotated_email_event: SendKeyRotatedEmailEvent):
"""
Send email to user after rotating key for the user
"""
@ -191,9 +207,7 @@ class BaseEmailLogger(CustomLogger):
)
# Check if API key should be included in email
include_api_key = get_secret_bool(
secret_name="EMAIL_INCLUDE_API_KEY", default_value=True
)
include_api_key = get_secret_bool(secret_name="EMAIL_INCLUDE_API_KEY", default_value=True)
if include_api_key is None:
include_api_key = True # Default to True if not set
key_token_display = (
@ -218,7 +232,6 @@ class BaseEmailLogger(CustomLogger):
subject=email_params.subject,
html_body=email_html_content,
)
pass
async def send_soft_budget_alert_email(self, event: WebhookEvent):
"""
@ -236,9 +249,7 @@ class BaseEmailLogger(CustomLogger):
)
# Format budget values
soft_budget_str = (
f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
)
soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
max_budget_info = ""
if event.max_budget is not None:
@ -260,7 +271,6 @@ class BaseEmailLogger(CustomLogger):
subject=email_params.subject,
html_body=email_html_content,
)
pass
async def send_team_soft_budget_alert_email(self, event: WebhookEvent):
"""
@ -268,7 +278,7 @@ class BaseEmailLogger(CustomLogger):
Supports multiple recipients via alert_emails field from team metadata
"""
# Collect all recipient emails
recipient_emails: List[str] = []
recipient_emails: list[str] = []
# Add additional alert emails from team metadata.soft_budget_alert_emails
if hasattr(event, "alert_emails") and event.alert_emails:
@ -306,9 +316,7 @@ class BaseEmailLogger(CustomLogger):
)
# Format budget values
soft_budget_str = (
f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
)
soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
max_budget_info = ""
if event.max_budget is not None:
@ -335,13 +343,12 @@ class BaseEmailLogger(CustomLogger):
subject=email_params.subject,
html_body=email_html_content,
)
pass
async def send_max_budget_alert_email(
self,
event: WebhookEvent,
threshold_pct: Optional[int] = None,
recipient_emails: Optional[List[str]] = None,
threshold_pct: int | None = None,
recipient_emails: list[str] | None = None,
):
"""
Send email to user when max budget alert threshold is reached.
@ -359,20 +366,14 @@ class BaseEmailLogger(CustomLogger):
# Format budget values
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
max_budget_str = (
f"${event.max_budget}" if event.max_budget is not None else "N/A"
)
max_budget_str = f"${event.max_budget}" if event.max_budget is not None else "N/A"
# Calculate percentage and alert threshold
percentage = threshold_pct if threshold_pct is not None else int(
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
percentage = (
threshold_pct if threshold_pct is not None else int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
)
threshold_fraction = percentage / 100.0
alert_threshold_str = (
f"${event.max_budget * threshold_fraction:.2f}"
if event.max_budget is not None
else "N/A"
)
alert_threshold_str = f"${event.max_budget * threshold_fraction:.2f}" if event.max_budget is not None else "N/A"
if recipient_emails:
# Multi-threshold path: batch send with generic key-based greeting
@ -382,9 +383,7 @@ class BaseEmailLogger(CustomLogger):
user_email=event.user_email or recipient_emails[0],
event_message=event.event_message,
)
greeting = html.escape(
event.user_email or event.key_alias or event.token or ""
)
greeting = html.escape(event.user_email or event.key_alias or event.token or "")
if event.event_group == Litellm_EntityType.TEAM_MEMBER:
email_html_content = TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
@ -488,10 +487,7 @@ class BaseEmailLogger(CustomLogger):
# For non-team alerts, require either max_budget or soft_budget
if user_info.max_budget is None and user_info.soft_budget is None:
return
if (
user_info.soft_budget is not None
and user_info.spend >= user_info.soft_budget
):
if user_info.soft_budget is not None and user_info.spend >= user_info.soft_budget:
# Generate cache key based on event type and identifier
# Use appropriate ID based on event_group to ensure unique cache keys per entity type
if user_info.event_group == Litellm_EntityType.TEAM:
@ -503,13 +499,8 @@ class BaseEmailLogger(CustomLogger):
else:
# For KEY and other types, use token or user_id
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}"
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
soft_alert: Final = SoftBudgetAlert(_id)
send_count = await self._claim_budget_alert(_cache, soft_alert)
if send_count is None or send_count <= 1:
# Create WebhookEvent for soft budget alert
event_message = f"Soft Budget Crossed - Total Soft Budget: ${user_info.soft_budget}"
@ -544,7 +535,7 @@ class BaseEmailLogger(CustomLogger):
f"Error sending soft budget alert email: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
await self._release_budget_alert_claim(_cache, soft_alert)
return
# For max_budget_alert, check if we've already sent an alert
@ -552,34 +543,20 @@ class BaseEmailLogger(CustomLogger):
if user_info.max_budget is not None and user_info.spend is not None:
if user_info.max_budget_alert_emails:
# New path: multi-threshold alerts
await self._handle_multi_threshold_max_budget_alert(
user_info=user_info, _cache=_cache
)
await self._handle_multi_threshold_max_budget_alert(user_info=user_info, _cache=_cache)
return
alert_threshold = (
user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
alert_threshold = user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
if (
user_info.spend >= alert_threshold
and user_info.spend < user_info.max_budget
):
if user_info.spend >= alert_threshold and user_info.spend < user_info.max_budget:
# Generate cache key based on event type and identifier
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:max_budget_alert:{_id}"
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
max_alert: Final = MaxBudgetAlert(_id)
send_count = await self._claim_budget_alert(_cache, max_alert)
if send_count is None or send_count <= 1:
# Calculate percentage
percentage = int(
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
)
percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
# Create WebhookEvent for max budget alert
event_message = f"Max Budget Alert - {percentage}% of Maximum Budget Reached"
@ -609,7 +586,7 @@ class BaseEmailLogger(CustomLogger):
f"Error sending max budget alert email: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
await self._release_budget_alert_claim(_cache, max_alert)
return
@with_service_target(_BUDGET_ALERT_CLAIMS_TARGET)
@ -636,9 +613,7 @@ class BaseEmailLogger(CustomLogger):
continue
_id = _max_budget_alert_id(user_info)
_cache_key = (
f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}"
)
threshold_alert = ThresholdBudgetAlert(_id, threshold_pct)
# Parse emails + auto-include owner
emails = _parse_email_list(raw_emails)
@ -653,11 +628,7 @@ class BaseEmailLogger(CustomLogger):
continue
recipient_emails = list(set(emails))
send_count = await _cache.async_increment_cache(
key=_cache_key,
value=1,
ttl=EMAIL_BUDGET_ALERT_TTL,
)
send_count = await self._claim_budget_alert(_cache, threshold_alert)
if send_count is not None and send_count > 1:
continue
@ -696,24 +667,44 @@ class BaseEmailLogger(CustomLogger):
f"Error sending multi-threshold max budget alert email for {threshold_pct}%: {e}",
exc_info=True,
)
await self._release_budget_alert_claim(_cache, _cache_key)
await self._release_budget_alert_claim(_cache, threshold_alert)
@with_service_target(_BUDGET_ALERT_CLAIMS_TARGET)
async def _release_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None:
async def _claim_budget_alert(self, cache: _BudgetAlertCache, alert: EmailBudgetAlert) -> float | None:
try:
await cache.async_delete_cache(key=cache_key)
except Exception:
verbose_proxy_logger.debug(
"Failed to release budget alert claim for %s; it expires with the TTL",
cache_key,
local_count: Final = (
await cache.in_memory_cache.async_increment(alert_key(alert), 1, ttl=EMAIL_BUDGET_ALERT_TTL)
if cache.in_memory_cache is not None
else None
)
except Exception:
return None
if cache.redis_cache is None:
return local_count
try:
store: Final[EmailAlertClaimsStore] = TransportEmailAlertClaimsStore(cache.redis_cache.transport())
count: Final = await store.claim(alert, ttl=timedelta(seconds=EMAIL_BUDGET_ALERT_TTL))
return local_count if isinstance(count, RedisFailure) else count
except Exception:
return local_count
@with_service_target(_BUDGET_ALERT_CLAIMS_TARGET)
async def _release_budget_alert_claim(self, cache: _BudgetAlertCache, alert: EmailBudgetAlert) -> None:
try:
if cache.in_memory_cache is not None:
cache.in_memory_cache.delete_cache(alert_key(alert))
if cache.redis_cache is not None:
store: Final[EmailAlertClaimsStore] = TransportEmailAlertClaimsStore(cache.redis_cache.transport())
await store.release_failed_send(alert)
except Exception:
return
async def _get_email_params(
self,
email_event: EmailEvent,
user_id: Optional[str] = None,
user_email: Optional[str] = None,
event_message: Optional[str] = None,
user_id: str | None = None,
user_email: str | None = None,
event_message: str | None = None,
) -> EmailParams:
"""
Get common email parameters used across different email sending methods
@ -738,12 +729,8 @@ class BaseEmailLogger(CustomLogger):
unused_custom_fields = []
# Function to safely get custom value or default
def get_custom_or_default(
custom_value: Optional[str], default_value: str, field_name: str
) -> str:
if (
custom_value is not None
): # Only check premium if trying to use custom value
def get_custom_or_default(custom_value: str | None, default_value: str, field_name: str) -> str:
if custom_value is not None: # Only check premium if trying to use custom value
from litellm.proxy.proxy_server import premium_user
if premium_user is not True:
@ -754,15 +741,9 @@ class BaseEmailLogger(CustomLogger):
# Get parameters, falling back to defaults if custom values aren't allowed
logo_url = get_custom_or_default(custom_logo, LITELLM_LOGO_URL, "logo URL")
support_contact = get_custom_or_default(
custom_support, self.DEFAULT_SUPPORT_EMAIL, "support contact"
)
base_url = os.getenv(
"PROXY_BASE_URL", "http://0.0.0.0:4000"
) # Not a premium feature
signature = get_custom_or_default(
custom_signature, EMAIL_FOOTER, "email signature"
)
support_contact = get_custom_or_default(custom_support, self.DEFAULT_SUPPORT_EMAIL, "support contact")
base_url = os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000") # Not a premium feature
signature = get_custom_or_default(custom_signature, EMAIL_FOOTER, "email signature")
# Get custom subject template based on email event type
if email_event == EmailEvent.new_user_invitation:
@ -787,25 +768,15 @@ class BaseEmailLogger(CustomLogger):
else:
subject_template = "LiteLLM: {event_message}"
subject = (
subject_template.format(event_message=event_message)
if event_message
else "LiteLLM Notification"
)
subject = subject_template.format(event_message=event_message) if event_message else "LiteLLM Notification"
recipient_email: Optional[str] = (
user_email or await self._lookup_user_email_from_db(user_id=user_id)
)
recipient_email: str | None = user_email or await self._lookup_user_email_from_db(user_id=user_id)
if recipient_email is None:
raise ValueError(
f"User email not found for user_id: {user_id}. User email is required to send email."
)
raise ValueError(f"User email not found for user_id: {user_id}. User email is required to send email.")
# if user invited event then send invitation link
if email_event == EmailEvent.new_user_invitation:
base_url = await self._get_invitation_link(
user_id=user_id, base_url=base_url
)
base_url = await self._get_invitation_link(user_id=user_id, base_url=base_url)
# If any custom fields were not applied, log a warning
if unused_custom_fields:
@ -826,7 +797,7 @@ class BaseEmailLogger(CustomLogger):
signature=signature,
)
def _format_key_budget(self, max_budget: Optional[float]) -> str:
def _format_key_budget(self, max_budget: float | None) -> str:
"""
Format the key budget to be displayed in the email
"""
@ -834,27 +805,23 @@ class BaseEmailLogger(CustomLogger):
return "No budget"
return f"${max_budget}"
async def _lookup_user_email_from_db(self, user_id: Optional[str]) -> Optional[str]:
async def _lookup_user_email_from_db(self, user_id: str | None) -> str | None:
"""
Lookup user email from user_id
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
verbose_proxy_logger.debug(
f"Prisma client not found. Unable to lookup user email for user_id: {user_id}"
)
verbose_proxy_logger.debug(f"Prisma client not found. Unable to lookup user email for user_id: {user_id}")
return None
user_row = await UserRepository(prisma_client).table.find_unique(
where={"user_id": user_id}
)
user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
if user_row is not None:
return user_row.user_email
return None
async def _get_invitation_link(self, user_id: Optional[str], base_url: str) -> str:
async def _get_invitation_link(self, user_id: str | None, base_url: str) -> str:
"""
Get invitation link for the user
"""
@ -883,9 +850,7 @@ class BaseEmailLogger(CustomLogger):
return base_url
invitation_link = self._construct_invitation_link(invitation.id, base_url)
verbose_proxy_logger.info(
f"Successfully created invitation link for user_id: {user_id}"
)
verbose_proxy_logger.info(f"Successfully created invitation link for user_id: {user_id}")
return invitation_link
async def _is_prisma_client_available(self) -> bool:
@ -893,9 +858,7 @@ class BaseEmailLogger(CustomLogger):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
verbose_proxy_logger.debug(
"Prisma client not found. Unable to lookup invitation"
)
verbose_proxy_logger.debug("Prisma client not found. Unable to lookup invitation")
return False
return True
@ -923,39 +886,29 @@ class BaseEmailLogger(CustomLogger):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
verbose_proxy_logger.error(
"Prisma client is None in _get_or_create_invitation"
)
verbose_proxy_logger.error("Prisma client is None in _get_or_create_invitation")
return None
try:
# Try to get existing invitation
existing_invitations = (
await InvitationLinkRepository(prisma_client).table.find_many(
where={"user_id": user_id},
order={"created_at": "desc"},
)
existing_invitations = await InvitationLinkRepository(prisma_client).table.find_many(
where={"user_id": user_id},
order={"created_at": "desc"},
)
if existing_invitations and len(existing_invitations) > 0:
verbose_proxy_logger.debug(
f"Found existing invitation for user_id: {user_id}"
)
verbose_proxy_logger.debug(f"Found existing invitation for user_id: {user_id}")
return existing_invitations[0]
# Create new invitation if none exists
verbose_proxy_logger.debug(
f"Creating new invitation for user_id: {user_id}"
)
verbose_proxy_logger.debug(f"Creating new invitation for user_id: {user_id}")
return await create_invitation_for_user(
data=InvitationNew(user_id=user_id),
user_api_key_dict=UserAPIKeyAuth(user_id=user_id),
)
except Exception as e:
verbose_proxy_logger.error(
f"Error getting/creating invitation for user_id {user_id}: {e}"
)
verbose_proxy_logger.error(f"Error getting/creating invitation for user_id {user_id}: {e}")
return None
def _construct_invitation_link(self, invitation_id: str, base_url: str) -> str:
@ -971,7 +924,7 @@ class BaseEmailLogger(CustomLogger):
async def send_email(
self,
from_email: str,
to_email: List[str],
to_email: list[str],
subject: str,
html_body: str,
):

77
litellm/redis/stores/email_alert_claims.py Normal file
View file

@ -0,0 +1,77 @@
from dataclasses import dataclass
from datetime import timedelta
from typing import Final, Protocol, TypeAlias
from litellm.redis.transport.contract import RedisFailure, RedisTransport
@dataclass(frozen=True, slots=True)
class SoftBudgetAlert:
identifier: str
@dataclass(frozen=True, slots=True)
class MaxBudgetAlert:
identifier: str
@dataclass(frozen=True, slots=True)
class TeamMemberAlert:
user_id: str | None
team_id: str | None
@dataclass(frozen=True, slots=True)
class ThresholdBudgetAlert:
identifier: str | TeamMemberAlert
threshold_pct: int
EmailBudgetAlert: TypeAlias = SoftBudgetAlert | MaxBudgetAlert | ThresholdBudgetAlert
def alert_key(alert: EmailBudgetAlert) -> str:
match alert:
case SoftBudgetAlert(identifier):
return f"email_budget_alerts:soft_budget_crossed:{identifier}"
case MaxBudgetAlert(identifier):
return f"email_budget_alerts:max_budget_alert:{identifier}"
case ThresholdBudgetAlert(identifier, threshold_pct):
subject: Final = (
f"team_member:{identifier.user_id}:{identifier.team_id}"
if isinstance(identifier, TeamMemberAlert)
else identifier
)
return f"email_budget_alerts:max_budget_alert:{threshold_pct}:{subject}"
class EmailAlertClaimsStore(Protocol):
async def claim(self, alert: EmailBudgetAlert, *, ttl: timedelta) -> float | RedisFailure: ...
async def release_failed_send(self, alert: EmailBudgetAlert) -> bool | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportEmailAlertClaimsStore:
transport: RedisTransport
async def claim(self, alert: EmailBudgetAlert, *, ttl: timedelta) -> float | RedisFailure:
key: Final = alert_key(alert)
batch: Final = self.transport.batch(transaction=False)
increment: Final = batch.incrbyfloat(key, 1)
expiry: Final = batch.ttl(key)
count: Final = await increment
current_ttl: Final = await expiry
if isinstance(count, RedisFailure):
return count
if isinstance(current_ttl, RedisFailure):
return current_ttl
if current_ttl == -1:
expired: Final = await self.transport.expire_seconds(key, ttl)
if isinstance(expired, RedisFailure):
return expired
return count
async def release_failed_send(self, alert: EmailBudgetAlert) -> bool | RedisFailure:
deleted: Final = await self.transport.delete(alert_key(alert))
return deleted if isinstance(deleted, RedisFailure) else deleted > 0

467
tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py
View file

@ -1,33 +1,52 @@
import asyncio
import json
import os
import unittest.mock as mock
from collections.abc import AsyncIterator
from types import SimpleNamespace
from typing import Final
from unittest import mock
from unittest.mock import patch
import fakeredis
import pytest
from fastapi.testclient import TestClient
from litellm.caching.caching import DualCache
import pytest_asyncio
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
BaseEmailLogger,
)
from litellm_enterprise.types.enterprise_callbacks.send_emails import (
EmailEvent,
SendKeyCreatedEmailEvent,
SendKeyRotatedEmailEvent,
)
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import EMAIL_BUDGET_ALERT_TTL
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
from litellm.proxy._types import (
CallInfo,
LiteLLM_UserTable,
Litellm_EntityType,
LiteLLM_UserTable,
WebhookEvent,
)
from litellm.constants import EMAIL_BUDGET_ALERT_TTL
from tests._support.fake_redis import fake_redis_cache, local_get
@pytest_asyncio.fixture(loop_scope="function")
async def alert_redis() -> AsyncIterator[tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer]]:
server: Final = fakeredis.FakeServer()
cache: Final = fake_redis_cache(server, namespace="alerts")
raw: Final = fakeredis.FakeAsyncRedis(server=server)
try:
yield cache, raw, server
finally:
await raw.aclose() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # fakeredis stubs predate aclose
@pytest.fixture
def redis_email_logger(
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
) -> BaseEmailLogger:
return BaseEmailLogger(internal_usage_cache=DualCache(redis_cache=alert_redis[0]))
@pytest.fixture(autouse=True)
@ -51,9 +70,7 @@ def mock_send_email():
@pytest.fixture
def mock_lookup_user_email():
with mock.patch.object(
BaseEmailLogger, "_lookup_user_email_from_db"
) as mock_lookup:
with mock.patch.object(BaseEmailLogger, "_lookup_user_email_from_db") as mock_lookup:
yield mock_lookup
@ -66,9 +83,7 @@ def test_format_key_budget(base_email_logger):
@pytest.mark.asyncio
async def test_send_key_created_email(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_key_created_email(base_email_logger, mock_send_email, mock_lookup_user_email):
# Setup test data
event = SendKeyCreatedEmailEvent(
user_id="test_user",
@ -104,9 +119,7 @@ async def test_send_key_created_email(
@pytest.mark.asyncio
async def test_send_user_invitation_email(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_user_invitation_email(base_email_logger, mock_send_email, mock_lookup_user_email):
# Setup test data
event = WebhookEvent(
user_id="test_user",
@ -139,9 +152,7 @@ async def test_send_user_invitation_email(
@pytest.mark.asyncio
async def test_send_user_invitation_email_from_db(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_user_invitation_email_from_db(base_email_logger, mock_send_email, mock_lookup_user_email):
# Setup test data with no direct email but one in the database
event = WebhookEvent(
user_id="test_user",
@ -177,9 +188,7 @@ async def test_send_user_invitation_email_from_db(
@pytest.mark.asyncio
async def test_send_user_invitation_email_no_email(
base_email_logger, mock_lookup_user_email
):
async def test_send_user_invitation_email_no_email(base_email_logger, mock_lookup_user_email):
# Setup test data with no email
event = WebhookEvent(
user_id="test_user",
@ -198,9 +207,7 @@ async def test_send_user_invitation_email_no_email(
@pytest.mark.asyncio
async def test_send_key_created_email_no_email(
base_email_logger, mock_lookup_user_email
):
async def test_send_key_created_email_no_email(base_email_logger, mock_lookup_user_email):
# Setup test data with no email
event = SendKeyCreatedEmailEvent(
user_id="test_user",
@ -222,9 +229,7 @@ async def test_send_key_created_email_no_email(
@pytest.mark.asyncio
async def test_send_key_rotated_email(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_key_rotated_email(base_email_logger, mock_send_email, mock_lookup_user_email):
"""
Test that send_key_rotated_email sends an email with the correct parameters and content
"""
@ -262,9 +267,7 @@ async def test_send_key_rotated_email(
@pytest.mark.asyncio
async def test_send_key_created_email_without_key(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_key_created_email_without_key(base_email_logger, mock_send_email, mock_lookup_user_email):
"""
Test that send_key_created_email hides the API key when EMAIL_INCLUDE_API_KEY is false
"""
@ -291,16 +294,11 @@ 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
async def test_send_key_rotated_email_without_key(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_key_rotated_email_without_key(base_email_logger, mock_send_email, mock_lookup_user_email):
"""
Test that send_key_rotated_email hides the API key when EMAIL_INCLUDE_API_KEY is false
"""
@ -328,10 +326,7 @@ 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
@ -350,17 +345,11 @@ async def test_get_invitation_link(base_email_logger):
with mock.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
# Test with valid user_id
result = await base_email_logger._get_invitation_link(
user_id="test-user", base_url="http://test.com"
)
assert (
result == "http://test.com/ui/onboarding?invitation_id=test-invitation-id"
)
result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com")
assert result == "http://test.com/ui/onboarding?invitation_id=test-invitation-id"
# Test with None user_id
result = await base_email_logger._get_invitation_link(
user_id=None, base_url="http://test.com"
)
result = await base_email_logger._get_invitation_link(user_id=None, base_url="http://test.com")
assert result == "http://test.com"
# Test with no invitation links
@ -368,17 +357,13 @@ async def test_get_invitation_link(base_email_logger):
return []
mock_prisma.db.litellm_invitationlink.find_many = mock_find_many_empty
result = await base_email_logger._get_invitation_link(
user_id="test-user", base_url="http://test.com"
)
result = await base_email_logger._get_invitation_link(user_id="test-user", base_url="http://test.com")
assert result == "http://test.com"
def test_construct_invitation_link(base_email_logger):
# Test invitation link construction
result = base_email_logger._construct_invitation_link(
invitation_id="test-id-123", base_url="http://test.com"
)
result = base_email_logger._construct_invitation_link(invitation_id="test-id-123", base_url="http://test.com")
assert result == "http://test.com/ui/onboarding?invitation_id=test-id-123"
@ -404,9 +389,7 @@ async def test_get_invitation_link_creates_new_when_none_exist(base_email_logger
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"
)
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()
@ -415,10 +398,7 @@ async def test_get_invitation_link_creates_new_when_none_exist(base_email_logger
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/onboarding?invitation_id=new-invitation-id"
)
assert result == "http://test.com/ui/onboarding?invitation_id=new-invitation-id"
@pytest.mark.asyncio
@ -441,18 +421,13 @@ async def test_get_invitation_link_uses_existing_when_available(base_email_logge
"litellm.proxy.management_helpers.user_invitation.create_invitation_for_user"
) as mock_create_invitation:
# Execute
result = await base_email_logger._get_invitation_link(
user_id="test-user", base_url="http://test.com"
)
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/onboarding?invitation_id=existing-invitation-id"
)
assert result == "http://test.com/ui/onboarding?invitation_id=existing-invitation-id"
@pytest.mark.asyncio
@ -477,9 +452,7 @@ async def test_get_invitation_link_creates_new_when_list_is_none(base_email_logg
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"
)
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()
@ -488,16 +461,11 @@ async def test_get_invitation_link_creates_new_when_list_is_none(base_email_logg
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/onboarding?invitation_id=new-invitation-from-none"
)
assert result == "http://test.com/ui/onboarding?invitation_id=new-invitation-from-none"
@pytest.mark.asyncio
async def test_get_email_params_user_invitation(
base_email_logger, mock_lookup_user_email
):
async def test_get_email_params_user_invitation(base_email_logger, mock_lookup_user_email):
# Mock environment variables
with mock.patch.dict(
os.environ,
@ -520,14 +488,9 @@ 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/onboarding?invitation_id=test-id"
)
assert result.base_url == "http://test.com/ui/onboarding?invitation_id=test-id"
assert result.recipient_email == "test@example.com"
@ -593,10 +556,7 @@ async def test_get_email_params_non_premium_user(mock_env_vars):
# 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
@ -644,9 +604,7 @@ async def test_get_email_params_default_templates(monkeypatch):
@pytest.mark.asyncio
async def test_send_soft_budget_alert_email(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_soft_budget_alert_email(base_email_logger, mock_send_email, mock_lookup_user_email):
"""Test that send_soft_budget_alert_email sends an email with the correct parameters and content"""
event = WebhookEvent(
user_id="test_user",
@ -673,19 +631,14 @@ 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
@pytest.mark.asyncio
async def test_send_soft_budget_alert_email_no_max_budget(
base_email_logger, mock_send_email, mock_lookup_user_email
):
async def test_send_soft_budget_alert_email_no_max_budget(base_email_logger, mock_send_email, mock_lookup_user_email):
"""Test that send_soft_budget_alert_email handles missing max_budget correctly"""
event = WebhookEvent(
user_id="test_user",
@ -710,13 +663,15 @@ 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(
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that budget_alerts sends email when soft budget is crossed"""
user_info = CallInfo(
user_id="test_user",
@ -727,39 +682,29 @@ async def test_budget_alerts_soft_budget_crossed(base_email_logger, mock_send_em
event_group=Litellm_EntityType.USER,
)
# Mock the cache so the claim is won (increment returns 1)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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="soft_budget", user_info=user_info)
await redis_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify email was sent
mock_send_email.assert_called_once()
call_args = mock_send_email.call_args[1]
assert call_args["to_email"] == ["test@example.com"]
# Verify the send slot was claimed to prevent duplicate alerts
mock_cache.async_increment_cache.assert_called_once()
cache_call_args = mock_cache.async_increment_cache.call_args[1]
assert await alert_redis[1].get("alerts:email_budget_alerts:soft_budget_crossed:test_user") == b"1"
assert (
cache_call_args["key"]
== "email_budget_alerts:soft_budget_crossed:test_user"
0
< await alert_redis[1].pttl("alerts:email_budget_alerts:soft_budget_crossed:test_user")
<= EMAIL_BUDGET_ALERT_TTL * 1000
)
assert cache_call_args["value"] == 1
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@pytest.mark.asyncio
async def test_budget_alerts_soft_budget_not_crossed(
base_email_logger, mock_send_email
):
async def test_budget_alerts_soft_budget_not_crossed(base_email_logger, mock_send_email):
"""Test that budget_alerts does not send email when soft budget is not crossed"""
user_info = CallInfo(
user_id="test_user",
@ -770,9 +715,6 @@ async def test_budget_alerts_soft_budget_not_crossed(
event_group=Litellm_EntityType.USER,
)
mock_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify email was NOT sent
@ -781,7 +723,9 @@ async def test_budget_alerts_soft_budget_not_crossed(
@pytest.mark.asyncio
async def test_budget_alerts_soft_budget_duplicate_prevention(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that budget_alerts does not send duplicate alerts within TTL period"""
user_info = CallInfo(
@ -793,12 +737,10 @@ async def test_budget_alerts_soft_budget_duplicate_prevention(
event_group=Litellm_EntityType.USER,
)
# Mock the cache so the slot is already claimed (increment returns > 1)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=2)
base_email_logger.internal_usage_cache = mock_cache
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
await alert_redis[0].async_increment(
"email_budget_alerts:soft_budget_crossed:test_user", 1, ttl=EMAIL_BUDGET_ALERT_TTL
)
await redis_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify email was NOT sent (duplicate prevention)
mock_send_email.assert_not_called()
@ -824,7 +766,9 @@ async def test_budget_alerts_no_budgets(base_email_logger, mock_send_email):
@pytest.mark.asyncio
async def test_budget_alerts_uses_token_for_cache_key(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that budget_alerts uses token for cache key when available"""
user_info = CallInfo(
@ -837,32 +781,19 @@ async def test_budget_alerts_uses_token_for_cache_key(
event_group=Litellm_EntityType.KEY,
)
# Mock the cache so the claim is won (increment returns 1)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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="soft_budget", user_info=user_info)
await redis_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify cache key uses token instead of user_id
mock_cache.async_increment_cache.assert_called_once()
cache_call_args = mock_cache.async_increment_cache.call_args[1]
assert (
cache_call_args["key"]
== "email_budget_alerts:soft_budget_crossed:hashed_token_123"
)
assert await alert_redis[1].get("alerts:email_budget_alerts:soft_budget_crossed:hashed_token_123") == b"1"
@pytest.mark.asyncio
async def test_get_email_params_soft_budget_crossed(
base_email_logger, mock_lookup_user_email
):
async def test_get_email_params_soft_budget_crossed(base_email_logger, mock_lookup_user_email):
"""Test that _get_email_params handles soft_budget_crossed event correctly"""
with mock.patch.dict(
os.environ,
@ -877,16 +808,16 @@ 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"
@pytest.mark.asyncio
async def test_budget_alerts_max_budget_alert_crossed(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that budget_alerts sends email when max budget alert threshold is crossed"""
user_info = CallInfo(
@ -897,37 +828,32 @@ async def test_budget_alerts_max_budget_alert_crossed(
event_group=Litellm_EntityType.USER,
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_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_increment_cache.assert_called_once()
cache_call_args = mock_cache.async_increment_cache.call_args[1]
assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:test_user") == b"1"
assert (
cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user"
0
< await alert_redis[1].pttl("alerts:email_budget_alerts:max_budget_alert:test_user")
<= EMAIL_BUDGET_ALERT_TTL * 1000
)
assert cache_call_args["value"] == 1
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@pytest.mark.asyncio
async def test_multi_threshold_sends_crossed_thresholds(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that multi-threshold path sends emails for all crossed thresholds"""
user_info = CallInfo(
@ -944,29 +870,23 @@ async def test_multi_threshold_sends_crossed_thresholds(
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_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_increment_cache.call_args_list
assert sorted(await alert_redis[1].keys("*")) == [
b"alerts:email_budget_alerts:max_budget_alert:50:hashed_key_1",
b"alerts:email_budget_alerts:max_budget_alert:75:hashed_key_1",
]
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
@pytest.mark.asyncio
async def test_multi_threshold_dedup_cache_prevents_resend(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that cached thresholds are not re-sent"""
user_info = CallInfo(
@ -982,31 +902,22 @@ async def test_multi_threshold_dedup_cache_prevents_resend(
},
)
# Simulate 50% already claimed (increment returns >1), 75% first send (returns 1)
async def cache_increment(key, value, ttl=None):
if "50:" in key:
return 2
return 1
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(side_effect=cache_increment)
base_email_logger.internal_usage_cache = mock_cache
await alert_redis[0].async_increment(
"email_budget_alerts:max_budget_alert:50:hashed_key_1", 1, ttl=EMAIL_BUDGET_ALERT_TTL
)
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 redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
# Only 75% should fire
assert mock_send_email.call_count == 1
cache_key = mock_cache.async_increment_cache.call_args[1]["key"]
assert "75:" in cache_key
assert "75%" in mock_send_email.call_args.kwargs["subject"]
assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:50:hashed_key_1") == b"2"
assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:75:hashed_key_1") == b"1"
@pytest.mark.asyncio
async def test_multi_threshold_owner_email_auto_included(
base_email_logger, mock_send_email
):
async def test_multi_threshold_owner_email_auto_included(redis_email_logger: BaseEmailLogger, mock_send_email):
"""Test that the owner email is auto-appended and deduplicated"""
user_info = CallInfo(
token="hashed_key_1",
@ -1020,14 +931,8 @@ async def test_multi_threshold_owner_email_auto_included(
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_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"]
@ -1036,9 +941,7 @@ async def test_multi_threshold_owner_email_auto_included(
@pytest.mark.asyncio
async def test_multi_threshold_malformed_keys_skipped(
base_email_logger, mock_send_email
):
async def test_multi_threshold_malformed_keys_skipped(redis_email_logger: BaseEmailLogger, mock_send_email):
"""Test that non-numeric threshold keys are skipped"""
user_info = CallInfo(
token="hashed_key_1",
@ -1053,23 +956,15 @@ async def test_multi_threshold_malformed_keys_skipped(
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_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
@pytest.mark.asyncio
async def test_multi_threshold_empty_emails_only_owner(
base_email_logger, mock_send_email
):
async def test_multi_threshold_empty_emails_only_owner(redis_email_logger: BaseEmailLogger, mock_send_email):
"""Test that empty email list for a threshold sends only to owner"""
user_info = CallInfo(
token="hashed_key_1",
@ -1083,14 +978,8 @@ async def test_multi_threshold_empty_emails_only_owner(
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_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"]
@ -1099,7 +988,7 @@ async def test_multi_threshold_empty_emails_only_owner(
@pytest.mark.asyncio
async def test_multi_threshold_team_member_alert_renders_member_template_per_team(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger, mock_send_email, alert_redis
):
"""A team member budget alert is keyed per member and team, names the member and team,
and goes to the member plus the threshold's configured recipients"""
@ -1114,22 +1003,15 @@ async def test_multi_threshold_team_member_alert_renders_member_template_per_tea
max_budget_alert_emails={"50": [], "100": ["finance@co.com"]},
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
cache_keys = sorted(c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list)
assert cache_keys == [
"email_budget_alerts:max_budget_alert:100:team_member:member_1:team_a",
"email_budget_alerts:max_budget_alert:50:team_member:member_1:team_a",
assert sorted(await alert_redis[1].keys("*")) == [
b"alerts:email_budget_alerts:max_budget_alert:100:team_member:member_1:team_a",
b"alerts:email_budget_alerts:max_budget_alert:50:team_member:member_1:team_a",
]
assert mock_send_email.call_count == 2
hundred = next(
c.kwargs for c in mock_send_email.call_args_list if "100%" in c.kwargs["subject"]
)
hundred = next(c.kwargs for c in mock_send_email.call_args_list if "100%" in c.kwargs["subject"])
assert hundred["subject"] == "LiteLLM: Team Member Budget Alert - 100% of Team Member Budget Reached"
assert sorted(hundred["to_email"]) == ["finance@co.com", "member@co.com"]
assert "member@co.com" in hundred["html_body"] and "Platform" in hundred["html_body"]
@ -1140,7 +1022,9 @@ async def test_multi_threshold_team_member_alert_renders_member_template_per_tea
@pytest.mark.asyncio
async def test_no_map_preserves_old_single_threshold(
base_email_logger, mock_send_email
redis_email_logger: BaseEmailLogger,
mock_send_email,
alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
):
"""Test that without max_budget_alert_emails, the old 80% single-threshold path works"""
user_info = CallInfo(
@ -1151,21 +1035,13 @@ async def test_no_map_preserves_old_single_threshold(
event_group=Litellm_EntityType.USER,
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
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 redis_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_increment_cache.call_args[1]["key"]
assert cache_key == "email_budget_alerts:max_budget_alert:test_user"
assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:test_user") == b"1"
CUSTOM_SIGNATURE = "<div>Best,<br/>The Acme Platform Team</div>"
@ -1186,10 +1062,13 @@ async def test_send_soft_budget_alert_email_uses_custom_signature(
max_budget=200.0,
soft_budget=100.0,
)
with mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
), patch("litellm.proxy.proxy_server.premium_user", True):
with (
mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
),
patch("litellm.proxy.proxy_server.premium_user", True),
):
await base_email_logger.send_soft_budget_alert_email(event)
html_body = mock_send_email.call_args[1]["html_body"]
@ -1213,10 +1092,13 @@ async def test_send_team_soft_budget_alert_email_uses_custom_signature(
team_alias="Acme",
alert_emails=["teamlead@example.com"],
)
with mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
), patch("litellm.proxy.proxy_server.premium_user", True):
with (
mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
),
patch("litellm.proxy.proxy_server.premium_user", True),
):
await base_email_logger.send_team_soft_budget_alert_email(event)
html_body = mock_send_email.call_args[1]["html_body"]
@ -1238,10 +1120,13 @@ async def test_send_max_budget_alert_email_single_recipient_uses_custom_signatur
spend=165.0,
max_budget=200.0,
)
with mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
), patch("litellm.proxy.proxy_server.premium_user", True):
with (
mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
),
patch("litellm.proxy.proxy_server.premium_user", True),
):
await base_email_logger.send_max_budget_alert_email(event)
html_body = mock_send_email.call_args[1]["html_body"]
@ -1263,10 +1148,13 @@ async def test_send_max_budget_alert_email_multi_recipient_uses_custom_signature
spend=165.0,
max_budget=200.0,
)
with mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
), patch("litellm.proxy.proxy_server.premium_user", True):
with (
mock.patch.dict(
os.environ,
{"PROXY_BASE_URL": "http://test.com", "EMAIL_SIGNATURE": CUSTOM_SIGNATURE},
),
patch("litellm.proxy.proxy_server.premium_user", True),
):
await base_email_logger.send_max_budget_alert_email(
event, threshold_pct=75, recipient_emails=["a@example.com", "b@example.com"]
)
@ -1354,12 +1242,8 @@ async def test_budget_alert_no_duplicate_on_concurrent_crossing(
with mock.patch.object(base_email_logger, send_method, side_effect=slow_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await asyncio.gather(
base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
),
base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
),
base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)),
base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)),
)
assert len(sends) == 1
@ -1388,42 +1272,37 @@ async def test_budget_alert_failed_send_releases_claim_for_retry(
with mock.patch.object(base_email_logger, send_method, side_effect=flaky_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
)
await base_email_logger.budget_alerts(
type=alert_type, user_info=_budget_alert_user_info(ci_kwargs)
)
await base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs))
await base_email_logger.budget_alerts(type=alert_type, user_info=_budget_alert_user_info(ci_kwargs))
assert len(attempts) == 2
@pytest.mark.asyncio
async def test_budget_alert_release_failure_does_not_propagate(base_email_logger):
async def test_budget_alert_release_failure_does_not_propagate(
redis_email_logger: BaseEmailLogger, alert_redis: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer]
):
"""If the send fails and releasing the claim also fails (transient cache
error), budget_alerts must swallow it and still log the send failure rather
than letting the exception escape the fire-and-forget task."""
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
mock_cache.async_delete_cache = mock.AsyncMock(
side_effect=RuntimeError("cache backend unavailable")
)
base_email_logger.internal_usage_cache = mock_cache
async def failing_send(*args, **kwargs):
alert_redis[2].connected = False
raise ValueError("smtp backend down")
with mock.patch.object(
base_email_logger, "send_max_budget_alert_email", side_effect=failing_send
):
with mock.patch.object(redis_email_logger, "send_email", side_effect=failing_send):
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
# Must not raise even though both the send and the release fail.
await base_email_logger.budget_alerts(
await redis_email_logger.budget_alerts(
type="max_budget_alert",
user_info=_budget_alert_user_info(dict(max_budget=100.0, spend=85.0)),
)
mock_cache.async_delete_cache.assert_awaited_once()
alert_redis[2].connected = True
assert await alert_redis[1].get("alerts:email_budget_alerts:max_budget_alert:hashed_key_1") == b"1"
assert (
local_get(redis_email_logger.internal_usage_cache, "email_budget_alerts:max_budget_alert:hashed_key_1") is None
)
class _RecordingEmailLogger(BaseEmailLogger):
@ -1443,12 +1322,12 @@ class _UserTable:
return self._rows.get(where["user_id"])
def _prisma_client_with_users(*rows):
def _prisma_client_with_users(*rows: LiteLLM_UserTable) -> SimpleNamespace:
table: Final = _UserTable({row.user_id: row for row in rows})
return SimpleNamespace(db=SimpleNamespace(litellm_usertable=table))
def _key_created_event(user_id):
def _key_created_event(user_id: str) -> SendKeyCreatedEmailEvent:
return SendKeyCreatedEmailEvent(
user_id=user_id,
user_email=None,
@ -1463,9 +1342,7 @@ def _key_created_event(user_id):
@pytest.mark.asyncio
async def test_key_created_email_goes_to_the_address_stored_for_the_user(monkeypatch):
stored_user: Final = LiteLLM_UserTable(
user_id="user-1", user_email="stored@example.com"
)
stored_user: Final = LiteLLM_UserTable(user_id="user-1", user_email="stored@example.com")
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
_prisma_client_with_users(stored_user),
@ -1481,9 +1358,7 @@ async def test_key_created_email_goes_to_the_address_stored_for_the_user(monkeyp
async def test_key_created_email_is_refused_for_a_user_the_database_does_not_know(
monkeypatch,
):
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client", _prisma_client_with_users()
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_client_with_users())
email_logger: Final = _RecordingEmailLogger()
with pytest.raises(ValueError, match="User email not found for user_id: user-2"):

293
tests/unit/redis/stores/test_email_alert_claims.py Normal file
View file

@ -0,0 +1,293 @@
import asyncio
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass
from datetime import timedelta
from typing import Final, Literal
from unittest.mock import patch
import fakeredis
import pytest
import pytest_asyncio
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import BaseEmailLogger
from redis.exceptions import ConnectionError as RedisConnectionError
from redis.exceptions import ResponseError
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import EMAIL_BUDGET_ALERT_TTL
from litellm.proxy._types import CallInfo, Litellm_EntityType
from litellm.redis.stores.email_alert_claims import SoftBudgetAlert, TransportEmailAlertClaimsStore
from litellm.redis.transport.contract import RedisFailure
from tests._support.fake_redis import async_client, fake_redis_cache, legacy_increment, local_get
@dataclass(frozen=True, slots=True)
class SentEmail:
recipients: tuple[str, ...]
subject: str
class EmailSink(BaseEmailLogger):
def __init__(self, cache: DualCache, on_send: Callable[[], Awaitable[None]] | None = None) -> None:
super().__init__(internal_usage_cache=cache) # pyright: ignore[reportUnknownMemberType] # legacy logger init is untyped
self.sent: tuple[SentEmail, ...] = ()
self.on_send = on_send
async def send_email(self, from_email: str, to_email: list[str], subject: str, html_body: str) -> None:
self.sent += (SentEmail(tuple(sorted(to_email)), subject),)
if self.on_send is not None:
await self.on_send()
@pytest_asyncio.fixture(loop_scope="function")
async def redis_env() -> AsyncIterator[tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer]]:
with patch("time.time", return_value=1700000000.0), patch.dict("os.environ", {"PROXY_BASE_URL": "http://test.com"}):
server: Final = fakeredis.FakeServer()
cache: Final = fake_redis_cache(server, namespace="alerts")
raw: Final = fakeredis.FakeAsyncRedis(server=server)
try:
yield cache, raw, server
finally:
await raw.aclose() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # fakeredis stubs predate aclose
_AlertType = Literal[
"token_budget", "soft_budget", "max_budget_alert", "user_budget", "team_budget", "organization_budget"
]
_ALERTS: Final[tuple[tuple[_AlertType, CallInfo, str], ...]] = (
(
"soft_budget",
CallInfo(
user_id="user:雪",
user_email="owner@example.com",
spend=60,
soft_budget=50,
event_group=Litellm_EntityType.USER,
),
"email_budget_alerts:soft_budget_crossed:user:雪",
),
(
"max_budget_alert",
CallInfo(
token="key:雪", user_email="owner@example.com", spend=85, max_budget=100, event_group=Litellm_EntityType.KEY
),
"email_budget_alerts:max_budget_alert:key:雪",
),
(
"max_budget_alert",
CallInfo(
user_id="member:雪",
team_id="team:a",
user_email="owner@example.com",
spend=85,
max_budget=100,
max_budget_alert_emails={"50": ["finance@example.com"]},
event_group=Litellm_EntityType.TEAM_MEMBER,
),
"email_budget_alerts:max_budget_alert:50:team_member:member:雪:team:a",
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("alert_type,info,key", _ALERTS)
async def test_old_and_new_claims_share_numeric_bytes_without_refreshing_expiry(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
alert_type: Literal["soft_budget", "max_budget_alert"],
info: CallInfo,
key: str,
) -> None:
cache, raw, _ = redis_env
logger: Final = EmailSink(DualCache(redis_cache=cache))
assert await legacy_increment(cache, key, 0.25, ttl=37) == 0.25
await logger.budget_alerts(type=alert_type, user_info=info)
assert logger.sent == ()
assert await raw.get(f"alerts:{key}") == b"1.25"
assert await raw.pttl(f"alerts:{key}") == 37000
await raw.delete(f"alerts:{key}")
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == 1
assert await raw.get(f"alerts:{key}") == b"1"
assert await raw.pttl(f"alerts:{key}") == EMAIL_BUDGET_ALERT_TTL * 1000
assert await legacy_increment(cache, key, 1, ttl=2) == 2
assert await raw.get(f"alerts:{key}") == b"2"
assert await raw.pttl(f"alerts:{key}") == EMAIL_BUDGET_ALERT_TTL * 1000
assert await raw.keys("*") == [f"alerts:{key}".encode()]
@pytest.mark.asyncio
@pytest.mark.parametrize("alert_type,info,key", _ALERTS)
async def test_concurrent_loggers_claim_once_and_failed_send_deletes_both_tiers(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
alert_type: Literal["soft_budget", "max_budget_alert"],
info: CallInfo,
key: str,
) -> None:
cache, raw, _ = redis_env
first: Final = EmailSink(DualCache(redis_cache=cache))
second: Final = EmailSink(DualCache(redis_cache=cache))
await asyncio.gather(
first.budget_alerts(type=alert_type, user_info=info), second.budget_alerts(type=alert_type, user_info=info)
)
assert len(first.sent) + len(second.sent) == 1
assert await raw.get(f"alerts:{key}") == b"2"
await raw.delete(f"alerts:{key}")
async def fail_send() -> None:
raise ValueError("email backend failed")
failed: Final = EmailSink(DualCache(redis_cache=cache), fail_send)
await failed.budget_alerts(type=alert_type, user_info=info)
assert len(failed.sent) == 1
assert await raw.get(f"alerts:{key}") is None
assert local_get(failed.internal_usage_cache, key) is None
retry: Final = EmailSink(failed.internal_usage_cache)
await retry.budget_alerts(type=alert_type, user_info=info)
assert len(retry.sent) == 1
assert await raw.get(f"alerts:{key}") == b"1"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"seed,sends,stored", [(b"-1", 1, b"0"), (b"0", 1, b"1"), (b"1", 0, b"2"), (b"broken", 1, b"broken")]
)
async def test_counter_result_and_server_error_keep_existing_send_policy(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
seed: bytes,
sends: int,
stored: bytes,
) -> None:
cache, raw, _ = redis_env
logger: Final = EmailSink(DualCache(redis_cache=cache))
alert_type, info, key = _ALERTS[0]
await raw.set(f"alerts:{key}", seed)
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == sends
assert await raw.get(f"alerts:{key}") == stored
assert await cache.transport().ping() is True
assert await raw.pttl(f"alerts:{key}") == (-1 if seed == b"broken" else EMAIL_BUDGET_ALERT_TTL * 1000)
@pytest.mark.asyncio
@pytest.mark.parametrize("local_cache", [True, False])
async def test_outage_and_open_breaker_use_local_count_or_none(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
local_cache: bool,
) -> None:
cache, raw, server = redis_env
dual: Final = DualCache(redis_cache=cache)
if not local_cache:
dual.in_memory_cache = None # pyright: ignore[reportAttributeAccessIssue] # the logger tolerates a cache without a local layer
logger: Final = EmailSink(dual)
alert_type, info, key = _ALERTS[0]
server.connected = False
for _ in range(6):
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == (1 if local_cache else 6)
result: Final = await cache.transport().get(key)
assert isinstance(result, RedisFailure)
assert result.kind == "breaker_open"
server.connected = True
assert await raw.get(f"alerts:{key}") is None
@pytest.mark.asyncio
async def test_applied_increment_with_lost_reply_falls_back_without_retrying_or_rolling_back(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
) -> None:
cache, raw, _ = redis_env
logger: Final = EmailSink(DualCache(redis_cache=cache))
alert_type, info, key = _ALERTS[0]
client: Final = async_client(cache)
pipeline: Final = client.pipeline(transaction=False) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # fakeredis pipeline is untyped
execute: Final = pipeline.execute
async def lose_reply(raise_on_error: bool = True) -> list[object]:
await execute(raise_on_error=raise_on_error)
assert local_get(logger.internal_usage_cache, key) == 1
raise RedisConnectionError("reply lost after write applied")
with (
patch.object(client, "pipeline", return_value=pipeline),
patch.object(pipeline, "execute", side_effect=lose_reply),
):
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == 1
assert await raw.get(f"alerts:{key}") == b"1"
assert await raw.pttl(f"alerts:{key}") == -1
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == 1
assert await raw.get(f"alerts:{key}") == b"2"
assert await raw.pttl(f"alerts:{key}") == EMAIL_BUDGET_ALERT_TTL * 1000
@pytest.mark.asyncio
async def test_expiry_failure_leaves_the_increment_and_uses_local_fallback(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
) -> None:
cache, raw, _ = redis_env
logger: Final = EmailSink(DualCache(redis_cache=cache))
alert_type, info, key = _ALERTS[0]
client: Final = async_client(cache)
with patch.object(client, "expire", side_effect=RedisConnectionError("lost before expiry")):
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == 1
assert await raw.get(f"alerts:{key}") == b"1"
assert await raw.pttl(f"alerts:{key}") == -1
@pytest.mark.asyncio
async def test_rejected_expiry_uses_local_result_instead_of_applied_redis_count(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
) -> None:
cache, raw, _ = redis_env
logger: Final = EmailSink(DualCache(redis_cache=cache))
alert_type, info, key = _ALERTS[0]
await raw.set(f"alerts:{key}", b"1")
with patch.object(cache.init_async_client(), "expire", side_effect=ResponseError("NOPERM expire denied")):
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == 1
assert await raw.get(f"alerts:{key}") == b"2"
assert await raw.pttl(f"alerts:{key}") == -1
@pytest.mark.asyncio
async def test_unknown_applied_release_deletes_locally_and_keeps_retry_policy(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
) -> None:
cache, raw, _ = redis_env
async def fail_send() -> None:
raise ValueError("email backend failed")
logger: Final = EmailSink(DualCache(redis_cache=cache), fail_send)
alert_type, info, key = _ALERTS[0]
client: Final = async_client(cache)
delete: Final = client.delete
async def lose_reply(*keys: str) -> int:
await delete(*keys)
raise RedisConnectionError("release reply lost after deletion")
with patch.object(client, "delete", side_effect=lose_reply):
await logger.budget_alerts(type=alert_type, user_info=info)
assert len(logger.sent) == 1
assert await raw.get(f"alerts:{key}") is None
assert local_get(logger.internal_usage_cache, key) is None
retry: Final = EmailSink(logger.internal_usage_cache)
await retry.budget_alerts(type=alert_type, user_info=info)
assert len(retry.sent) == 1
assert await raw.get(f"alerts:{key}") == b"1"
@pytest.mark.asyncio
@pytest.mark.parametrize("seconds", [0, -1])
async def test_nonpositive_expiry_deletes_the_counter_like_legacy_expire(
redis_env: tuple[RedisCache, fakeredis.FakeAsyncRedis, fakeredis.FakeServer],
seconds: int,
) -> None:
cache, raw, _ = redis_env
store: Final = TransportEmailAlertClaimsStore(cache.transport())
assert await store.claim(SoftBudgetAlert("expired"), ttl=timedelta(seconds=seconds)) == 1
with patch("time.time", return_value=1700000000.001):
assert await raw.get("alerts:email_budget_alerts:soft_budget_crossed:expired") is None