mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
ccfbb77b73
* use mget_nonatomic * redis cluster override MGET op * fix redis cluster + MGET * test redis cluster
66 lines
2 KiB
Python
66 lines
2 KiB
Python
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
|