mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix: guard cache set and numeric operations
This commit is contained in:
parent
d5d8bca2d1
commit
8c938fd3e4
2 changed files with 20 additions and 0 deletions
|
|
@ -193,6 +193,8 @@ class InMemoryCache(BaseCache):
|
|||
"""
|
||||
# get the value
|
||||
init_value: Final = self.get_cache(key=key) or set()
|
||||
if not isinstance(init_value, set):
|
||||
raise TypeError("Cached value is not a set")
|
||||
for val in value:
|
||||
init_value.add(val)
|
||||
self.set_cache(key, init_value, ttl=ttl)
|
||||
|
|
@ -234,6 +236,8 @@ class InMemoryCache(BaseCache):
|
|||
with self._increment_lock:
|
||||
# keep read-modify-write atomic
|
||||
init_value: Final = self.get_cache(key=key) or 0
|
||||
if not isinstance(init_value, (int, float)):
|
||||
raise TypeError("Cached value is not numeric")
|
||||
value = init_value + value
|
||||
self.set_cache(key, value, **kwargs)
|
||||
return value
|
||||
|
|
|
|||
|
|
@ -86,6 +86,22 @@ def test_increment_cache_is_atomic_under_thread_concurrency():
|
|||
assert cache.get_cache("counter") == seed + thread_count
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_value", [{"count": 1}, [1], _CachedBudget(spend=1.0, events=[1])])
|
||||
def test_increment_cache_rejects_non_numeric_cached_values(bad_value):
|
||||
cache = InMemoryCache()
|
||||
cache.set_cache("counter", bad_value)
|
||||
with pytest.raises(TypeError, match="Cached value is not numeric"):
|
||||
cache.increment_cache("counter", 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_cache_sadd_rejects_non_set_cached_value():
|
||||
cache = InMemoryCache()
|
||||
cache.set_cache("members", ["one"])
|
||||
with pytest.raises(TypeError, match="Cached value is not a set"):
|
||||
await cache.async_set_cache_sadd("members", ["two"], ttl=None)
|
||||
|
||||
|
||||
async def test_async_increment_delegates_to_locked_sync_path():
|
||||
cache = InMemoryCache()
|
||||
assert await cache.async_increment("counter", 2) == 2
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue