unit testing for provider budgets

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