Added getter and setter for the client's retry object and added more tests

Author: barshaul

redis/asyncio/client.py

@@ -273,6 +273,12 @@ def get_connection_kwargs(self):
         """Get the connection's key-word arguments"""
         return self.connection_pool.connection_kwargs
 
+    def get_retry(self) -> Optional["Retry"]:
+        return self.get_connection_kwargs().get("retry")
+
+    def set_retry(self, retry: "Retry") -> None:
+        self.get_connection_kwargs().update({"retry": retry})
+
     def load_external_module(self, funcname, func):
         """
         This function can be used to add externally defined redis modules,

redis/asyncio/cluster.py

@@ -495,6 +495,14 @@ def get_connection_kwargs(self) -> Dict[str, Optional[Any]]:
         """Get the kwargs passed to :class:`~redis.asyncio.connection.Connection`."""
         return self.connection_kwargs
 
+    def get_retry(self) -> Optional["Retry"]:
+        return self.retry
+
+    def set_retry(self, retry: "Retry") -> None:
+        self.retry = retry
+        for node in self.get_nodes():
+            node.connection_kwargs.update({"retry": retry})
+
     def set_response_callback(self, command: str, callback: ResponseCallbackT) -> None:
         """Set a custom response callback."""
         self.response_callbacks[command] = callback

redis/client.py

@@ -5,6 +5,7 @@
 import time
 import warnings
 from itertools import chain
+from typing import Optional
 
 from redis.commands import (
     CoreCommands,
@@ -1043,6 +1044,12 @@ def get_connection_kwargs(self):
         """Get the connection's key-word arguments"""
         return self.connection_pool.connection_kwargs
 
+    def get_retry(self) -> Optional["Retry"]:
+        return self.get_connection_kwargs().get("retry")
+
+    def set_retry(self, retry: "Retry") -> None:
+        self.get_connection_kwargs().update({"retry": retry})
+
     def set_response_callback(self, command, callback):
         """Set a custom Response Callback"""
         self.response_callbacks[command] = callback

redis/cluster.py

@@ -687,6 +687,14 @@ def set_default_node(self, node):
         self.nodes_manager.default_node = node
         return True
 
+    def get_retry(self) -> Optional["Retry"]:
+        return self.retry
+
+    def set_retry(self, retry: "Retry") -> None:
+        self.retry = retry
+        for node in self.get_nodes():
+            node.redis_connection.set_retry(retry)
+
     def monitor(self, target_node=None):
         """
         Returns a Monitor object for the specified target node.

tests/test_asyncio/test_cluster.py

@@ -249,9 +249,32 @@ async def test_startup_nodes(self) -> None:
                 ]
             )
 
-    async def test_cluster_retry_object(self) -> None:
+    async def test_cluster_get_set_retry_object(self, request):
+        retry = Retry(NoBackoff(), 2)
+        url = request.config.getoption("--redis-url")
+        r = await RedisCluster.from_url(url, retry=retry)
+        assert r.get_retry()._retries == retry._retries
+        assert isinstance(r.get_retry()._backoff, NoBackoff)
+        for node in r.get_nodes():
+            n_retry = node.connection_kwargs.get("retry")
+            assert n_retry is not None
+            assert n_retry._retries == retry._retries
+            assert isinstance(n_retry._backoff, NoBackoff)
+        # Change retry policy
+        new_retry = Retry(ExponentialBackoff(), 3)
+        r.set_retry(new_retry)
+        assert r.get_retry()._retries == new_retry._retries
+        assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
+        for node in r.get_nodes():
+            n_retry = node.connection_kwargs.get("retry")
+            assert n_retry is not None
+            assert n_retry._retries == new_retry._retries
+            assert isinstance(n_retry._backoff, ExponentialBackoff)
+
+    async def test_cluster_retry_object(self, request) -> None:
+        url = request.config.getoption("--redis-url")
+        rc_default = await RedisCluster.from_url(url)
         # Test default retry
-        rc_default = await RedisCluster("127.0.0.1", 16379)
         retry = rc_default.connection_kwargs.get("retry")
         assert isinstance(retry, Retry)
         assert retry._retries == 3
@@ -262,23 +285,21 @@ async def test_cluster_retry_object(self) -> None:
 
         # Test custom retry
         retry = Retry(ExponentialBackoff(10, 5), 5)
-        rc_custom_retry = await RedisCluster("127.0.0.1", 16379, retry=retry)
+        rc_custom_retry = await RedisCluster.from_url(url, retry=retry)
         assert (
             rc_custom_retry.get_node("127.0.0.1", 16379).connection_kwargs.get("retry")
             == retry
         )
 
         # Test no connection retries
-        rc_no_retries = await RedisCluster(
-            "127.0.0.1", 16379, connection_error_retry_attempts=0
+        rc_no_retries = await RedisCluster.from_url(
+            url, connection_error_retry_attempts=0
         )
         assert (
             rc_no_retries.get_node("127.0.0.1", 16379).connection_kwargs.get("retry")
             is None
         )
-        rc_no_retries = await RedisCluster(
-            "127.0.0.1", 16379, retry=Retry(NoBackoff(), 0)
-        )
+        rc_no_retries = await RedisCluster.from_url(url, retry=Retry(NoBackoff(), 0))
         assert (
             rc_no_retries.get_node("127.0.0.1", 16379)
             .connection_kwargs.get("retry")

tests/test_asyncio/test_retry.py

@@ -1,8 +1,9 @@
 import pytest
 
+from redis.asyncio import Redis
 from redis.asyncio.connection import Connection, UnixDomainSocketConnection
 from redis.asyncio.retry import Retry
-from redis.backoff import AbstractBackoff, NoBackoff
+from redis.backoff import AbstractBackoff, ExponentialBackoff, NoBackoff
 from redis.exceptions import ConnectionError, TimeoutError
 
 
@@ -114,3 +115,18 @@ async def test_infinite_retry(self):
 
         assert self.actual_attempts == 5
         assert self.actual_failures == 5
+
+
+class TestRedisClientRetry:
+    "Test the Redis client behavior with retries"
+
+    async def test_get_set_retry_object(self, request):
+        retry = Retry(NoBackoff(), 2)
+        url = request.config.getoption("--redis-url")
+        r = await Redis.from_url(url, retry_on_timeout=True, retry=retry)
+        assert r.get_retry()._retries == retry._retries
+        assert isinstance(r.get_retry()._backoff, NoBackoff)
+        new_retry_policy = Retry(ExponentialBackoff(), 3)
+        r.set_retry(new_retry_policy)
+        assert r.get_retry()._retries == new_retry_policy._retries
+        assert isinstance(r.get_retry()._backoff, ExponentialBackoff)

tests/test_cluster.py

@@ -360,6 +360,60 @@ def ok_response(connection, *args, **options):
 
             assert r.execute_command("SET", "foo", "bar") == "MOCK_OK"
 
+    def test_handling_cluster_failover_to_a_replica(self, r):
+        # Set the key we'll test for
+        key = "key"
+        r.set("key", "value")
+        primary = r.get_node_from_key(key, replica=False)
+        assert r.get("key") == "value"
+        # Get the current output of cluster slots
+        cluster_slots = primary.redis_connection.execute_command("CLUSTER SLOTS")
+        replica_host = ""
+        replica_port = 0
+        # Replace one of the replicas to be the new primary based on the
+        # cluster slots output
+        for slot_range in cluster_slots:
+            primary_port = slot_range[2][1]
+            if primary_port == primary.port:
+                if len(slot_range) <= 3:
+                    # cluster doesn't have a replica, return
+                    return
+                replica_host = str_if_bytes(slot_range[3][0])
+                replica_port = slot_range[3][1]
+                # replace replica and primary in the cluster slots output
+                tmp_node = slot_range[2]
+                slot_range[2] = slot_range[3]
+                slot_range[3] = tmp_node
+                break
+
+        def raise_connection_error():
+            raise ConnectionError("error")
+
+        def mock_execute_command(*_args, **_kwargs):
+            if _args[0] == "CLUSTER SLOTS":
+                return cluster_slots
+            else:
+                raise Exception("Failed to mock cluster slots")
+
+        # Mock connection error for the current primary
+        mock_node_resp_func(primary, raise_connection_error)
+        primary.redis_connection.set_retry(Retry(NoBackoff(), 1))
+
+        # Mock the cluster slots response for all other nodes
+        redis_mock_node = Mock()
+        redis_mock_node.execute_command.side_effect = mock_execute_command
+        # Mock response value for all other commands
+        redis_mock_node.parse_response.return_value = "MOCK_OK"
+        for node in r.get_nodes():
+            if node.port != primary.port:
+                node.redis_connection = redis_mock_node
+
+        assert r.get(key) == "MOCK_OK"
+        new_primary = r.get_node_from_key(key, replica=False)
+        assert new_primary.host == replica_host
+        assert new_primary.port == replica_port
+        assert r.get_node(primary.host, primary.port).server_type == REPLICA
+
     def test_moved_redirection(self, request):
         """
         Test that the client handles MOVED response.
@@ -693,44 +747,57 @@ def moved_redirect_effect(connection, *args, **options):
                     cur_node = r.get_node(node_name=node_name)
                     assert conn == r.get_redis_connection(cur_node)
 
-    def test_cluster_retry_object(self) -> None:
+    def test_cluster_get_set_retry_object(self, request):
+        retry = Retry(NoBackoff(), 2)
+        r = _get_client(Redis, request, retry=retry)
+        assert r.get_retry()._retries == retry._retries
+        assert isinstance(r.get_retry()._backoff, NoBackoff)
+        for node in r.get_nodes():
+            assert node.redis_connection.get_retry()._retries == retry._retries
+            assert isinstance(node.redis_connection.get_retry()._backoff, NoBackoff)
+        # Change retry policy
+        new_retry = Retry(ExponentialBackoff(), 3)
+        r.set_retry(new_retry)
+        assert r.get_retry()._retries == new_retry._retries
+        assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
+        for node in r.get_nodes():
+            assert node.redis_connection.get_retry()._retries == new_retry._retries
+            assert isinstance(
+                node.redis_connection.get_retry()._backoff, ExponentialBackoff
+            )
+
+    def test_cluster_retry_object(self, r) -> None:
         # Test default retry
-        rc_default = RedisCluster("127.0.0.1", 16379)
-        retry = rc_default.get_connection_kwargs().get("retry")
+        retry = r.get_connection_kwargs().get("retry")
         assert isinstance(retry, Retry)
         assert retry._retries == 3
         assert isinstance(retry._backoff, type(get_default_backoff()))
-        assert rc_default.get_node("127.0.0.1", 16379).get_connection_kwargs().get(
-            "retry"
-        ) == rc_default.get_node("127.0.0.1", 16380).get_connection_kwargs().get(
-            "retry"
-        )
+        node1 = r.get_node("127.0.0.1", 16379).redis_connection
+        node2 = r.get_node("127.0.0.1", 16380).redis_connection
+        assert node1.get_retry()._retries == node2.get_retry()._retries
 
         # Test custom retry
         retry = Retry(ExponentialBackoff(10, 5), 5)
         rc_custom_retry = RedisCluster("127.0.0.1", 16379, retry=retry)
         assert (
             rc_custom_retry.get_node("127.0.0.1", 16379)
-            .get_connection_kwargs()
-            .get("retry")
-            == retry
+            .redis_connection.get_retry()
+            ._retries
+            == retry._retries
         )
 
         # Test no connection retries
         rc_no_retries = RedisCluster(
             "127.0.0.1", 16379, connection_error_retry_attempts=0
         )
         assert (
-            rc_no_retries.get_node("127.0.0.1", 16379)
-            .get_connection_kwargs()
-            .get("retry")
+            rc_no_retries.get_node("127.0.0.1", 16379).redis_connection.get_retry()
             is None
         )
         rc_no_retries = RedisCluster("127.0.0.1", 16379, retry=Retry(NoBackoff(), 0))
         assert (
             rc_no_retries.get_node("127.0.0.1", 16379)
-            .get_connection_kwargs()
-            .get("retry")
+            .redis_connection.get_retry()
             ._retries
             == 0
         )

tests/test_retry.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from redis.backoff import NoBackoff
+from redis.backoff import ExponentialBackoff, NoBackoff
 from redis.client import Redis
 from redis.connection import Connection, UnixDomainSocketConnection
 from redis.exceptions import (
@@ -203,3 +203,13 @@ def test_client_retry_on_timeout(self, request):
                     r.get("foo")
                 finally:
                     assert parse_response.call_count == retries + 1
+
+    def test_get_set_retry_object(self, request):
+        retry = Retry(NoBackoff(), 2)
+        r = _get_client(Redis, request, retry_on_timeout=True, retry=retry)
+        assert r.get_retry()._retries == retry._retries
+        assert isinstance(r.get_retry()._backoff, NoBackoff)
+        new_retry_policy = Retry(ExponentialBackoff(), 3)
+        r.set_retry(new_retry_policy)
+        assert r.get_retry()._retries == new_retry_policy._retries
+        assert isinstance(r.get_retry()._backoff, ExponentialBackoff)