From 4308a54cda0f2e5d735c9bfe4e7cffdb5253f6fc Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:53:56 +0000 Subject: [PATCH] refactor(caching): deprecate raw redis-py accessors on RedisCache --- litellm/caching/redis_batch.py | 2 +- litellm/caching/redis_cache.py | 217 ++++++++++++----- litellm/caching/redis_cluster_cache.py | 4 +- .../proxy/common_utils/config_sync_pubsub.py | 2 +- litellm/proxy/common_utils/debug_utils.py | 6 +- litellm/proxy/mcp_server/gateway_dcr_flow.py | 4 +- .../runtime_refresh_coordinator.py | 4 +- .../test_redis_increment_with_floor.py | 14 +- .../integration/sdk/test_dual_cache_redis.py | 36 +-- tests/unit/caching/test_caching.py | 4 +- tests/unit/caching/test_dual_cache.py | 2 +- tests/unit/caching/test_redis_batch.py | 2 +- tests/unit/caching/test_redis_cache.py | 219 +++++++++++++++--- .../unit/caching/test_redis_cluster_cache.py | 60 ++--- .../test_gateway_endpoints.py | 5 +- tests/unit/proxy/auth/test_auth_checks.py | 5 +- ...st_auth_checks_object_access_and_lookup.py | 10 +- .../test_auth_cache_invalidation_pubsub.py | 2 +- .../common_utils/test_config_sync_pubsub.py | 10 +- .../test_user_provider_credentials.py | 2 +- .../test_e2e_pod_lock_manager.py | 2 +- .../test_internal_user_endpoints.py | 5 +- .../test_key_management_endpoints.py | 5 +- .../test_team_endpoints.py | 7 +- .../test_team_roster_sync.py | 5 +- .../test_redis_refresh_coordinator.py | 32 +++ .../mcp_server/test_byok_credential_cache.py | 5 +- .../proxy/mcp_server/test_gateway_dcr_flow.py | 10 +- .../proxy/proxy_server/test_spend_counters.py | 2 +- .../test_response_polling_handler.py | 10 +- tests/unit/proxy/test_proxy_server.py | 6 +- ...test_proxy_server_endpoints_and_startup.py | 5 +- .../test_session_id_affinity.py | 2 +- .../test_router/test_io_token_rate_limits.py | 2 +- tests/unit/test_router_redis_init.py | 2 +- 35 files changed, 509 insertions(+), 201 deletions(-) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..53cb1ab5e35 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -373,7 +373,7 @@ class RedisBatch: widths: list[int] = [] # mutable-ok: filled while enqueuing async def run() -> list[object]: - client: Final = self.redis_cache.init_async_client() + client: Final = self.redis_cache._async_client() # pyright: ignore[reportPrivateUsage] # raw redis-py client, no typed alternative yet async with client.pipeline(transaction=False) as pipe: widths.extend(op.enqueue(pipe) for op in ops) return await pipe.execute(raise_on_error=False) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..69596d33304 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,6 +18,7 @@ import json import logging import threading import time +import warnings import weakref from collections import Counter from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence @@ -25,9 +26,10 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta from types import FrameType, MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypeVar, cast from pydantic import TypeAdapter +from typing_extensions import deprecated import litellm from litellm._internal_context import current_service_caller @@ -53,34 +55,52 @@ from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from redis.asyncio import Redis, RedisCluster + from redis import Redis as SyncRedis + from redis.asyncio import BlockingConnectionPool, Redis, RedisCluster from redis.asyncio.client import Pipeline from redis.asyncio.cluster import ClusterPipeline + from redis.asyncio.connection import Connection + from redis.cluster import RedisCluster as SyncRedisCluster from litellm._service_logger import ServiceLogging - pipeline = Pipeline - cluster_pipeline = ClusterPipeline - async_redis_client = Redis - async_redis_cluster_client = RedisCluster - Span = _Span + pipeline: TypeAlias = Pipeline[bytes] + cluster_pipeline: TypeAlias = ClusterPipeline[bytes] + redis_client: TypeAlias = SyncRedis[str] + redis_cluster_client: TypeAlias = SyncRedisCluster[str] + async_redis_client: TypeAlias = Redis[bytes] + async_redis_cluster_client: TypeAlias = RedisCluster[bytes] + redis_conn_pool: TypeAlias = BlockingConnectionPool[Connection] + Span: TypeAlias = _Span else: - pipeline = Any - cluster_pipeline = Any - async_redis_client = Any - async_redis_cluster_client = Any - Span = Any + pipeline: TypeAlias = Any + cluster_pipeline: TypeAlias = Any + redis_client: TypeAlias = Any + redis_cluster_client: TypeAlias = Any + async_redis_client: TypeAlias = Any + async_redis_cluster_client: TypeAlias = Any + redis_conn_pool: TypeAlias = Any + Span: TypeAlias = Any + + +_SyncRedisClient: TypeAlias = redis_client | redis_cluster_client + +_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE: Final = ( + "This accessor exposes the raw redis-py client and will be removed in a future release." +) class _AsyncRedisCommands(Protocol): """Async redis commands this cache issues. redis-py's type stubs omit these methods on RedisCluster, so the union returned by - init_async_client() is untyped at every call site without this protocol. + _async_client() is untyped at every call site without this protocol. """ def ping(self) -> Awaitable[bool]: ... + def get(self, name: str) -> Awaitable[bytes | None]: ... + def delete(self, *names: str) -> Awaitable[int]: ... def ttl(self, name: str) -> Awaitable[int]: ... @@ -637,7 +657,15 @@ def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) - class RedisCache(BaseCache): - # if users don't provider one, use the default litellm cache + def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks + super().__init_subclass__(**kwargs) + if "init_async_client" in vars(cls) or "init_pubsub_client" in vars(cls): + warnings.warn( + f"{cls.__qualname__} overrides a RedisCache raw client accessor. " + + _RAW_REDIS_CLIENT_DEPRECATION_MESSAGE, + DeprecationWarning, + stacklevel=2, + ) def __init__( self, @@ -672,10 +700,10 @@ class RedisCache(BaseCache): self.service_logger_obj = kwargs.pop("service_logger_obj") redis_kwargs.update(kwargs) - self.redis_client = get_redis_client(**redis_kwargs) - self.redis_async_client: async_redis_client | async_redis_cluster_client | None = None + self._redis_client: redis_client | redis_cluster_client = get_redis_client(**redis_kwargs) + self._redis_async_client: async_redis_client | async_redis_cluster_client | None = None self.redis_kwargs = redis_kwargs - self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs) + self._async_redis_conn_pool: redis_conn_pool | None = get_redis_connection_pool(**redis_kwargs) # redis namespaces self.namespace = namespace @@ -687,8 +715,8 @@ class RedisCache(BaseCache): self.redis_flush_size = redis_flush_size self.redis_version: str = "Unknown" try: - if not coroutine_checker.is_async_callable(self.redis_client): - self.redis_version = self.redis_client.info()["redis_version"] + if not coroutine_checker.is_async_callable(self._redis_client): + self.redis_version = self._redis_client.info()["redis_version"] except Exception: pass @@ -706,6 +734,36 @@ class RedisCache(BaseCache): else: super().__init__() # defaults to 60s + @property + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def redis_client(self) -> _SyncRedisClient: + return self._redis_client + + @redis_client.setter + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def redis_client(self, value: _SyncRedisClient) -> None: + self._redis_client = value + + @property + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def redis_async_client(self) -> async_redis_client | async_redis_cluster_client | None: + return self._redis_async_client + + @redis_async_client.setter + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def redis_async_client(self, value: async_redis_client | async_redis_cluster_client | None) -> None: + self._redis_async_client = value + + @property + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def async_redis_conn_pool(self) -> redis_conn_pool | None: + return self._async_redis_conn_pool + + @async_redis_conn_pool.setter + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def async_redis_conn_pool(self, value: redis_conn_pool | None) -> None: + self._async_redis_conn_pool = value + def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: return call() @@ -727,8 +785,8 @@ class RedisCache(BaseCache): # SYNC HEALTH PING try: - if hasattr(self.redis_client, "ping"): - self.redis_client.ping() + if hasattr(self._redis_client, "ping"): + self._redis_client.ping() except Exception as e: verbose_logger.error("Error connecting to Sync Redis client", extra={"error": str(e)}) self._handle_sync_ping_error(e) @@ -778,7 +836,7 @@ class RedisCache(BaseCache): kwargs_hash: Final = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16] return f"async-redis-client-{kwargs_hash}" - def init_async_client( + def _init_async_client( self, ) -> async_redis_client | async_redis_cluster_client: from redis.asyncio import RedisCluster @@ -790,26 +848,29 @@ class RedisCache(BaseCache): cache_key: Final = self._get_async_client_cache_key() cached_client: Final = in_memory_llm_clients_cache.get_cache(key=cache_key) if cached_client is not None: - redis_async_client = cast(async_redis_client | async_redis_cluster_client, cached_client) - else: - # Create new connection pool and client for current event loop - self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) - redis_async_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) - in_memory_llm_clients_cache.set_cache( - key=cache_key, - value=redis_async_client, - litellm_owned_client=isinstance(redis_async_client, RedisCluster), - ) + cached_redis_async_client: Final = cast(async_redis_client | async_redis_cluster_client, cached_client) + self._redis_async_client = cached_redis_async_client + return cached_redis_async_client - self.redis_async_client = redis_async_client - return redis_async_client + self._async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) + created_redis_async_client: Final = get_redis_async_client( + connection_pool=self._async_redis_conn_pool, **self.redis_kwargs + ) + in_memory_llm_clients_cache.set_cache( + key=cache_key, + value=created_redis_async_client, + litellm_owned_client=isinstance(created_redis_async_client, RedisCluster), + ) - def init_pubsub_client(self) -> async_redis_client: # pyright: ignore[reportUnknownParameterType] # redis generics + self._redis_async_client = created_redis_async_client + return created_redis_async_client + + def _init_pubsub_client(self) -> async_redis_client: # pyright: ignore[reportUnknownParameterType] # redis generics from redis.asyncio import RedisCluster from litellm import in_memory_llm_clients_cache - client: Final = self.init_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis generics + client: Final = self._async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis generics if not isinstance(client, RedisCluster): return client # pyright: ignore[reportUnknownVariableType] # redis generics cache_key: Final = f"{self._get_async_client_cache_key()}-pubsub" @@ -826,8 +887,34 @@ class RedisCache(BaseCache): in_memory_llm_clients_cache.set_cache(key=cache_key, value=pubsub_client, litellm_owned_client=True) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped in-memory client cache return pubsub_client # pyright: ignore[reportUnknownVariableType] # redis generics + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def init_async_client( + self, + ) -> async_redis_client | async_redis_cluster_client: + return self._init_async_client() + + @deprecated(_RAW_REDIS_CLIENT_DEPRECATION_MESSAGE) + def init_pubsub_client(self) -> async_redis_client: # pyright: ignore[reportUnknownParameterType] # redis generics + return self._init_pubsub_client() + + def _async_client(self) -> async_redis_client | async_redis_cluster_client: + if ( + "init_async_client" in vars(self) + or type(self).init_async_client is not _BASE_INIT_ASYNC_CLIENT # pyright: ignore[reportDeprecated] # override probe + ): + return self.init_async_client() # pyright: ignore[reportDeprecated] # honors a legacy subclass override + return self._init_async_client() + + def _pubsub_client(self) -> async_redis_client: # pyright: ignore[reportUnknownParameterType] # redis generics + if ( + "init_pubsub_client" in vars(self) + or type(self).init_pubsub_client is not _BASE_INIT_PUBSUB_CLIENT # pyright: ignore[reportDeprecated] # override probe + ): + return self.init_pubsub_client() # pyright: ignore[reportDeprecated, reportUnknownVariableType] # honors a legacy subclass override + return self._init_pubsub_client() # pyright: ignore[reportUnknownVariableType] # redis generics + def _async_commands(self) -> _AsyncRedisCommands: - return self.init_async_client() + return self._async_client() def check_and_fix_namespace(self, key: str) -> str: """ @@ -881,7 +968,7 @@ class RedisCache(BaseCache): key = self.check_and_fix_namespace(key=key) try: start_time: Final = time.time() - self.redis_client.set(name=key, value=str(value), ex=ttl) + self._redis_client.set(name=key, value=str(value), ex=ttl) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( @@ -897,7 +984,7 @@ class RedisCache(BaseCache): print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: - _redis_client: Final = self.redis_client + _redis_client: Final = self._redis_client start_time = time.time() set_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) @@ -968,7 +1055,7 @@ class RedisCache(BaseCache): takes traffic. Returns the resulting count. """ namespaced_key: Final = self.check_and_fix_namespace(key=key) - count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + count: Final[object] = self._redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl ) return _LUA_COUNT.validate_python(count) @@ -995,7 +1082,7 @@ class RedisCache(BaseCache): start_time: Final = time.time() try: keys: Final = [] - _redis_client: Final = self.init_async_client() + _redis_client: Final = self._async_client() if not hasattr(_redis_client, "scan_iter"): verbose_logger.debug( "Redis client does not support scan_iter, potentially using Redis Cluster. Returning empty list." @@ -1055,7 +1142,7 @@ class RedisCache(BaseCache): registered script is bound to the connection of the loop it was created on; awaiting it from another loop raises "got Future attached to a different loop". Binding lazily on the calling loop gives the script the - same per-loop scoping init_async_client already gives the clients, so a + same per-loop scoping _init_async_client already gives the clients, so a script registered once at startup is never reused across loops. Args: @@ -1102,7 +1189,7 @@ class RedisCache(BaseCache): Kept separate from async_register_script so each loop caches its own executor; see that method for why the binding must be per loop. """ - _redis_client: Final[Any] = self.init_async_client() + _redis_client: Final[Any] = self._async_client() if hasattr(_redis_client, "register_script"): registered_script: Final = _redis_client.register_script(script) @@ -1144,7 +1231,7 @@ class RedisCache(BaseCache): start_time: Final = time.time() try: - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final[Redis] = self._async_client() except Exception as e: end_time = time.time() _duration = end_time - start_time @@ -1253,7 +1340,7 @@ class RedisCache(BaseCache): if len(cache_list) == 0: return - _redis_client: Final = self.init_async_client() + _redis_client: Final = self._async_client() start_time: Final = time.time() print_verbose( @@ -1319,7 +1406,7 @@ class RedisCache(BaseCache): ) start_time: Final = time.time() try: - async with self.init_async_client().pipeline(transaction=False) as pipe: + async with self._async_client().pipeline(transaction=False) as pipe: for cache_key, json_cache_value, ttl in commands: pipe.set(name=cache_key, value=json_cache_value, ex=None if ttl is None else timedelta(seconds=ttl)) await pipe.execute() @@ -1376,7 +1463,7 @@ class RedisCache(BaseCache): start_time: Final = time.time() try: - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final[Redis] = self._async_client() except Exception as e: end_time = time.time() _duration = end_time - start_time @@ -1479,7 +1566,7 @@ class RedisCache(BaseCache): ) -> float: from redis.asyncio import Redis - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final[Redis] = self._async_client() start_time: Final = time.time() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) @@ -1544,7 +1631,7 @@ class RedisCache(BaseCache): GET/compare/SET runs in a single Lua call, so it is also atomic across racing callers and pods. Returns the resulting value. """ - _redis_client: Final = self.init_async_client() + _redis_client: Final = self._async_client() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) lua: Final = ( @@ -1636,7 +1723,7 @@ class RedisCache(BaseCache): key = self.check_and_fix_namespace(key=key) print_verbose(f"Get Redis Cache: key: {key}") start_time: Final = time.time() - cached_response: Final = self.redis_client.get(key) + cached_response: Final = self._redis_client.get(key) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( @@ -1662,7 +1749,7 @@ class RedisCache(BaseCache): We use a wrapper so RedisCluster can override this method """ - return self.redis_client.mget(keys=keys) + return self._redis_client.mget(keys=keys) async def _async_run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ @@ -1670,7 +1757,7 @@ class RedisCache(BaseCache): We use a wrapper so RedisCluster can override this method """ - async_redis_client: Final = self.init_async_client() + async_redis_client: Final = self._async_client() return await async_redis_client.mget(keys=keys) def batch_get_cache( @@ -1743,7 +1830,7 @@ class RedisCache(BaseCache): ): from redis.asyncio import Redis - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final[Redis] = self._async_client() key = self.check_and_fix_namespace(key=key) start_time: Final = time.time() @@ -1869,7 +1956,7 @@ class RedisCache(BaseCache): print_verbose("Pinging Sync Redis Cache") start_time: Final = time.time() try: - response: Final[bool] = self.redis_client.ping() + response: Final[bool] = self._redis_client.ping() print_verbose(f"Redis Cache PING: {response}") ## LOGGING ## end_time = time.time() @@ -1941,24 +2028,24 @@ class RedisCache(BaseCache): await _redis_client.delete(*keys) def client_list(self) -> list: - client_list: Final[list] = self.redis_client.client_list() + client_list: Final[list] = self._redis_client.client_list() return client_list def info(self): - info: Final = self.redis_client.info() + info: Final = self._redis_client.info() return info def flush_cache(self): - self.redis_client.flushall() + self._redis_client.flushall() def flushall(self): - self.redis_client.flushall() + self._redis_client.flushall() async def disconnect(self): from litellm import in_memory_llm_clients_cache - if self.async_redis_conn_pool is not None: - await self.async_redis_conn_pool.disconnect(inuse_connections=True) + if self._async_redis_conn_pool is not None: + await self._async_redis_conn_pool.disconnect(inuse_connections=True) cached_pubsub_client: Final = cast( # cast-ok: only this module stores clients under this key # pyright: ignore[reportUnknownVariableType] # redis generics async_redis_client | None, in_memory_llm_clients_cache.get_cache( # pyright: ignore[reportUnknownMemberType] # untyped in-memory client cache @@ -1971,7 +2058,7 @@ class RedisCache(BaseCache): except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection verbose_logger.debug("Error closing cached pub/sub Redis client: %s", e) try: - self.redis_client.close() + self._redis_client.close() except Exception as e: verbose_logger.debug("Error closing sync Redis client: %s", e) @@ -2021,7 +2108,7 @@ class RedisCache(BaseCache): def delete_cache(self, key): key = self.check_and_fix_namespace(key=key) - self.redis_client.delete(key) + self._redis_client.delete(key) async def _pipeline_increment_helper( self, @@ -2063,7 +2150,7 @@ class RedisCache(BaseCache): from redis.asyncio import Redis - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final[Redis] = self._async_client() start_time: Final = time.time() print_verbose("Increment Async Redis Cache Pipeline: increment list: %s", increment_list) @@ -2513,3 +2600,7 @@ class RedisCache(BaseCache): e, ) raise e + + +_BASE_INIT_ASYNC_CLIENT: Final = RedisCache.init_async_client # pyright: ignore[reportDeprecated] # identity for the override probe +_BASE_INIT_PUBSUB_CLIENT: Final = RedisCache.init_pubsub_client # pyright: ignore[reportDeprecated] # identity for the override probe diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 6fb84df331e..c8001766d22 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -130,13 +130,13 @@ class RedisClusterCache(RedisCache): """ Overrides `_run_redis_mget_operation` in redis_cache.py """ - return self.redis_client.mget_nonatomic(keys=keys) + return self._redis_client.mget_nonatomic(keys=keys) async def _async_run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ Overrides `_async_run_redis_mget_operation` in redis_cache.py """ - async_redis_cluster_client: Final = self.init_async_client() + async_redis_cluster_client: Final = self._async_client() return await async_redis_cluster_client.mget_nonatomic(keys=keys) async def test_connection(self) -> dict: diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index e9aaa0e81e4..d7ff8d789ec 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -99,7 +99,7 @@ def config_sync_channel(redis_cache: "RedisCache") -> str: def pubsub_capable_client(redis_cache: "RedisCache") -> _ConfigSyncPubSubClient: return cast( # cast-ok: protocol view of the pub/sub-capable async redis client _ConfigSyncPubSubClient, - redis_cache.init_pubsub_client(), # pyright: ignore[reportUnknownMemberType] # redis generics + redis_cache._pubsub_client(), # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # raw redis-py client, no typed alternative yet ) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 884aacb9fe1..813fcbff4e7 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -570,9 +570,9 @@ def _get_cache_memory_stats( } # Try to get Redis connection pool info if available try: - if hasattr(redis_usage_cache, "redis_client") and redis_usage_cache.redis_client: - if hasattr(redis_usage_cache.redis_client, "connection_pool"): - pool_info: Final = redis_usage_cache.redis_client.connection_pool + if hasattr(redis_usage_cache, "_redis_client") and redis_usage_cache._redis_client: + if hasattr(redis_usage_cache._redis_client, "connection_pool"): + pool_info: Final = redis_usage_cache._redis_client.connection_pool cache_stats["redis_usage_cache"]["connection_pool"] = { "max_connections": ( pool_info.max_connections if hasattr(pool_info, "max_connections") else None diff --git a/litellm/proxy/mcp_server/gateway_dcr_flow.py b/litellm/proxy/mcp_server/gateway_dcr_flow.py index ae74d878b36..18b8250c4b7 100644 --- a/litellm/proxy/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/mcp_server/gateway_dcr_flow.py @@ -1403,7 +1403,9 @@ class _SingleUseGuard: redis_cache: Final = redis_usage_cache or getattr(self._cache, "redis_cache", None) if redis_cache is not None: try: - value = await redis_cache.init_async_client().get(redis_cache.check_and_fix_namespace(key=key)) + value: Final = await redis_cache._async_commands().get( # pyright: ignore[reportPrivateUsage] # internal raw-client access + redis_cache.check_and_fix_namespace(key=key) + ) except Exception as e: # noqa: BLE001 # ANY Redis fault fails the read closed verbose_logger.warning("mcp gateway single-use peek: shared cache backend unavailable: %s", e) return "unavailable" diff --git a/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py index 9376d1bda50..a7d061d592a 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py +++ b/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -31,9 +31,9 @@ def runtime_refresh_coordinator() -> RefreshCoordinator | None: redis_cache: Final = user_api_key_cache.redis_cache if redis_cache is None: return None - # The Redis client from init_async_client() is only partially typed; the lock validates every + # The Redis client from _async_client() is only partially typed; the lock validates every # reply it depends on, so the untyped boundary is contained here. - redis_client: Final = redis_cache.init_async_client() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm redis wrapper is untyped + redis_client: Final = redis_cache._async_client() # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType] # raw redis-py client, no typed alternative yet lock: Final = RedisDistributedLock( redis_client, # pyright: ignore[reportArgumentType,reportUnknownArgumentType] # litellm redis wrapper is untyped namespace_key=redis_cache.check_and_fix_namespace, diff --git a/tests/integration/routing/test_redis_increment_with_floor.py b/tests/integration/routing/test_redis_increment_with_floor.py index d5535d30715..f0b36b27181 100644 --- a/tests/integration/routing/test_redis_increment_with_floor.py +++ b/tests/integration/routing/test_redis_increment_with_floor.py @@ -38,22 +38,22 @@ def test_traffic_never_pushes_a_counters_expiry_back_out(counter: tuple[RedisCac cache, key, namespaced_key = counter cache.increment_with_floor(key, 1, TTL) - assert cache.redis_client.ttl(namespaced_key) > TTL - 60 + assert cache._redis_client.ttl(namespaced_key) > TTL - 60 - cache.redis_client.expire(namespaced_key, 30) + cache._redis_client.expire(namespaced_key, 30) cache.increment_with_floor(key, 1, TTL) - assert cache.redis_client.ttl(namespaced_key) <= 30 + assert cache._redis_client.ttl(namespaced_key) <= 30 def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter: tuple[RedisCache, str, str]) -> None: cache, key, namespaced_key = counter cache.increment_with_floor(key, 1, TTL) - cache.redis_client.expire(namespaced_key, 30) + cache._redis_client.expire(namespaced_key, 30) assert cache.increment_with_floor(key, -5, TTL) == 0 - assert cache.redis_client.ttl(namespaced_key) <= 30 + assert cache._redis_client.ttl(namespaced_key) <= 30 @pytest.mark.asyncio @@ -63,7 +63,7 @@ async def test_the_async_counter_behaves_the_same_way(counter: tuple[RedisCache, assert await cache.async_increment_with_floor(key, 2, TTL) == 2 assert await cache.async_batch_get_counts([key]) == (2,) - cache.redis_client.expire(namespaced_key, 30) + cache._redis_client.expire(namespaced_key, 30) assert await cache.async_increment_with_floor(key, -9, TTL) == 0 - assert cache.redis_client.ttl(namespaced_key) <= 30 + assert cache._redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index 7d05eb2795b..4f55f2c2640 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/tests/integration/sdk/test_dual_cache_redis.py @@ -134,7 +134,7 @@ async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unav key: Final = f"session-repin-legacy-{uuid.uuid4().hex}" replacement: Final = {"model_id": "replacement"} try: - redis.redis_client.set(redis.check_and_fix_namespace(key), stored, ex=60) + redis._redis_client.set(redis.check_and_fix_namespace(key), stored, ex=60) preserved: Final = await claim_affinity_pin( cache, key, @@ -169,12 +169,12 @@ async def test_tier_pin_reads_shared_state_refreshes_ttl_and_forgets_deleted_pin try: for worker in workers: worker.in_memory_cache.set_cache(key, {"model": "old", "tier": "SIMPLE"}, ttl=60) - redis.redis_client.set(namespaced, stored, ex=10) + redis._redis_client.set(namespaced, stored, ex=10) first: Final = await read_affinity_pin(workers[0], key, 60) second: Final = await read_affinity_pin(workers[1], key, 60) assert first == second and first in ("current", {"model": "current", "tier": "REASONING"}) assert all(worker.in_memory_cache.get_cache(key) == first for worker in workers) - assert 10 < redis.redis_client.ttl(namespaced) <= 60 + assert 10 < redis._redis_client.ttl(namespaced) <= 60 replacement: Final = {"model": "newer", "tier": "SIMPLE"} await redis.async_set_cache(key, replacement, ttl=10) assert await read_affinity_pin(workers[0], key, 60) == replacement @@ -183,7 +183,7 @@ async def test_tier_pin_reads_shared_state_refreshes_ttl_and_forgets_deleted_pin for worker in workers: assert await read_affinity_pin(worker, key, 60) is None assert worker.in_memory_cache.get_cache(key) is None - assert redis.redis_client.exists(namespaced) == 0 + assert redis._redis_client.exists(namespaced) == 0 finally: await redis.async_delete_cache(key) @@ -230,14 +230,14 @@ async def test_tier_pin_joins_pending_reads_without_reusing_request_state(mode: current: Final = {"model": "current", "tier": "REASONING"} newer: Final = {"model": "newer", "tier": "COMPLEX"} try: - redis.redis_client.set(key, json.dumps(current), ex=10) - redis.redis_client.set(other_key, json.dumps("other-value"), ex=60) + redis._redis_client.set(key, json.dumps(current), ex=10) + redis._redis_client.set(other_key, json.dumps("other-value"), ex=60) assert await read_affinity_pin(cache, "pin", 60) == current cache.in_memory_cache.set_cache("pin", stale, ttl=60) assert cache.in_memory_cache.get_cache("pin") == stale if mode == "script_error": - redis.redis_client.delete(key) - redis.redis_client.rpush(key, "wrong-type") + redis._redis_client.delete(key) + redis._redis_client.rpush(key, "wrong-type") start: Final = len(sends) with request_redis_batch_scope() as request: batch: Final = request.batch(redis) @@ -250,27 +250,27 @@ async def test_tier_pin_joins_pending_reads_without_reusing_request_state(mode: assert len(wire) == (2 if mode == "separate" else 1) assert sum(b"MGET" in data for data in wire) == sum(b"EVALSHA" in data for data in wire) == 1 if mode == "script_error": - assert redis.redis_client.type(key) == b"list" - redis.redis_client.delete(key) + assert redis._redis_client.type(key) == b"list" + redis._redis_client.delete(key) else: - assert 10 < redis.redis_client.ttl(key) <= 60 - redis.redis_client.set(key, json.dumps(newer), ex=10) + assert 10 < redis._redis_client.ttl(key) <= 60 + redis._redis_client.set(key, json.dumps(newer), ex=10) later_pending: Final = batch.mget(("other",)) assert await read_affinity_pin(cache, "pin", 60) == newer assert await later_pending == {"other": "other-value"} assert cache.in_memory_cache.get_cache("pin") == newer - assert 10 < redis.redis_client.ttl(key) <= 60 - assert redis.redis_client.get(key) == json.dumps(newer).encode() - redis.redis_client.delete(key) + assert 10 < redis._redis_client.ttl(key) <= 60 + assert redis._redis_client.get(key) == json.dumps(newer).encode() + redis._redis_client.delete(key) assert await read_affinity_pin(cache, "pin", 60) is None assert cache.in_memory_cache.get_cache("pin") is None - assert redis.redis_client.exists(key) == 0 + assert redis._redis_client.exists(key) == 0 finally: - redis.redis_client.delete(key, other_key) + redis._redis_client.delete(key, other_key) for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict): if client_key in cached_key: litellm.in_memory_llm_clients_cache.delete_cache(cached_key) await client.aclose() await pool.aclose() await redis.disconnect() - redis.redis_client.close() + redis._redis_client.close() diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index c4727cbeff0..78c2b0f375d 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -1236,7 +1236,7 @@ def test_caching_with_reasoning_content(monkeypatch: pytest.MonkeyPatch) -> None def test_custom_redis_cache_params(redis_client_without_network: None) -> None: cache: Final = Cache(type=LiteLLMCacheType.REDIS, host="127.0.0.1", port="6379", db=4) - assert cache.cache.redis_client.connection_pool.connection_kwargs["db"] == 4 + assert cache.cache._redis_client.connection_pool.connection_kwargs["db"] == 4 def test_redis_cache_ignores_a_leftover_disk_cache_dir( @@ -1530,7 +1530,7 @@ def test_redis_caching_llm_caching_ttl(monkeypatch: pytest.MonkeyPatch, redis_cl monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") set_cache: Final = MagicMock() - monkeypatch.setattr(cache.redis_client, "set", set_cache) + monkeypatch.setattr(cache._redis_client, "set", set_cache) cache.set_cache(key="cache-key", value="cache-value") diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index 5ebc3bf7336..3601cfa9dfd 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -74,7 +74,7 @@ def _assert_sync_batch_used_blocking_client(dual_cache: DualCache, mock_redis: M assert result == ["redis_value"] mock_redis.batch_get_cache.assert_called_once_with(key_list=["lit6729_key"], parent_otel_span=None) mock_redis.async_batch_get_cache.assert_not_called() - mock_redis.init_async_client.assert_not_called() + mock_redis._init_async_client.assert_not_called() assert dual_cache.in_memory_cache.get_cache("lit6729_key") == "redis_value" diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..c5702bc72b9 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -99,7 +99,7 @@ class FakeRedisCache(RedisCache): self.alone: list[tuple[str, Any]] = [] self.store: dict[str, Any] = {} - def init_async_client(self) -> FakeClient: # pyright: ignore[reportIncompatibleMethodOverride] # fake client, no server + def _init_async_client(self) -> FakeClient: # pyright: ignore[reportIncompatibleMethodOverride] # fake client, no server return self.client async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, Any]: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct read diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..f5134c736dd 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -2,7 +2,9 @@ import asyncio import gc import time import types +import warnings from collections.abc import Iterator +from dataclasses import dataclass from datetime import timedelta from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -29,11 +31,162 @@ def sync_batch_redis_cache(redis_no_ping): "litellm._redis.get_redis_client", return_value=MagicMock() ) as get_client: cache = RedisCache(host="127.0.0.1", port=6379) - cache.redis_client.mget.side_effect = OSError("redis unavailable") + cache._redis_client.mget.side_effect = OSError("redis unavailable") get_client.assert_called_once() yield cache +@dataclass(frozen=True, slots=True) +class _FactoryBackends: + cache: RedisCache + sync_client: MagicMock + async_client: AsyncMock + pool: MagicMock + + +@pytest.fixture +def factory_backends(redis_no_ping) -> Iterator[_FactoryBackends]: + sync_client: Final = MagicMock() + async_client: Final = AsyncMock() + pool: Final = MagicMock() + litellm.in_memory_llm_clients_cache.flush_cache() + with patch( # test-quality-ok: RedisCache builds its clients from these factories, with no injection point + "litellm._redis.get_redis_client", return_value=sync_client + ), patch("litellm._redis.get_redis_connection_pool", return_value=pool), patch( + "litellm._redis.get_redis_async_client", return_value=async_client + ): + yield _FactoryBackends(RedisCache(host="127.0.0.1", port=6379), sync_client, async_client, pool) + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.parametrize( + ("read", "backend"), + [ + pytest.param(lambda cache: cache.redis_client, "sync_client", id="redis_client"), + pytest.param(lambda cache: cache.async_redis_conn_pool, "pool", id="async_redis_conn_pool"), + pytest.param(lambda cache: cache.init_async_client(), "async_client", id="init_async_client"), + pytest.param(lambda cache: cache.init_pubsub_client(), "async_client", id="init_pubsub_client"), + ], +) +def test_deprecated_accessor_warns_once_and_returns_the_backend_in_use(factory_backends, read, backend: str) -> None: + with pytest.warns(DeprecationWarning, match="raw redis-py client") as caught: + returned: Final = read(factory_backends.cache) + + assert returned is getattr(factory_backends, backend) + assert len(caught) == 1, [str(w.message) for w in caught] + + +@pytest.mark.asyncio +async def test_redis_async_client_reports_the_client_serving_async_reads(factory_backends) -> None: + factory_backends.async_client.get.return_value = '"async-value"' + + assert await factory_backends.cache.async_get_cache("async-key") == "async-value" + with pytest.warns(DeprecationWarning, match="raw redis-py client"): + assert factory_backends.cache.redis_async_client is factory_backends.async_client + + +@pytest.mark.asyncio +async def test_no_deprecation_warning_for_internal_redis_operations(factory_backends) -> None: + from litellm.proxy.common_utils.config_sync_pubsub import pubsub_capable_client + + cache: Final = factory_backends.cache + sync_client: Final = factory_backends.sync_client + async_client: Final = factory_backends.async_client + sync_client.reset_mock() + sync_client.get.return_value = '"sync-value"' + sync_client.ping.return_value = True + async_client.get.return_value = '"async-value"' + async_client.set.return_value = True + async_client.ping.return_value = True + + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + cache.set_cache("sync-key", "sync-value") + assert cache.get_cache("sync-key") == "sync-value" + assert cache.sync_ping() is True + assert await cache.async_set_cache("async-key", "async-value") is True + assert await cache.async_get_cache("async-key") == "async-value" + assert await cache.ping() is True + assert pubsub_capable_client(cache) is async_client + + sync_client.set.assert_called_once_with(name="sync-key", value="sync-value", ex=60) + sync_client.get.assert_called_once_with("sync-key") + sync_client.ping.assert_called_once_with() + async_client.set.assert_awaited_once_with(name="async-key", value='"async-value"', nx=False, ex=60) + async_client.get.assert_awaited_once_with("async-key") + async_client.ping.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_legacy_init_async_client_subclass_override_serves_cache_reads(factory_backends) -> None: + override_client: Final = AsyncMock() + override_client.get.return_value = '"from-override"' + + with pytest.warns(DeprecationWarning, match="overrides a RedisCache raw client accessor"): + + class LegacyAsyncCache(RedisCache): + def init_async_client(self) -> AsyncMock: + return override_client + + cache: Final = LegacyAsyncCache(host="127.0.0.1", port=6379) + + assert await cache.async_get_cache("legacy-key") == "from-override" + override_client.get.assert_awaited_once_with("legacy-key") + factory_backends.async_client.get.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("patched_on", ["class", "instance"]) +async def test_patched_init_async_client_serves_cache_reads(factory_backends, patched_on: str) -> None: + override_client: Final = AsyncMock() + override_client.get.return_value = '"from-patch"' + target: Final = RedisCache if patched_on == "class" else factory_backends.cache + + with patch.object(target, "init_async_client", return_value=override_client): + assert await factory_backends.cache.async_get_cache("patched-key") == "from-patch" + + override_client.get.assert_awaited_once_with("patched-key") + factory_backends.async_client.get.assert_not_awaited() + + +def test_legacy_init_pubsub_client_override_still_serves_config_sync(factory_backends) -> None: + from litellm.proxy.common_utils.config_sync_pubsub import pubsub_capable_client + + override_client: Final = object() + + with pytest.warns(DeprecationWarning, match="overrides a RedisCache raw client accessor"): + + class LegacyPubSubCache(RedisCache): + def init_pubsub_client(self) -> object: + return override_client + + cache: Final = LegacyPubSubCache(host="127.0.0.1", port=6379) + + assert pubsub_capable_client(cache) is override_client + + +@pytest.mark.parametrize("name", ["redis_client", "redis_async_client", "async_redis_conn_pool"]) +def test_deprecated_attribute_assignment_warns_and_reads_back(factory_backends, name: str) -> None: + replacement: Final = object() + + with pytest.warns(DeprecationWarning, match="raw redis-py client"): + setattr(factory_backends.cache, name, replacement) + with pytest.warns(DeprecationWarning, match="raw redis-py client"): + assert getattr(factory_backends.cache, name) is replacement + + +def test_assigned_sync_redis_client_serves_sync_reads(factory_backends) -> None: + replacement: Final = MagicMock() + replacement.get.return_value = '"from-assigned-client"' + + with pytest.warns(DeprecationWarning, match="raw redis-py client"): + factory_backends.cache.redis_client = replacement + + assert factory_backends.cache.get_cache("assigned-key") == "from-assigned-client" + replacement.get.assert_called_once_with("assigned-key") + factory_backends.sync_client.get.assert_not_called() + + @pytest.mark.parametrize( ("namespace", "key", "expected"), [ @@ -69,7 +222,7 @@ async def test_async_delete_cache_applies_namespace(namespace, monkeypatch, redi redis_cache = RedisCache(namespace=namespace) mock_redis_instance = AsyncMock() - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): await redis_cache.async_delete_cache(key="3997c4abcdef") expected_key = "litellm:3997c4abcdef" if namespace else "3997c4abcdef" @@ -83,7 +236,7 @@ def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace=namespace) mock_redis_client = MagicMock() - redis_cache.redis_client = mock_redis_client + redis_cache._redis_client = mock_redis_client redis_cache.delete_cache(key="3997c4abcdef") @@ -111,7 +264,7 @@ async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping, monkeypatch.delenv("REDIS_HOST", raising=False) redis_cache = RedisCache(socket_timeout=1.0, **redis_config) assert redis_cache.redis_kwargs["socket_timeout"] == 1.0 - client = redis_cache.init_async_client() + client = redis_cache._init_async_client() assert client is not None assert client.connection_pool.connection_kwargs["socket_timeout"] == 1.0 @@ -148,7 +301,7 @@ async def test_async_rpush_pipeline_empty_list_returns_empty(monkeypatch, redis_ mock_redis_instance = AsyncMock() - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): result = await redis_cache.async_rpush_pipeline(rpush_list=[]) assert result == [] @@ -163,7 +316,7 @@ async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): mock_redis_instance = AsyncMock() - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): result = await redis_cache.async_lpop_pipeline(lpop_list=[]) assert result == [] @@ -199,7 +352,7 @@ async def test_async_register_script_namespaces_keys(namespace, raw_keys, expect mock_redis_instance = MagicMock() mock_redis_instance.register_script = MagicMock(return_value=registered_script) - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): script = redis_cache.async_register_script("return 1") result = await script(keys=raw_keys, args=[60]) @@ -231,7 +384,7 @@ def test_async_register_script_binds_per_event_loop(namespace, monkeypatch): unique_script = "return 'lit3298'" - with patch.object(redis_cache, "init_async_client", side_effect=make_client): + with patch.object(redis_cache, "_init_async_client", side_effect=make_client): script = redis_cache.async_register_script(unique_script) # Registration is deferred: no client is touched until the script runs. @@ -273,8 +426,8 @@ async def test_async_register_script_not_shared_across_namespaces(monkeypatch, r same_script = "return redis.call('GET', KEYS[1])" with ( - patch.object(cache_a, "init_async_client", return_value=client_a), - patch.object(cache_b, "init_async_client", return_value=client_b), + patch.object(cache_a, "_init_async_client", return_value=client_a), + patch.object(cache_b, "_init_async_client", return_value=client_b), ): script_a = cache_a.async_register_script(same_script) script_b = cache_b.async_register_script(same_script) @@ -297,7 +450,7 @@ async def test_async_register_script_cluster_path_uses_evalsha(monkeypatch, redi cluster_client.script_load = MagicMock(return_value="sha123") cluster_client.evalsha = AsyncMock(return_value="cluster-ok") - with patch.object(redis_cache, "init_async_client", return_value=cluster_client): + with patch.object(redis_cache, "_init_async_client", return_value=cluster_client): script = redis_cache.async_register_script("return 'cluster'") result = await script(keys=["{k:v}:tokens"], args=[5, 60]) @@ -314,7 +467,7 @@ async def test_async_register_script_raises_for_unsupported_client(monkeypatch, redis_cache = RedisCache() bad_client = MagicMock(spec=[]) - with patch.object(redis_cache, "init_async_client", return_value=bad_client): + with patch.object(redis_cache, "_init_async_client", return_value=bad_client): script = redis_cache.async_register_script("return 'x'") with pytest.raises(ValueError, match="does not support Lua script"): await script(keys=["k"], args=[1]) @@ -326,7 +479,7 @@ async def test_async_delete_cache_namespaces_key(namespace, expected, monkeypatc monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace=namespace) mock_redis_instance = AsyncMock() - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): await redis_cache.async_delete_cache("k") mock_redis_instance.delete.assert_awaited_once_with(expected) @@ -337,7 +490,7 @@ async def test_delete_cache_keys_namespaces_keys(namespace, expected, monkeypatc monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace=namespace) mock_redis_instance = AsyncMock() - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): await redis_cache.delete_cache_keys(["k"]) mock_redis_instance.delete.assert_awaited_once_with(expected) @@ -349,7 +502,7 @@ async def test_async_get_ttl_namespaces_key(namespace, expected, monkeypatch, re redis_cache = RedisCache(namespace=namespace) mock_redis_instance = AsyncMock() mock_redis_instance.ttl = AsyncMock(return_value=42) - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): ttl = await redis_cache.async_get_ttl("k") assert ttl == 42 mock_redis_instance.ttl.assert_awaited_once_with(expected) @@ -362,7 +515,7 @@ async def test_async_lpop_namespaces_key(namespace, expected, monkeypatch, redis redis_cache = RedisCache(namespace=namespace) mock_redis_instance = AsyncMock() mock_redis_instance.lpop = AsyncMock(return_value=b"value") - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): await redis_cache.async_lpop(key="k") mock_redis_instance.lpop.assert_awaited_once_with(expected, None) @@ -374,7 +527,7 @@ async def test_async_rpush_namespaces_key(namespace, expected, monkeypatch, redi redis_cache = RedisCache(namespace=namespace) mock_redis_instance = AsyncMock() mock_redis_instance.rpush = AsyncMock(return_value=1) - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): await redis_cache.async_rpush("k", ["v"]) mock_redis_instance.rpush.assert_awaited_once_with(expected, "v") @@ -398,7 +551,7 @@ async def test_async_scan_iter_namespaces_pattern(namespace, expected_match, mon mock_redis_instance = MagicMock() mock_redis_instance.scan_iter = scan_iter - with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with patch.object(redis_cache, "_init_async_client", return_value=mock_redis_instance): await redis_cache.async_scan_iter(pattern="k") assert captured["match"] == expected_match @@ -410,7 +563,7 @@ def test_increment_cache_namespaces_key(namespace, expected, monkeypatch, redis_ mock_client = MagicMock() mock_client.incr.return_value = 5 mock_client.ttl.return_value = 100 - redis_cache.redis_client = mock_client + redis_cache._redis_client = mock_client redis_cache.increment_cache(key="k", value=1) mock_client.incr.assert_called_once_with(name=expected, amount=1) @@ -420,7 +573,7 @@ def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace=namespace) mock_client = MagicMock() - redis_cache.redis_client = mock_client + redis_cache._redis_client = mock_client redis_cache.delete_cache(key="k") mock_client.delete.assert_called_once_with(expected) @@ -480,7 +633,7 @@ def test_circuit_breaker_open_makes_sync_batch_get_cache_fast_fail(sync_batch_re with caplog.at_level("INFO"): with pytest.raises(RedisCircuitBreakerOpenError): sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) - sync_batch_redis_cache.redis_client.mget.assert_called() + sync_batch_redis_cache._redis_client.mget.assert_called() assert caplog.records == [] @@ -493,7 +646,7 @@ def test_sync_get_cache_failure_feeds_the_breaker_and_logs_a_well_formed_record( """ from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - sync_batch_redis_cache.redis_client.get.side_effect = OSError("redis unavailable") + sync_batch_redis_cache._redis_client.get.side_effect = OSError("redis unavailable") with caplog.at_level("ERROR"): for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): @@ -529,7 +682,7 @@ async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports ): cache = RedisCache(host="127.0.0.1", port=6379) - with patch.object(cache, "init_async_client", return_value=failing_client): + with patch.object(cache, "_init_async_client", return_value=failing_client): assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} with pytest.raises(OSError, match="redis unavailable"): @@ -545,7 +698,7 @@ def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(sto "litellm._redis.get_redis_client", return_value=MagicMock() ): cache = RedisCache(host="127.0.0.1", port=6379) - cache.redis_client.mget.return_value = [stored, None, b"0"] + cache._redis_client.mget.return_value = [stored, None, b"0"] assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) @@ -726,7 +879,7 @@ async def test_circuit_breaker_success_still_resets_the_failure_streak(): reachable_redis = AsyncMock() reachable_redis.get.return_value = None - with patch.object(cache, "init_async_client", return_value=reachable_redis): + with patch.object(cache, "_init_async_client", return_value=reachable_redis): await cache.async_get_cache("lit4930") for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - 1): @@ -1204,8 +1357,8 @@ def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_bat monkeypatch.setattr( redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) ) - sync_batch_redis_cache.redis_client.get.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") - sync_batch_redis_cache.redis_client.mget.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + sync_batch_redis_cache._redis_client.get.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + sync_batch_redis_cache._redis_client.mget.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") with caplog.at_level(logging.DEBUG, logger="LiteLLM"): for _ in range(200): @@ -1234,7 +1387,7 @@ def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_bat ] caplog.clear() - sync_batch_redis_cache.redis_client.get.side_effect = OSError("redis unavailable") + sync_batch_redis_cache._redis_client.get.side_effect = OSError("redis unavailable") with caplog.at_level(logging.DEBUG, logger="LiteLLM"): for _ in range(3): assert sync_batch_redis_cache.get_cache("lit7520") is None @@ -1290,7 +1443,7 @@ async def test_write_path_timeouts_inside_the_interval_stay_at_debug(call_method cache = RedisCache() with ( - patch.object(cache, "init_async_client", return_value=client), + patch.object(cache, "_init_async_client", return_value=client), caplog.at_level(logging.DEBUG, logger="LiteLLM"), ): with contextlib.suppress(RedisTimeoutError): @@ -1464,7 +1617,7 @@ async def test_async_increment_pipelines_the_ttl_command( redis_cache = RedisCache(namespace="ns") client = _RoundTripCountingRedis(ttl=existing_ttl) - with patch.object(redis_cache, "init_async_client", return_value=client): + with patch.object(redis_cache, "_init_async_client", return_value=client): first = await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=refresh_ttl) second = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=refresh_ttl) @@ -1502,7 +1655,7 @@ async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypat client = MagicMock() client.pipeline = MagicMock(return_value=pipe) - with patch.object(redis_cache, "init_async_client", return_value=client): + with patch.object(redis_cache, "_init_async_client", return_value=client): await redis_cache.async_set_cache_pipeline_with_ttls( (("team_id:t1", {"team_id": "t1"}, 60), ("u1", {"user_id": "u1"}, 7), ("org_id:o1", {"a": 1}, None)) ) @@ -1555,7 +1708,7 @@ async def test_async_rpush_and_trim_runs_push_and_trim_in_one_transaction(monkey client = MagicMock() client.pipeline = MagicMock(return_value=pipe) - with patch.object(redis_cache, "init_async_client", return_value=client): + with patch.object(redis_cache, "_init_async_client", return_value=client): pushed_len = await redis_cache.async_rpush_and_trim(key="buf", values=["c", "d"], max_len=3) client.pipeline.assert_called_once_with(transaction=True) @@ -1692,7 +1845,7 @@ async def test_async_set_cache_pipeline_does_not_format_cached_values_for_loggin client.pipeline = MagicMock(return_value=pipe) _FormatCountingStr.format_calls = 0 - with patch.object(redis_cache, "init_async_client", return_value=client): + with patch.object(redis_cache, "_init_async_client", return_value=client): await redis_cache.async_set_cache_pipeline([("k1", _FormatCountingStr("embedding"))], ttl=60) assert _FormatCountingStr.format_calls == 0 diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 85422819926..14aed9bc1f1 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -60,8 +60,8 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): password="hello", ) - # Mock the init_async_client to return our mock redis client - cache.init_async_client = MagicMock(return_value=mock_redis) + # 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"] @@ -210,9 +210,9 @@ def test_init_pubsub_client_derives_a_node_client_for_cluster_backend() -> None: cluster = _cluster_for_pubsub() node = ClusterNode(host="10.1.2.3", port=7001) cluster.nodes_manager.default_node = node - cache.init_async_client = MagicMock(return_value=cluster) + cache._init_async_client = MagicMock(return_value=cluster) - client = cache.init_pubsub_client() + client = cache._init_pubsub_client() assert isinstance(client, Redis) and not isinstance(client, RedisCluster) kwargs = client.connection_pool.connection_kwargs @@ -228,9 +228,9 @@ def test_init_pubsub_client_falls_back_to_first_startup_node() -> None: startup node is a valid pub/sub target.""" cache = _isolated_redis_cache("cluster-pubsub-startup-fallback") cluster = _cluster_for_pubsub(startup_node_host="10.8.8.8") - cache.init_async_client = MagicMock(return_value=cluster) + cache._init_async_client = MagicMock(return_value=cluster) - client = cache.init_pubsub_client() + client = cache._init_pubsub_client() assert isinstance(client, Redis) assert client.connection_pool.connection_kwargs["host"] == "10.8.8.8" @@ -240,10 +240,10 @@ def test_init_pubsub_client_returns_the_same_cached_client_on_repeat_calls() -> cache = _isolated_redis_cache("cluster-pubsub-caching") cluster = _cluster_for_pubsub() cluster.nodes_manager.default_node = ClusterNode(host="10.1.2.3", port=7001) - cache.init_async_client = MagicMock(return_value=cluster) + cache._init_async_client = MagicMock(return_value=cluster) - first = cache.init_pubsub_client() - second = cache.init_pubsub_client() + first = cache._init_pubsub_client() + second = cache._init_pubsub_client() assert first is second @@ -251,9 +251,9 @@ def test_init_pubsub_client_returns_the_same_cached_client_on_repeat_calls() -> def test_init_pubsub_client_returns_the_shared_async_client_for_standalone() -> None: cache = _isolated_redis_cache("standalone-pubsub") standalone = Redis() - cache.init_async_client = MagicMock(return_value=standalone) + cache._init_async_client = MagicMock(return_value=standalone) - assert cache.init_pubsub_client() is standalone + assert cache._init_pubsub_client() is standalone def test_init_pubsub_client_preserves_tls_and_authentication() -> None: @@ -268,9 +268,9 @@ def test_init_pubsub_client_preserves_tls_and_authentication() -> None: socket_connect_timeout=3.0, socket_keepalive=True, ) - cache.init_async_client = MagicMock(return_value=cluster) + cache._init_async_client = MagicMock(return_value=cluster) - client = cache.init_pubsub_client() + client = cache._init_pubsub_client() connection = client.connection_pool.make_connection() assert isinstance(connection, SSLConnection) @@ -286,13 +286,13 @@ def test_init_pubsub_client_rejects_missing_nodes_and_can_retry() -> None: cache = _isolated_redis_cache("cluster-pubsub-no-nodes") cluster = _cluster_for_pubsub() cluster.nodes_manager.startup_nodes = {} - cache.init_async_client = MagicMock(return_value=cluster) + cache._init_async_client = MagicMock(return_value=cluster) with pytest.raises(ValueError, match="no default node and no startup nodes"): - cache.init_pubsub_client() + cache._init_pubsub_client() cluster.nodes_manager.default_node = ClusterNode(host="recovered.example.test", port=7001) - client = cache.init_pubsub_client() + client = cache._init_pubsub_client() assert client.connection_pool.connection_kwargs["host"] == "recovered.example.test" @@ -303,11 +303,11 @@ def test_disconnect_closes_derived_pubsub_connections_even_when_pool_close_fails close_fails: bool, has_shared_pool: bool ) -> None: cache = _isolated_redis_cache(f"cluster-pubsub-close-{close_fails}-{has_shared_pool}") - cache.async_redis_conn_pool = AsyncMock() if has_shared_pool else None - cache.init_async_client = MagicMock(return_value=_cluster_for_pubsub()) + cache._async_redis_conn_pool = AsyncMock() if has_shared_pool else None + cache._init_async_client = MagicMock(return_value=_cluster_for_pubsub()) async def exercise() -> None: - client = cache.init_pubsub_client() + client = cache._init_pubsub_client() connection = AsyncMock() connection.disconnect.side_effect = ConnectionError("connection close failed") if close_fails else None client.connection_pool._available_connections.append(connection) @@ -316,19 +316,19 @@ def test_disconnect_closes_derived_pubsub_connections_even_when_pool_close_fails connection.disconnect.assert_awaited_once() if has_shared_pool: - cache.async_redis_conn_pool.disconnect.assert_awaited_once_with(inuse_connections=True) - cache.redis_client.close.assert_called_once() + cache._async_redis_conn_pool.disconnect.assert_awaited_once_with(inuse_connections=True) + cache._redis_client.close.assert_called_once() asyncio.run(exercise()) def test_expired_pubsub_client_closes_connections_after_eviction() -> None: cache = _isolated_redis_cache("cluster-pubsub-expired") - cache.init_async_client = MagicMock(return_value=_cluster_for_pubsub()) + cache._init_async_client = MagicMock(return_value=_cluster_for_pubsub()) clients = LLMClientCache(evicted_client_closer=EvictedClientCloser(grace_seconds=0)) async def exercise() -> None: - client = cache.init_pubsub_client() + client = cache._init_pubsub_client() closed = asyncio.Event() connection = AsyncMock() connection.disconnect.side_effect = closed.set @@ -336,7 +336,7 @@ def test_expired_pubsub_client_closes_connections_after_eviction() -> None: cache_key = clients.update_cache_key_with_event_loop(f"{cache._get_async_client_cache_key()}-pubsub") clients.ttl_dict[cache_key] = 0 - replacement = cache.init_pubsub_client() + replacement = cache._init_pubsub_client() await asyncio.wait_for(closed.wait(), timeout=1) assert replacement is not client @@ -358,10 +358,10 @@ def test_init_async_client_returns_one_cluster_client_per_event_loop() -> None: loop_b = asyncio.new_event_loop() async def get_client() -> RedisCluster: - return cache.init_async_client() + return cache._init_async_client() async def get_client_pair() -> tuple[RedisCluster, RedisCluster]: - return cache.init_async_client(), cache.init_async_client() + return cache._init_async_client(), cache._init_async_client() try: client_a, client_a_again = loop_a.run_until_complete(get_client_pair()) @@ -395,13 +395,13 @@ def test_evicted_async_client_is_closed_only_for_cluster(cluster: bool) -> None: ) async def exercise() -> None: - client = cache.init_async_client() + client = cache._init_async_client() closed = asyncio.Event() client.aclose = AsyncMock(side_effect=closed.set) cache_key = clients.update_cache_key_with_event_loop(cache._get_async_client_cache_key()) clients.ttl_dict[cache_key] = 0 - replacement = cache.init_async_client() + replacement = cache._init_async_client() assert replacement is not client if cluster: @@ -445,7 +445,7 @@ def _cluster_set_cache_with_client( max_connections=max_connections, connection_gate=connection_gate, ) - setattr(cache, "init_async_client", MagicMock(return_value=client)) + setattr(cache, "_init_async_client", MagicMock(return_value=client)) return cast( "Callable[[str, str], Awaitable[object]]", cache.async_set_cache ) # cast-ok: wrapped **kwargs are untyped @@ -492,7 +492,7 @@ async def test_cluster_cache_sends_after_the_connection_wait_timeout() -> None: @pytest.mark.asyncio async def test_cluster_caches_sharing_a_client_share_the_connection_cap() -> None: - """init_async_client hands caches with identical settings the same client, so two caches must not each admit + """_init_async_client hands caches with identical settings the same client, so two caches must not each admit a full max_connections worth of commands against that one pool.""" client = _InFlightCountingClient(command_seconds=0.01) first = _cluster_set_cache_with_client(client, max_connections=3) diff --git a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py index 8d4571c2f01..11c201e07af 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_gateway_endpoints.py @@ -107,9 +107,12 @@ class _SharedRedisFake: async def async_get_cache(self, key: str, **kwargs: object) -> object: return self.values.get(self.check_and_fix_namespace(key)) - def init_async_client(self) -> "_SharedRedisClientFake": + def _init_async_client(self) -> "_SharedRedisClientFake": return _SharedRedisClientFake(self) + def _async_commands(self) -> "_SharedRedisClientFake": + return self._init_async_client() + class _SharedRedisClientFake: def __init__(self, store: _SharedRedisFake) -> None: diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 9d4fe67614e..91e3d606039 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -2299,9 +2299,12 @@ class _PubSubRedisCache: def __init__(self, client: _SlowPublishRedisClient | _WedgedPublishRedisClient) -> None: self._client = client - def init_pubsub_client(self) -> _SlowPublishRedisClient | _WedgedPublishRedisClient: + def _init_pubsub_client(self) -> _SlowPublishRedisClient | _WedgedPublishRedisClient: return self._client + def _pubsub_client(self) -> _SlowPublishRedisClient | _WedgedPublishRedisClient: + return self._init_pubsub_client() + async def _drain_publishes() -> None: await asyncio.gather(*pubsub_module._pending_publishes) # pyright: ignore[reportPrivateUsage] # drain module-level tasks diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 6b59f6d40dc..1acc924f6c0 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10769,9 +10769,12 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t def __init__(self) -> None: self.namespace = None - def init_pubsub_client(self) -> object: + def _init_pubsub_client(self) -> object: return _RecordingRedisClient() + def _pubsub_client(self) -> object: + return self._init_pubsub_client() + local_spend_counter_cache = DualCache() remote_user_api_key_cache = UserApiKeyCache() @@ -10843,9 +10846,12 @@ async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_ def __init__(self) -> None: self.namespace = None - def init_pubsub_client(self) -> object: + def _init_pubsub_client(self) -> object: return _RecordingRedisClient() + def _pubsub_client(self) -> object: + return self._init_pubsub_client() + local_spend_counter_cache = DualCache() local_user_api_key_cache = UserApiKeyCache() diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 694b817ab86..7d801995771 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -141,7 +141,7 @@ class _FakeRedisCache: self._client = client self.namespace = namespace - def init_pubsub_client(self) -> object: + def _pubsub_client(self) -> object: return self._client diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index cbb37c964c6..f50adb26461 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -85,7 +85,7 @@ class _FailingPublishRedisClient(Redis): class _ScriptedPubSubClient: """Pub/sub-capable client that is not a redis.asyncio.Redis. - Mirrors what RedisCache.init_pubsub_client returns for a cluster backend: + Mirrors what RedisCache._pubsub_client returns for a cluster backend: a node-level client exposing publish/pubsub without being an instance of the standalone Redis class. """ @@ -203,14 +203,14 @@ class _FakeRedisCache: self._client = client self.namespace = namespace - def init_pubsub_client(self) -> object: + def _pubsub_client(self) -> object: return self._client class _ExplodingRedisCache: namespace: Optional[str] = None - def init_pubsub_client(self) -> object: + def _pubsub_client(self) -> object: raise ConnectionError("cannot connect") @@ -270,7 +270,7 @@ async def test_publish_swallows_client_init_errors() -> None: async def test_publish_reaches_cluster_derived_pubsub_clients() -> None: """LIT-8543: a cluster-backed cache returns a node-level client from - init_pubsub_client; publishes must go out on it instead of being skipped.""" + _pubsub_client; publishes must go out on it instead of being skipped.""" client = _ScriptedPubSubClient(pubsubs=[]) cache = _FakeRedisCache(client) @@ -775,7 +775,7 @@ async def test_stop_before_start_is_a_noop() -> None: async def test_subscriber_subscribes_on_cluster_derived_pubsub_client() -> None: """LIT-8543: the subscriber used to disable itself on cluster caches; now it - subscribes on the node-level client init_pubsub_client returns.""" + subscribes on the node-level client _pubsub_client returns.""" pubsub = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_proxymodeltable"})]) cache = _FakeRedisCache(_ScriptedPubSubClient(pubsubs=[pubsub])) resyncs: List[str] = [] diff --git a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py index dae3a0459e2..172d25cd363 100644 --- a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py +++ b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py @@ -48,7 +48,7 @@ def _fake_redis(monkeypatch): fake = fakeredis.FakeAsyncRedis() redis = RedisCache(host="localhost", port=6379) - monkeypatch.setattr(redis, "init_async_client", lambda: fake) + monkeypatch.setattr(redis, "_init_async_client", lambda: fake) return redis diff --git a/tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py b/tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py index d2c6b661068..c18f68aff38 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py @@ -389,7 +389,7 @@ async def test_e2e_size_of_redis_buffer(): password=os.getenv("REDIS_PASSWORD"), ) fake_redis_client = fakeredis.FakeAsyncRedis() - redis_cache.redis_async_client = fake_redis_client + redis_cache._redis_async_client = fake_redis_client setattr(litellm.proxy.proxy_server, "use_redis_transaction_buffer", True) db_writer = DBSpendUpdateWriter(redis_cache=redis_cache) diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 39586a1a204..a570d0129d9 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2539,9 +2539,12 @@ async def test_user_update_reaches_cached_user_on_every_worker( class _FakeRedisCache: namespace: str | None = None - def init_pubsub_client(self) -> _RecordingRedisClient: + def _init_pubsub_client(self) -> _RecordingRedisClient: return _RecordingRedisClient() + def _pubsub_client(self) -> _RecordingRedisClient: + return self._init_pubsub_client() + class _UserTableMocks: def __init__( self, diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index c26a4bf19b8..3647c88ba3b 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -5808,9 +5808,12 @@ class _SharedCoordinationRedis: self.published: tuple[str, ...] = () self._subscriptions: tuple[_ChannelSubscription, ...] = () - def init_pubsub_client(self) -> "_SharedCoordinationRedis": + def _init_pubsub_client(self) -> "_SharedCoordinationRedis": return self + def _pubsub_client(self) -> "_SharedCoordinationRedis": + return self._init_pubsub_client() + def pubsub(self) -> _ChannelSubscription: subscription: Final = _ChannelSubscription() self._subscriptions = (*self._subscriptions, subscription) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index e33743bb149..7d9d8338b43 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -2575,7 +2575,7 @@ async def test_refresh_cached_team_flushes_redis_set_before_broadcast(monkeypatc ) self.default_ttl: Final[int | None] = None - def init_async_client( + def _init_async_client( self, ) -> RecordingRedisClient: # pyright: ignore[reportIncompatibleMethodOverride] # in-memory Redis client return self.client @@ -10265,9 +10265,12 @@ class _BroadcastRecordingRedis: self.broadcast_keys: list[str] = [] self._round_trip_seconds = round_trip_seconds - def init_pubsub_client(self) -> "_BroadcastRecordingRedis": + def _init_pubsub_client(self) -> "_BroadcastRecordingRedis": return self + def _pubsub_client(self) -> "_BroadcastRecordingRedis": + return self._init_pubsub_client() + async def send(self, messages: Sequence[str]) -> None: if self._round_trip_seconds: await asyncio.sleep(self._round_trip_seconds) diff --git a/tests/unit/proxy/management_helpers/test_team_roster_sync.py b/tests/unit/proxy/management_helpers/test_team_roster_sync.py index 74c3dee9d56..b3bb1aa30d2 100644 --- a/tests/unit/proxy/management_helpers/test_team_roster_sync.py +++ b/tests/unit/proxy/management_helpers/test_team_roster_sync.py @@ -582,9 +582,12 @@ class _PubSubRedisCache: def __init__(self, client: "_SlowPublishRedisClient | _WedgedPublishRedisClient") -> None: self._client = client - def init_pubsub_client(self) -> "_SlowPublishRedisClient | _WedgedPublishRedisClient": + def _init_pubsub_client(self) -> "_SlowPublishRedisClient | _WedgedPublishRedisClient": return self._client + def _pubsub_client(self) -> "_SlowPublishRedisClient | _WedgedPublishRedisClient": + return self._init_pubsub_client() + async def _until_reached_or_done(signal: asyncio.Event, sync: "asyncio.Task[object]") -> None: reached: Final = asyncio.create_task(signal.wait()) diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py index 18bd4a66b65..897f49aed7a 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py @@ -1,9 +1,14 @@ """Tests for the cross-replica refresh coordinator: winner refreshes, losers wait then re-read.""" import asyncio +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import RedisDistributedLock from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) @@ -11,6 +16,9 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator imp LockAcquisition, RedisRefreshCoordinator, ) +from litellm.proxy.mcp_server.outbound_credentials.runtime_refresh_coordinator import ( + runtime_refresh_coordinator, +) _KEY = "mcp:refresh_lock:u:s" @@ -55,6 +63,30 @@ def _advancing_sleep(clock, step=0.5): return sleep +def test_runtime_factory_uses_legacy_redis_initializer(monkeypatch: pytest.MonkeyPatch) -> None: + override_client: Final = AsyncMock() + + with pytest.warns(DeprecationWarning, match="overrides a RedisCache raw client accessor"): + + class LegacyAsyncCache(RedisCache): + def init_async_client(self) -> AsyncMock: + return override_client + + with patch("litellm._redis.get_redis_client", return_value=MagicMock()), patch( + "litellm._redis.get_redis_connection_pool", return_value=None + ): + redis_cache: Final = LegacyAsyncCache(host="127.0.0.1", port=6379) + + cache_wrapper: Final = SimpleNamespace(redis_cache=redis_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache_wrapper) + + coordinator: Final = runtime_refresh_coordinator() + + assert isinstance(coordinator, RedisRefreshCoordinator) + assert isinstance(coordinator.lock, RedisDistributedLock) + assert coordinator.lock.client is override_client + + @pytest.mark.asyncio async def test_winner_acquires_and_releases_with_the_same_token_and_never_rereads(): lock = _FakeLock(acquired=LockAcquisition.ACQUIRED) diff --git a/tests/unit/proxy/mcp_server/test_byok_credential_cache.py b/tests/unit/proxy/mcp_server/test_byok_credential_cache.py index 7ca169366a6..ebfa330e9e6 100644 --- a/tests/unit/proxy/mcp_server/test_byok_credential_cache.py +++ b/tests/unit/proxy/mcp_server/test_byok_credential_cache.py @@ -16,9 +16,12 @@ from litellm.proxy.mcp_server.byok_credential_cache import ( class _FakeRedisCache: namespace = None - def init_pubsub_client(self) -> object: + def _init_pubsub_client(self) -> object: return object() + def _pubsub_client(self) -> object: + return self._init_pubsub_client() + @pytest.fixture(autouse=True) def _empty_cache(): diff --git a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py index a426cea131a..ceeec007cd0 100644 --- a/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py @@ -740,12 +740,12 @@ async def test_single_use_guard_peek_fails_closed_when_redis_errors(): cache.redis_cache = MagicMock() cache.redis_cache.async_get_cache = AsyncMock(return_value=None) cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) - cache.redis_cache.init_async_client.return_value.get = AsyncMock(side_effect=ConnectionError("redis down")) + cache.redis_cache._async_commands.return_value.get = AsyncMock(side_effect=ConnectionError("redis down")) guard = _SingleUseGuard(cache) assert await guard.peek("family-fault") == "unavailable" - cache.redis_cache.init_async_client.return_value.get = AsyncMock(return_value=b"1") + cache.redis_cache._async_commands.return_value.get = AsyncMock(return_value=b"1") assert await guard.peek("family-fault") == "claimed" @@ -774,7 +774,7 @@ async def test_single_use_guard_peek_reads_the_key_under_the_namespace_claim_wro cache.redis_cache = MagicMock() cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=namespaced) cache.redis_cache.async_increment = AsyncMock(side_effect=increment) - cache.redis_cache.init_async_client.return_value.get = AsyncMock(side_effect=raw_get) + cache.redis_cache._async_commands.return_value.get = AsyncMock(side_effect=raw_get) guard = _SingleUseGuard(cache) assert await guard.peek("family-ns") == "unclaimed" @@ -2254,7 +2254,7 @@ def _redis_that(async_increment, get=None): cache.redis_cache = MagicMock() cache.redis_cache.async_increment = async_increment cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) - cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None) + cache.redis_cache._async_commands.return_value.get = get or AsyncMock(return_value=None) cache.redis_cache.async_get_cache = AsyncMock( side_effect=AssertionError("peek must read the client, not the wrapper") ) @@ -2316,7 +2316,7 @@ class _RedisThatFaultsOnTheSecondWrite: def check_and_fix_namespace(self, key: str) -> str: return key - def init_async_client(self) -> "_RedisThatFaultsOnTheSecondWrite": + def _async_commands(self) -> "_RedisThatFaultsOnTheSecondWrite": return self async def async_increment(self, key: str, value: int, **kwargs: object) -> int: diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..b30c09d3e30 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -1812,7 +1812,7 @@ class _RedisCacheOnClient(RedisCache): self.service_logger_obj = ServiceLogging() self.default_ttl = None - def init_async_client(self) -> FakeAsyncRedis: + def _init_async_client(self) -> FakeAsyncRedis: return self._injected_client diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 087b273548a..2f135907cc7 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -550,13 +550,13 @@ class TestResponsePollingHandler: """Test that delete_polling removes the entry from Redis""" mock_redis = AsyncMock() mock_async_client = AsyncMock() - mock_redis.redis_async_client = True # hasattr check - # init_async_client is a sync method that returns an async client - mock_redis.init_async_client = Mock(return_value=mock_async_client) + mock_redis._redis_async_client = True # hasattr check + # _init_async_client is a sync method that returns an async client + mock_redis._init_async_client = Mock(return_value=mock_async_client) - # Mock async_delete_cache to actually call init_async_client and delete + # Mock async_delete_cache to actually call _init_async_client and delete async def mock_async_delete_cache(key): - client = mock_redis.init_async_client() + client = mock_redis._init_async_client() await client.delete(key) mock_redis.async_delete_cache = mock_async_delete_cache diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 139cfbbe58c..354c3747c4f 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -1326,7 +1326,7 @@ from litellm.proxy.proxy_server import ProxyConfig @pytest.mark.skip(reason="local variable conflicts. needs to be refactored.") @mock.patch("litellm.proxy.proxy_server.litellm.Cache") def test_load_router_config(mock_cache, fake_env_vars): - mock_cache.return_value.cache.__dict__ = {"redis_client": None} + mock_cache.return_value.cache.__dict__ = {"_redis_client": None} mock_cache.return_value.supported_call_types = [ "completion", "acompletion", @@ -1382,7 +1382,7 @@ def test_load_router_config(mock_cache, fake_env_vars): ) ) assert litellm.cache is not None - assert "redis_client" in vars( + assert "_redis_client" in vars( litellm.cache.cache ) # it should default to redis on proxy assert litellm.cache.supported_call_types == [ @@ -1411,7 +1411,7 @@ def test_load_router_config(mock_cache, fake_env_vars): print(litellm.cache) print(litellm.cache.supported_call_types) print(vars(litellm.cache.cache)) - assert "redis_client" in vars( + assert "_redis_client" in vars( litellm.cache.cache ) # it should default to redis on proxy assert litellm.cache.supported_call_types == [ diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index bc3a228321a..1bad7472250 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15678,9 +15678,12 @@ async def test_auth_cache_invalidation_subscriber_evicts_byok_credentials_cached def __init__(self, client: object) -> None: self._client = client - def init_pubsub_client(self) -> object: + def _init_pubsub_client(self) -> object: return self._client + def _pubsub_client(self) -> object: + return self._init_pubsub_client() + byok_credential_cache.flush_cache() cache_byok_credential("mallory", "srv-byok", "sk-revoked-elsewhere") message: Final = { diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..23ed7da2796 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -288,7 +288,7 @@ async def test_delayed_shared_pin_read_preserves_a_newer_local_write( await resume.wait() return value - reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool) + reader: Final = DelayedReader(connection_pool=redis._init_async_client().connection_pool) litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader) pending: Final = asyncio.create_task(_selected_session_ids(router)) try: diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 819c6b9d4f5..7aca3b63ae2 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -1734,7 +1734,7 @@ class _RedisCacheOnFakeServer(RedisCache): super().__init__(host="io-token-turn-redis", connection_pool=FakeRedis(server=server).connection_pool) self._fake_async_client: Final = FakeAsyncRedis(server=server) - def init_async_client(self) -> FakeAsyncRedis: + def _init_async_client(self) -> FakeAsyncRedis: return self._fake_async_client diff --git a/tests/unit/test_router_redis_init.py b/tests/unit/test_router_redis_init.py index 4a8a5b57622..9b429e869ff 100644 --- a/tests/unit/test_router_redis_init.py +++ b/tests/unit/test_router_redis_init.py @@ -38,7 +38,7 @@ async def test_router_uses_correct_redis_db(): try: if router.cache.redis_cache: # Check connection kwargs or internal client db - cache_client = router.cache.redis_cache.redis_client + cache_client = router.cache.redis_cache._redis_client # Redis client stores connection args in connection_pool.connection_kwargs conn_kwargs = cache_client.connection_pool.connection_kwargs