diff --git a/tests/local_testing/test_router_provider_budgets.py b/tests/local_testing/test_router_provider_budgets.py index bd970ea1b..87dc46a44 100644 --- a/tests/local_testing/test_router_provider_budgets.py +++ b/tests/local_testing/test_router_provider_budgets.py @@ -25,6 +25,25 @@ import litellm verbose_router_logger.setLevel(logging.DEBUG) +@pytest.fixture(autouse=True) +async def cleanup_redis(): + """Cleanup Redis cache before each test""" + try: + import redis + + redis_client = redis.Redis( + host=os.getenv("REDIS_HOST"), + port=int(os.getenv("REDIS_PORT")), + password=os.getenv("REDIS_PASSWORD"), + ) + # Delete all provider spend keys + for key in redis_client.scan_iter("provider_spend:*"): + redis_client.delete(key) + except Exception as e: + print(f"Error cleaning up Redis: {str(e)}") + yield + + @pytest.mark.asyncio async def test_provider_budgets_e2e_test(): """ @@ -272,87 +291,3 @@ 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_in_memory_redis_sync_e2e(): - """ - Test that the in-memory cache gets properly synced with Redis values through the periodic sync mechanism - - Critical test for using provider budgets in a multi-instance environment - """ - original_sync_interval = getattr( - litellm.router_strategy.provider_budgets, "DEFAULT_REDIS_SYNC_INTERVAL" - ) - - # Modify for test - setattr(litellm.router_strategy.provider_budgets, "DEFAULT_REDIS_SYNC_INTERVAL", 2) - - provider_budget_config = { - "openai": ProviderBudgetInfo(time_period="1d", budget_limit=100), - } - - router = Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo-very-new", - "litellm_params": { - "model": "openai/gpt-3.5-turbo", - }, - }, - ], - provider_budget_config=provider_budget_config, - redis_host=os.getenv("REDIS_HOST"), - redis_port=int(os.getenv("REDIS_PORT")), - redis_password=os.getenv("REDIS_PASSWORD"), - ) - - if router.cache is None: - raise ValueError("Router cache is not initialized") - if router.cache.redis_cache is None: - raise ValueError("Redis cache is not initialized") - - # Get the ProviderBudgetLimiting instance - spend_key = "provider_spend:openai:1d" - - # Set initial values - test_spend_1 = 50.0 - await router.cache.redis_cache.async_set_cache(key=spend_key, value=test_spend_1) - - # Make a completion call to trigger spend tracking - response = await router.acompletion( - model="gpt-3.5-turbo-very-new", - messages=[{"role": "user", "content": "Hello"}], - mock_response="Hello there!", - ) - - # Wait for periodic sync (should be less than DEFAULT_REDIS_SYNC_INTERVAL) - await asyncio.sleep(2.5) - - # Verify in-memory cache matches Redis - in_memory_spend = float(router.cache.in_memory_cache.get_cache(spend_key) or 0) - redis_spend = float(await router.cache.redis_cache.async_get_cache(spend_key) or 0) - assert ( - abs(in_memory_spend - redis_spend) < 0.01 - ) # Allow for small floating point differences - - # Update Redis with new value from a "different litellm proxy instance" - test_spend_2 = 75.0 - await router.cache.redis_cache.async_set_cache(key=spend_key, value=test_spend_2) - - # Wait for periodic sync - await asyncio.sleep(2.5) - - # Verify in-memory cache was updated - in_memory_spend = float(router.cache.in_memory_cache.get_cache(spend_key) or 0) - assert abs(in_memory_spend - test_spend_2) < 0.01 - - # clean up key from router cache - await router.cache.async_delete_cache(spend_key) - - # Restore original value - setattr( - litellm.router_strategy.provider_budgets, - "DEFAULT_REDIS_SYNC_INTERVAL", - original_sync_interval, - )