From 33a0744abe30f71baa43c4a690810c2240782768 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 23 Nov 2024 16:24:13 -0800 Subject: [PATCH] test_in_memory_redis_sync_e2e --- .../test_router_provider_budgets.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/local_testing/test_router_provider_budgets.py b/tests/local_testing/test_router_provider_budgets.py index 3ef8de64b..415a93f4d 100644 --- a/tests/local_testing/test_router_provider_budgets.py +++ b/tests/local_testing/test_router_provider_budgets.py @@ -2,6 +2,7 @@ import sys, os, asyncio, time, random from datetime import datetime import traceback from dotenv import load_dotenv +from httpx import delete load_dotenv() import os, copy @@ -276,6 +277,11 @@ async def test_in_memory_redis_sync_e2e(): 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 = { @@ -336,3 +342,13 @@ async def test_in_memory_redis_sync_e2e(): # 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, + )