diff --git a/tests/local_testing/test_router_provider_budgets.py b/tests/local_testing/test_router_provider_budgets.py index c16fc5d5c..430550632 100644 --- a/tests/local_testing/test_router_provider_budgets.py +++ b/tests/local_testing/test_router_provider_budgets.py @@ -17,7 +17,7 @@ from litellm.types.router import ( ProviderBudgetConfigType, ProviderBudgetInfo, ) -from litellm.caching.caching import DualCache +from litellm.caching.caching import DualCache, RedisCache import logging from litellm._logging import verbose_router_logger import litellm @@ -296,3 +296,183 @@ async def test_prometheus_metric_tracking(): # Verify the mock was called correctly mock_prometheus.track_provider_remaining_budget.assert_called_once() + + +@pytest.mark.asyncio +async def test_handle_new_budget_window(): + """ + Test _handle_new_budget_window helper method + + Current + """ + cleanup_redis() + provider_budget = ProviderBudgetLimiting( + router_cache=DualCache(), provider_budget_config={} + ) + + spend_key = "provider_spend:openai:7d" + start_time_key = "provider_budget_start_time:openai" + current_time = 1000.0 + response_cost = 0.5 + ttl_seconds = 86400 # 1 day + + # Test handling new budget window + new_start_time = await provider_budget._handle_new_budget_window( + spend_key=spend_key, + start_time_key=start_time_key, + current_time=current_time, + response_cost=response_cost, + ttl_seconds=ttl_seconds, + ) + + assert new_start_time == current_time + + # Verify the spend was set correctly + spend = await provider_budget.router_cache.async_get_cache(spend_key) + print("spend in cache for key", spend_key, "is", spend) + assert float(spend) == response_cost + + # Verify start time was set correctly + start_time = await provider_budget.router_cache.async_get_cache(start_time_key) + print("start time in cache for key", start_time_key, "is", start_time) + assert float(start_time) == current_time + + +@pytest.mark.asyncio +async def test_get_or_set_budget_start_time(): + """ + Test _get_or_set_budget_start_time helper method + + scenario 1: no existing start time in cache, should return current time + scenario 2: existing start time in cache, should return existing start time + """ + cleanup_redis() + provider_budget = ProviderBudgetLimiting( + router_cache=DualCache(), provider_budget_config={} + ) + + start_time_key = "test_start_time" + current_time = 1000.0 + ttl_seconds = 86400 # 1 day + + # When there is no existing start time, we should set it to the current time + start_time = await provider_budget._get_or_set_budget_start_time( + start_time_key=start_time_key, + current_time=current_time, + ttl_seconds=ttl_seconds, + ) + print("budget start time when no existing start time is in cache", start_time) + assert start_time == current_time + + # When there is an existing start time, we should return it even if the current time is later + new_current_time = 2000.0 + existing_start_time = await provider_budget._get_or_set_budget_start_time( + start_time_key=start_time_key, + current_time=new_current_time, + ttl_seconds=ttl_seconds, + ) + print( + "budget start time when existing start time is in cache, but current time is later", + existing_start_time, + ) + assert existing_start_time == current_time # Should return the original start time + + +@pytest.mark.asyncio +async def test_increment_spend_in_current_window(): + """ + Test _increment_spend_in_current_window helper method + + Expected behavior: + - Increment the spend in memory cache + - Queue the increment operation to Redis + """ + cleanup_redis() + provider_budget = ProviderBudgetLimiting( + router_cache=DualCache(), provider_budget_config={} + ) + + spend_key = "provider_spend:openai:1d" + response_cost = 0.5 + ttl = 86400 # 1 day + + # Set initial spend + await provider_budget.router_cache.async_set_cache( + key=spend_key, value=1.0, ttl=ttl + ) + + # Test incrementing spend + await provider_budget._increment_spend_in_current_window( + spend_key=spend_key, + response_cost=response_cost, + ttl=ttl, + ) + + # Verify the spend was incremented correctly in memory + spend = await provider_budget.router_cache.async_get_cache(spend_key) + assert float(spend) == 1.5 + + # Verify the increment operation was queued for Redis + print( + "redis_increment_operation_queue", + provider_budget.redis_increment_operation_queue, + ) + assert len(provider_budget.redis_increment_operation_queue) == 1 + queued_op = provider_budget.redis_increment_operation_queue[0] + assert queued_op["key"] == spend_key + assert queued_op["increment_value"] == response_cost + assert queued_op["ttl"] == ttl + + +@pytest.mark.asyncio +async def test_sync_in_memory_spend_with_redis(): + """ + Test _sync_in_memory_spend_with_redis helper method + + Expected behavior: + - Push all provider spend increments to Redis + - Fetch all current provider spend from Redis to update in-memory cache + """ + cleanup_redis() + provider_budget_config = { + "openai": ProviderBudgetInfo(time_period="1d", budget_limit=100), + "anthropic": ProviderBudgetInfo(time_period="1d", budget_limit=200), + } + + provider_budget = ProviderBudgetLimiting( + router_cache=DualCache( + redis_cache=RedisCache( + host=os.getenv("REDIS_HOST"), + port=int(os.getenv("REDIS_PORT")), + password=os.getenv("REDIS_PASSWORD"), + ) + ), + provider_budget_config=provider_budget_config, + ) + + # Set some values in Redis + spend_key_openai = "provider_spend:openai:1d" + spend_key_anthropic = "provider_spend:anthropic:1d" + + await provider_budget.router_cache.redis_cache.async_set_cache( + key=spend_key_openai, value=50.0 + ) + await provider_budget.router_cache.redis_cache.async_set_cache( + key=spend_key_anthropic, value=75.0 + ) + + # Test syncing with Redis + await provider_budget._sync_in_memory_spend_with_redis() + + # Verify in-memory cache was updated + openai_spend = await provider_budget.router_cache.in_memory_cache.async_get_cache( + spend_key_openai + ) + anthropic_spend = ( + await provider_budget.router_cache.in_memory_cache.async_get_cache( + spend_key_anthropic + ) + ) + + assert float(openai_spend) == 50.0 + assert float(anthropic_spend) == 75.0