From 05bf98c3ad9b243921b01614d4ff20b091112d93 Mon Sep 17 00:00:00 2001
From: "zach.lee" <zach.lee@sendbird.com>
Date: Wed, 28 Jun 2023 15:18:45 +0900
Subject: [PATCH] [GROW-2938] additional test to #1 (#2)

* [GROW-2938] add a test for invalid index returned from load balancer

* [GROW-2938] add a test for pipeline additional backoff

(cherry picked from commit 308c82b981742db06df0028a2d7e146ddc4af315)
---
 tests/test_cluster.py | 63 ++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index 9d729b3d57..3bfdedaca9 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -1,8 +1,10 @@
 import binascii
 import datetime
+import random
 import uuid
 import warnings
 from time import sleep
+from unittest import mock
 from unittest.mock import DEFAULT, Mock, call, patch
 
 import pytest
@@ -20,6 +22,7 @@
     REDIS_CLUSTER_HASH_SLOTS,
     REPLICA,
     ClusterNode,
+    LoadBalancer,
     NodesManager,
     RedisCluster,
     get_node_name,
@@ -810,7 +813,7 @@ def raise_error(target_node, *args, **kwargs):
                 rc = get_mocked_redis_client(
                     host=default_host,
                     port=default_port,
-                    retry=Retry(ConstantBackoff(1), 3),
+                    retry=Retry(ConstantBackoff(1), 10),
                 )
 
                 with pytest.raises(error):
@@ -2519,6 +2522,37 @@ def test_connection_pool_class(self, connection_pool_class):
                 node.redis_connection.connection_pool, connection_pool_class
             )
 
+    @pytest.mark.parametrize("invalid_index", [-10, 10])
+    def test_return_primary_if_invalid_node_index_is_returned(self, invalid_index):
+        rc = get_mocked_redis_client(
+            url="redis://my@DNS.com:7000",
+            cluster_slots=default_cluster_slots,
+        )
+        random_slot = random.randint(
+            default_cluster_slots[0][0], default_cluster_slots[0][1]
+        )
+
+        ports = set()
+        for _ in range(0, 10):
+            ports.add(
+                rc.nodes_manager.get_node_from_slot(
+                    random_slot, read_from_replicas=True
+                ).port
+            )
+        assert ports == {default_port, 7003}
+
+        ports = set()
+        with mock.patch.object(
+            LoadBalancer, "get_server_index", return_value=invalid_index
+        ):
+            for _ in range(0, 10):
+                ports.add(
+                    rc.nodes_manager.get_node_from_slot(
+                        random_slot, read_from_replicas=True
+                    ).port
+                )
+        assert ports == {default_port}
+
 
 @pytest.mark.onlycluster
 class TestClusterPubSubObject:
@@ -2930,6 +2964,33 @@ def test_empty_stack(self, r):
         result = p.execute()
         assert result == []
 
+    @pytest.mark.parametrize("error", [ConnectionError, TimeoutError])
+    def test_additional_backoff_cluster_pipeline(self, r, error):
+        with patch.object(ConstantBackoff, "compute") as compute:
+
+            def _compute(target_node, *args, **kwargs):
+                return 1
+
+            compute.side_effect = _compute
+            with patch("redis.cluster.get_connection") as get_connection:
+
+                def raise_error(target_node, *args, **kwargs):
+                    get_connection.failed_calls += 1
+                    raise error("mocked error")
+
+                get_connection.side_effect = raise_error
+
+                r.set_retry(Retry(ConstantBackoff(1), 10))
+                pipeline = r.pipeline()
+
+                with pytest.raises(error):
+                    pipeline.get("bar")
+                    pipeline.get("bar")
+                    pipeline.execute()
+                # cluster pipeline does one more back off than a single Redis command
+                #   this is not required, but it's just how it's implemented as of now
+                assert compute.call_count == r.cluster_error_retry_attempts + 1
+
 
 @pytest.mark.onlycluster
 class TestReadOnlyPipeline: