Added missing methods

redis/asyncio/cluster.py

@@ -3,7 +3,6 @@
 import random
 import socket
 import ssl
-import threading
 import warnings
 from typing import (
     Any,
@@ -49,7 +48,6 @@
 from redis.credentials import CredentialProvider
 from redis.event import (
     AfterAsyncClusterInstantiationEvent,
-    AsyncAfterConnectionReleasedEvent,
     EventDispatcher,
 )
 from redis.exceptions import (

redis/asyncio/connection.py

@@ -5,7 +5,6 @@
 import socket
 import ssl
 import sys
-import threading
 import warnings
 import weakref
 from abc import abstractmethod
@@ -44,7 +43,6 @@
 from redis.connection import DEFAULT_RESP_VERSION
 from redis.credentials import (
     CredentialProvider,
-    StreamingCredentialProvider,
     UsernamePasswordCredentialProvider,
 )
 from redis.exceptions import (

redis/auth/idp.py

@@ -9,7 +9,8 @@
 
 class IdentityProviderInterface(ABC):
     """
-    Receive a token from the identity provider. Receiving a token only works when being authenticated.
+    Receive a token from the identity provider.
+    Receiving a token only works when being authenticated.
     """
 
     @abstractmethod

redis/auth/token_manager.py

@@ -3,7 +3,7 @@
 import threading
 from datetime import datetime, timezone
 from time import sleep
-from typing import Any, Awaitable, Callable, Coroutine, Union
+from typing import Any, Awaitable, Callable, Union
 
 from redis.auth.err import RequestTokenErr, TokenRenewalErr
 from redis.auth.idp import IdentityProviderInterface
@@ -78,9 +78,9 @@ def __init__(
 
     def get_expiration_refresh_ratio(self) -> float:
         """
-        Represents the ratio of a token's lifetime at which a refresh should be triggered.
-        For example, a value of 0.75 means the token should be refreshed when 75% of its
-        lifetime has elapsed (or when 25% of its lifetime remains).
+        Represents the ratio of a token's lifetime at which a refresh should be triggered. # noqa: E501
+        For example, a value of 0.75 means the token should be refreshed
+        when 75% of its lifetime has elapsed (or when 25% of its lifetime remains).
 
         :return: float
         """
@@ -89,9 +89,10 @@ def get_expiration_refresh_ratio(self) -> float:
 
     def get_lower_refresh_bound_millis(self) -> int:
         """
-        Represents the minimum time in milliseconds before token expiration to trigger a refresh, in milliseconds.
-        This value sets a fixed lower bound for when a token refresh should occur, regardless
-        of the token's total lifetime.
+        Represents the minimum time in milliseconds before token expiration
+        to trigger a refresh, in milliseconds.
+        This value sets a fixed lower bound for when a token refresh should occur,
+        regardless of the token's total lifetime.
         If set to 0 there will be no lower bound and the refresh will be triggered
         based on the expirationRefreshRatio only.
 
@@ -101,7 +102,8 @@ def get_lower_refresh_bound_millis(self) -> int:
 
     def get_token_request_execution_timeout_in_ms(self) -> int:
         """
-        Represents the maximum time in milliseconds to wait for a token request to complete.
+        Represents the maximum time in milliseconds to wait
+        for a token request to complete.
 
         :return: int
         """

redis/client.py

@@ -1,4 +1,3 @@
-import asyncio
 import copy
 import re
 import threading

redis/cluster.py

@@ -20,7 +20,6 @@
     AfterPubSubConnectionInstantiationEvent,
     ClientType,
     EventDispatcher,
-    EventDispatcherInterface,
 )
 from redis.exceptions import (
     AskError,

redis/connection.py

@@ -26,13 +26,11 @@
 from .backoff import NoBackoff
 from .credentials import (
     CredentialProvider,
-    StreamingCredentialProvider,
     UsernamePasswordCredentialProvider,
 )
 from .event import (
     AfterConnectionReleasedEvent,
     EventDispatcher,
-    EventDispatcherInterface,
 )
 from .exceptions import (
     AuthenticationError,
@@ -790,6 +788,7 @@ def __init__(
         self.retry = self._conn.retry
         self.host = self._conn.host
         self.port = self._conn.port
+        self.credential_provider = conn.credential_provider
         self._pool_lock = pool_lock
         self._cache = cache
         self._cache_lock = threading.Lock()
@@ -973,6 +972,17 @@ def _on_invalidation_callback(self, data: List[Union[str, Optional[List[bytes]]]
             else:
                 self._cache.delete_by_redis_keys(data[1])
 
+    def get_protocol(self):
+        return self._conn.get_protocol()
+
+    @property
+    def set_re_auth_token(self, token: TokenInterface):
+        self._conn.set_re_auth_token(token)
+
+    @property
+    def re_auth(self):
+        self._conn.re_auth()
+
 
 class SSLConnection(Connection):
     """Manages SSL connections to and from the Redis server(s).

redis/credentials.py

@@ -15,7 +15,8 @@ def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
 
     async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
         logger.warning(
-            "This method is added for backward compatability. Please override it in your implementation."
+            "This method is added for backward compatability. "
+            "Please override it in your implementation."
         )
         return self.get_credentials()
 
@@ -28,7 +29,8 @@ class StreamingCredentialProvider(CredentialProvider, ABC):
     @abstractmethod
     def on_next(self, callback: Callable[[Any], None]):
         """
-        Specifies the callback that should be invoked when the next credentials will be retrieved.
+        Specifies the callback that should be invoked
+        when the next credentials will be retrieved.
 
         :param callback: Callback with
         :return:

redis/event.py

@@ -30,7 +30,8 @@ async def listen(self, event: object):
 
 class EventDispatcherInterface(ABC):
     """
-    Represents a dispatcher that dispatches events to listeners associated with given event.
+    Represents a dispatcher that dispatches events to listeners
+    associated with given event.
     """
 
     @abstractmethod
@@ -132,7 +133,8 @@ class AfterSingleConnectionInstantiationEvent:
     """
     Event that will be fired after single connection instances was created.
 
-    :param connection_lock: For sync client thread-lock should be provided, for async asyncio.Lock
+    :param connection_lock: For sync client thread-lock should be provided,
+    for async asyncio.Lock
     """
 
     def __init__(
@@ -192,7 +194,8 @@ class AfterAsyncClusterInstantiationEvent:
     """
     Event that will be fired after async cluster instance was created.
 
-    Async cluster doesn't use connection pools, instead ClusterNode object manages connections.
+    Async cluster doesn't use connection pools,
+    instead ClusterNode object manages connections.
     """
 
     def __init__(

tests/conftest.py

@@ -17,7 +17,6 @@
 from redis import Sentinel
 from redis.auth.idp import IdentityProviderInterface
 from redis.auth.token import JWToken
-from redis.auth.token_manager import TokenManager
 from redis.backoff import NoBackoff
 from redis.cache import (
     CacheConfig,
@@ -27,7 +26,7 @@
     EvictionPolicy,
 )
 from redis.connection import Connection, ConnectionInterface, SSLConnection, parse_url
-from redis.credentials import CredentialProvider, StreamingCredentialProvider
+from redis.credentials import CredentialProvider
 from redis.exceptions import RedisClusterException
 from redis.retry import Retry
 from redis_entraid.cred_provider import EntraIdCredentialsProvider, TokenAuthConfig
@@ -329,7 +328,8 @@ def skip_if_nocryptography() -> _TestDecorator:
     #
     #     return pytest.mark.skipif(False, reason="Cryptography dependency found")
     # except ImportError:
-    # TODO: Because JWT library depends on cryptography, now it's always true and tests should be fixed
+    # TODO: Because JWT library depends on cryptography,
+    #  now it's always true and tests should be fixed
     return pytest.mark.skipif(True, reason="No cryptography dependency")
 
 

tests/test_asyncio/test_credentials.py

@@ -17,8 +17,7 @@
 from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider
 from redis.exceptions import ConnectionError
 from redis.utils import str_if_bytes
-from redis_entraid.cred_provider import EntraIdCredentialsProvider, TokenAuthConfig
-from redis_entraid.identity_provider import create_provider_from_service_principal
+from redis_entraid.cred_provider import EntraIdCredentialsProvider
 from tests.conftest import get_endpoint, skip_if_redis_enterprise
 from tests.test_asyncio.conftest import get_credential_provider
 

tests/test_auth/test_token_manager.py

@@ -205,7 +205,8 @@ def on_next(token):
         config = TokenManagerConfig(1, 0, 1000, retry_policy)
         mgr = TokenManager(mock_provider, config)
         mgr.start(mock_listener, skip_initial=True)
-        # Should be less than a 0.1, or it will be flacky due to additional token renewal.
+        # Should be less than a 0.1, or it will be flacky due to
+        # additional token renewal.
         sleep(0.2)
 
         assert len(tokens) == 2
@@ -246,7 +247,8 @@ async def on_next(token):
         config = TokenManagerConfig(1, 0, 1000, retry_policy)
         mgr = TokenManager(mock_provider, config)
         await mgr.start_async(mock_listener, skip_initial=True)
-        # Should be less than a 0.1, or it will be flacky due to additional token renewal.
+        # Should be less than a 0.1, or it will be flacky
+        # due to additional token renewal.
         await asyncio.sleep(0.2)
 
         assert len(tokens) == 2
@@ -282,7 +284,8 @@ def on_next(token):
         config = TokenManagerConfig(1, 0, 1000, retry_policy)
         mgr = TokenManager(mock_provider, config)
         mgr.start(mock_listener)
-        # Should be less than a 0.1, or it will be flacky due to additional token renewal.
+        # Should be less than a 0.1, or it will be flacky
+        # due to additional token renewal.
         sleep(0.08)
 
         assert mock_provider.request_token.call_count in {3, 4}
@@ -321,7 +324,8 @@ async def on_next(token):
         config = TokenManagerConfig(1, 0, 1000, retry_policy)
         mgr = TokenManager(mock_provider, config)
         await mgr.start_async(mock_listener, block_for_initial=True)
-        # Should be less than a 0.1, or it will be flacky due to additional token renewal.
+        # Should be less than a 0.1, or it will be flacky
+        # due to additional token renewal.
         await asyncio.sleep(0.08)
 
         assert mock_provider.request_token.call_count in {3, 4}

tests/test_credentials.py

@@ -1,31 +1,22 @@
 import functools
-import json
-import os
 import random
 import string
 import threading
-from asyncio import Lock as AsyncioLock
-from asyncio import sleep as asyncio_sleep
 from time import sleep
 from typing import Optional, Tuple, Union
 
 import pytest
 import redis
 from mock.mock import Mock, call
 from redis import AuthenticationError, DataError, Redis, ResponseError, asyncio
-from redis.asyncio import Connection
-from redis.asyncio import ConnectionPool as AsyncConnectionPool
-from redis.asyncio import Redis as AsyncRedis
 from redis.auth.err import RequestTokenErr
-from redis.auth.idp import IdentityProviderInterface
 from redis.backoff import NoBackoff
 from redis.connection import ConnectionInterface, ConnectionPool
 from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider
 from redis.exceptions import ConnectionError, RedisError
 from redis.retry import Retry
 from redis.utils import str_if_bytes
-from redis_entraid.cred_provider import EntraIdCredentialsProvider, TokenAuthConfig
-from redis_entraid.identity_provider import create_provider_from_service_principal
+from redis_entraid.cred_provider import EntraIdCredentialsProvider
 from tests.conftest import (
     _get_client,
     get_credential_provider,