From 0233587fbafdfd98c6aab1fdd77f243d54ac45cd Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 19 Mar 2026 20:07:52 -0700 Subject: [PATCH] test(caching): add circuit breaker unit tests --- tests/test_litellm/caching/test_dual_cache.py | 218 ++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 tests/test_litellm/caching/test_dual_cache.py diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py new file mode 100644 index 00000000000..ca01096b048 --- /dev/null +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -0,0 +1,218 @@ +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache + + +@pytest.mark.asyncio +async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads(): + dual_cache = DualCache( + redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 + ) + keys = ["shared_a", "shared_b"] + start_gate = asyncio.Event() + + async def _mock_async_batch_get_cache(key_list, parent_otel_span=None): + await asyncio.sleep(0.05) + return {k: None for k in key_list} + + with patch.object( + dual_cache.redis_cache, + "async_batch_get_cache", + new=AsyncMock(side_effect=_mock_async_batch_get_cache), + ) as mock_async_batch_get_cache: + + async def worker(): + await start_gate.wait() + return await dual_cache.async_batch_get_cache(keys=keys) + + tasks = [asyncio.create_task(worker()) for _ in range(50)] + start_gate.set() + await asyncio.gather(*tasks) + + assert mock_async_batch_get_cache.call_count == 1 + + +@pytest.mark.asyncio +async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_error(): + dual_cache = DualCache( + redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 + ) + keys = ["shared_a", "shared_b"] + + with patch.object( + dual_cache.redis_cache, + "async_batch_get_cache", + new=AsyncMock(side_effect=RuntimeError("redis unavailable")), + ) as mock_async_batch_get_cache: + first_result = await dual_cache.async_batch_get_cache(keys=keys) + second_result = await dual_cache.async_batch_get_cache(keys=keys) + + assert first_result is None + assert second_result is None + assert mock_async_batch_get_cache.call_count == 2 + assert "shared_a" not in dual_cache.last_redis_batch_access_time + assert "shared_b" not in dual_cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl(): + """ + Test that async_set_cache injects default_in_memory_ttl into kwargs + when no explicit ttl is provided, matching the sync set_cache behavior. + + Regression test for: async_set_cache was missing the TTL injection that + sync set_cache has, causing InMemoryCache to use its own default_ttl (600s) + instead of DualCache's default_in_memory_ttl. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + await dual_cache.async_set_cache(key="test_key", value="test_value") + after = time.time() + + # The TTL stored should reflect default_in_memory_ttl (60s), not + # InMemoryCache's default_ttl (600s) + expiry = in_memory_cache.ttl_dict["test_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_respects_explicit_ttl(): + """ + Test that async_set_cache does NOT override an explicitly provided ttl. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + await dual_cache.async_set_cache(key="test_key", value="test_value", ttl=30) + after = time.time() + + # The explicit ttl=30 should be used, not default_in_memory_ttl (60) + expiry = in_memory_cache.ttl_dict["test_key"] + assert expiry >= before + 30 + assert expiry <= after + 30 + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_pipeline_injects_default_in_memory_ttl(): + """ + Test that async_set_cache_pipeline injects default_in_memory_ttl into kwargs + when no explicit ttl is provided. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + cache_list = [("key_a", "value_a"), ("key_b", "value_b")] + + before = time.time() + await dual_cache.async_set_cache_pipeline(cache_list=cache_list) + after = time.time() + + for key in ["key_a", "key_b"]: + expiry = in_memory_cache.ttl_dict[key] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_sync_and_async_set_cache_use_same_ttl(): + """ + Test that sync set_cache and async async_set_cache produce the same TTL + when no explicit ttl is provided, ensuring parity between the two paths. + """ + in_memory_sync = InMemoryCache(default_ttl=600) + dual_cache_sync = DualCache( + in_memory_cache=in_memory_sync, + default_in_memory_ttl=60, + ) + + in_memory_async = InMemoryCache(default_ttl=600) + dual_cache_async = DualCache( + in_memory_cache=in_memory_async, + default_in_memory_ttl=60, + ) + + dual_cache_sync.set_cache(key="test_key", value="test_value") + await dual_cache_async.async_set_cache(key="test_key", value="test_value") + + sync_expiry = in_memory_sync.ttl_dict["test_key"] + async_expiry = in_memory_async.ttl_dict["test_key"] + + # Both should use default_in_memory_ttl=60, so their expiry times + # should be within a small tolerance of each other + assert abs(sync_expiry - async_expiry) < 1.0 + + +def test_circuit_breaker_opens_after_threshold(): + """Circuit opens after N consecutive Redis failures.""" + from litellm.caching.redis_cache import RedisCircuitBreaker + + cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + for _ in range(3): + cb.record_failure() + + assert cb._state == "open" + + +@pytest.mark.asyncio +async def test_circuit_breaker_open_skips_redis(): + """When circuit is open, the guard decorator raises immediately without calling the method.""" + from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _redis_circuit_breaker_guard, + ) + + class FakeRedis: + def __init__(self): + self._circuit_breaker = RedisCircuitBreaker( + failure_threshold=3, recovery_timeout=60 + ) + self._circuit_breaker._state = "open" + self._circuit_breaker._opened_at = time.time() + self.call_count = 0 + + @_redis_circuit_breaker_guard + async def do_thing(self): + self.call_count += 1 + return "result" + + fr = FakeRedis() + with pytest.raises(Exception, match="circuit breaker is open"): + await fr.do_thing() + + assert fr.call_count == 0 # method body never executed + + +def test_circuit_breaker_closes_on_recovery(): + """After recovery_timeout expires, probe is allowed and success closes the circuit.""" + from litellm.caching.redis_cache import RedisCircuitBreaker + + cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + cb._state = "open" + cb._opened_at = time.time() - 9999 # recovery timeout long expired + + # is_open() should return False to allow a probe through, and transition to HALF_OPEN + assert cb.is_open() is False + assert cb._state == "half_open" + + # Successful probe closes the circuit + cb.record_success() + assert cb._state == "closed"