Spaces:
Configuration error
Configuration error
File size: 6,724 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
import datetime
import json
import os
import sys
import unittest
from typing import List, Optional, Tuple
from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system-path
import litellm
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.proxy._types import CallInfo, Litellm_EntityType
class TestSlackAlerting(unittest.TestCase):
def setUp(self):
self.slack_alerting = SlackAlerting()
def test_get_percent_of_max_budget_left(self):
# Test case 1: When max_budget is None
user_info = CallInfo(
max_budget=None, spend=50.0, event_group=Litellm_EntityType.KEY
)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.0)
# Test case 2: When max_budget is 0
user_info = CallInfo(
max_budget=0.0, spend=50.0, event_group=Litellm_EntityType.KEY
)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.0)
# Test case 3: When spend is less than max_budget
user_info = CallInfo(
max_budget=100.0, spend=75.0, event_group=Litellm_EntityType.KEY
)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.25)
# Test case 4: When spend equals max_budget
user_info = CallInfo(
max_budget=100.0, spend=100.0, event_group=Litellm_EntityType.KEY
)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, 0.0)
# Test case 5: When spend exceeds max_budget
user_info = CallInfo(
max_budget=100.0, spend=120.0, event_group=Litellm_EntityType.KEY
)
result = self.slack_alerting._get_percent_of_max_budget_left(user_info)
self.assertEqual(result, -0.2)
def test_get_event_and_event_message_max_budget(self):
# Initial setup with no event
event = None
event_message = "Test Message: "
# Test case 1: When spend exceeds max_budget
user_info = CallInfo(
max_budget=100.0,
spend=120.0,
soft_budget=None,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=event, event_message=event_message
)
self.assertEqual(event, "budget_crossed")
self.assertTrue("Budget Crossed" in event_message)
# Test case 2: When 5% of max_budget is left
user_info = CallInfo(
max_budget=100.0,
spend=95.0,
soft_budget=None,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=event, event_message=event_message
)
self.assertEqual(event, "threshold_crossed")
self.assertTrue("5% Threshold Crossed" in event_message)
# Test case 3: When 15% of max_budget is left
user_info = CallInfo(
max_budget=100.0,
spend=85.0,
soft_budget=None,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=event, event_message=event_message
)
self.assertEqual(event, "threshold_crossed")
self.assertTrue("15% Threshold Crossed" in event_message)
def test_get_event_and_event_message_soft_budget(self):
# Initial setup with no event
event = None
event_message = "Test Message: "
# Test case 1: When spend exceeds soft_budget
user_info = CallInfo(
max_budget=None,
spend=120.0,
soft_budget=100.0,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=event, event_message=event_message
)
self.assertEqual(event, "soft_budget_crossed")
self.assertTrue("Total Soft Budget" in event_message)
# Test case 2: When spend is less than soft_budget
user_info = CallInfo(
max_budget=None,
spend=90.0,
soft_budget=100.0,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=None, event_message=event_message
)
print("got event", event)
print("got event_message", event_message)
self.assertEqual(event, None) # No event should be triggered
def test_get_event_and_event_message_both_budgets(self):
# Initial setup with no event
event = None
event_message = "Test Message: "
# Test case 1: When spend exceeds both max_budget and soft_budget
user_info = CallInfo(
max_budget=150.0,
spend=160.0,
soft_budget=100.0,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=event, event_message=event_message
)
# budget_crossed has higher priority
self.assertEqual(event, "budget_crossed")
self.assertTrue("Budget Crossed" in event_message)
# Test case 2: When spend exceeds soft_budget but not max_budget
user_info = CallInfo(
max_budget=150.0,
spend=120.0,
soft_budget=100.0,
event_group=Litellm_EntityType.KEY,
)
event, event_message = self.slack_alerting._get_event_and_event_message(
user_info=user_info, event=event, event_message=event_message
)
self.assertEqual(event, "soft_budget_crossed")
self.assertTrue("Total Soft Budget" in event_message)
# Calling update_values with alerting args should try to start the periodic task
@patch("asyncio.create_task")
def test_update_values_starts_periodic_task(self, mock_create_task):
# Make it do nothing (or return a dummy future)
mock_create_task.return_value = AsyncMock() # prevents awaiting errors
assert self.slack_alerting.periodic_started == False
self.slack_alerting.update_values(alerting_args={"slack_alerting": "True"})
assert self.slack_alerting.periodic_started == True
|