test_in_memory_redis_sync_e2e

BerriAI · Nov 24, 2024 · 33a0744 · 33a0744
1 parent 5f04c04
commit 33a0744
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/tests/local_testing/test_router_provider_budgets.py b/tests/local_testing/test_router_provider_budgets.py
@@ -2,6 +2,7 @@
 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,
+    )