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

Author: zach-iee

tests/test_cluster.py

@@ -1,5 +1,6 @@
 import binascii
 import datetime
+import random
 import select
 import socket
 import socketserver
@@ -8,11 +9,12 @@
 import warnings
 from queue import LifoQueue, Queue
 from time import sleep
+from unittest import mock
 from unittest.mock import DEFAULT, Mock, call, patch
 
 import pytest
 
-from redis import Redis
+from redis import Redis, crc
 from redis.backoff import (
     ConstantBackoff,
     ExponentialBackoff,
@@ -26,7 +28,7 @@
     ClusterNode,
     NodesManager,
     RedisCluster,
-    get_node_name,
+    get_node_name, LoadBalancer,
 )
 from redis.commands import CommandsParser
 from redis.connection import BlockingConnectionPool, Connection, ConnectionPool
@@ -2654,6 +2656,25 @@ def test_allow_custom_queue_class(self, queue_class):
         for node in rc.nodes_manager.nodes_cache.values():
             assert node.redis_connection.connection_pool.queue_class == queue_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://[email protected]: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: