mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 11:14:04 +00:00
(Redis fix) - use mget_non_atomic (#8682)
* use mget_nonatomic * redis cluster override MGET op * fix redis cluster + MGET * test redis cluster
This commit is contained in:
parent
bb6f43d12e
commit
ccfbb77b73
4 changed files with 106 additions and 8 deletions
66
tests/litellm/caching/test_redis_cluster_cache.py
Normal file
66
tests/litellm/caching/test_redis_cluster_cache.py
Normal file
|
@ -0,0 +1,66 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
|
||||
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
def test_redis_cluster_batch_get(mock_init_redis_cluster):
|
||||
"""
|
||||
Test that RedisClusterCache uses mget_nonatomic instead of mget for batch operations
|
||||
"""
|
||||
# Create a mock Redis client
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.mget_nonatomic.return_value = [None, None] # Simulate no cache hits
|
||||
mock_init_redis_cluster.return_value = mock_redis
|
||||
|
||||
# Create RedisClusterCache instance with mock client
|
||||
cache = RedisClusterCache(
|
||||
startup_nodes=[{"host": "localhost", "port": 6379}],
|
||||
password="hello",
|
||||
)
|
||||
|
||||
# Test batch_get_cache
|
||||
keys = ["key1", "key2"]
|
||||
cache.batch_get_cache(keys)
|
||||
|
||||
# Verify mget_nonatomic was called instead of mget
|
||||
mock_redis.mget_nonatomic.assert_called_once()
|
||||
assert not mock_redis.mget.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
async def test_redis_cluster_async_batch_get(mock_init_redis_cluster):
|
||||
"""
|
||||
Test that RedisClusterCache uses mget_nonatomic instead of mget for async batch operations
|
||||
"""
|
||||
# Create a mock Redis client
|
||||
mock_redis = MagicMock()
|
||||
mock_redis.mget_nonatomic.return_value = [None, None] # Simulate no cache hits
|
||||
|
||||
# Create RedisClusterCache instance with mock client
|
||||
cache = RedisClusterCache(
|
||||
startup_nodes=[{"host": "localhost", "port": 6379}],
|
||||
password="hello",
|
||||
)
|
||||
|
||||
# Mock the init_async_client to return our mock redis client
|
||||
cache.init_async_client = MagicMock(return_value=mock_redis)
|
||||
|
||||
# Test async_batch_get_cache
|
||||
keys = ["key1", "key2"]
|
||||
await cache.async_batch_get_cache(keys)
|
||||
|
||||
# Verify mget_nonatomic was called instead of mget
|
||||
mock_redis.mget_nonatomic.assert_called_once()
|
||||
assert not mock_redis.mget.called
|
Loading…
Add table
Add a link
Reference in a new issue