diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index c648f7faea0..6eed6e6b66e 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -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, ): diff --git a/litellm/redis/stores/email_alert_claims.py b/litellm/redis/stores/email_alert_claims.py new file mode 100644 index 00000000000..c44d1390332 --- /dev/null +++ b/litellm/redis/stores/email_alert_claims.py @@ -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 diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py index 6cfdb75fd01..11323a87ca8 100644 --- a/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -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 = "
Best,
The Acme Platform Team
" @@ -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"): diff --git a/tests/unit/redis/stores/test_email_alert_claims.py b/tests/unit/redis/stores/test_email_alert_claims.py new file mode 100644 index 00000000000..71d80ad93e1 --- /dev/null +++ b/tests/unit/redis/stores/test_email_alert_claims.py @@ -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