make mock class have same behavior as actual class

Author: agnesnatasya

redis/sentinel.py

@@ -58,6 +58,9 @@ def _connect_retry(self):
         if self._is_address_fixed:
             self.connect_to((self.host, self.port))
             return
+        self._connect_to_sentinel()
+    
+    def _connect_to_sentinel(self):
         # If same_server is False, connnect to master in master mode
         # and rotate to the next slave in slave mode
         if self.connection_pool.is_master:

tests/test_asyncio/test_sentinel_managed_connection.py

@@ -43,41 +43,21 @@ async def mock_connect():
 
 
 class SentinelManagedConnectionMock(SentinelManagedConnection):
-    async def connect_to_same_address(self) -> None:
-        pass
-
-    async def can_read_destructive(self) -> bool:
-        # Mock this function to always return False.
-        # Trrue means there's still data to be read and hence we can't reconnect
-        # to this connection yet
-        return False
-
-
-class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock):
-    async def connect(self) -> None:
+    async def _connect_to_sentinel(self) -> None:
         """
-        This simulates the behavior of connect when
+        This simulates the behavior of _connect_to_sentinel when
         :py:class:`~redis.SentinelConnectionPool`
         is in replica mode.
         It'll call rotate_slaves and connect to the next replica.
         """
-        import random
-        import time
-
-        self.host = f"host-{random.randint(0, 10)}"
-        self.port = time.time()
-
+        if self.connection_pool.is_master:
+            self.host, self.port = ("master", 1)
+        else:
+            import random
+            import time
 
-class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock):
-    async def connect_to(self, address: Tuple[str, int]) -> None:
-        """
-        This simulates the behavior of connect_to when
-        :py:class:`~redis.SentinelConnectionPool`
-        is in master mode.
-        It'll try to connect to master. In this mock class,
-        it'll just set the host and port without actually connecting.
-        """
-        self.host, self.port = address
+            self.host = f"host-{random.randint(0, 10)}"
+            self.port = time.time()
 
 
 @pytest.fixture()
@@ -90,7 +70,7 @@ def connection_pool_replica_mock() -> SentinelConnectionPool:
         "usasm",
         sentinel_manager,
         is_master=False,
-        connection_class=SentinelManagedConnectionMockForReplicaMode,
+        connection_class=SentinelManagedConnectionMock,
     )
     return connection_pool
 
@@ -105,7 +85,7 @@ def connection_pool_master_mock() -> SentinelConnectionPool:
         "usasm",
         sentinel_manager,
         is_master=True,
-        connection_class=SentinelManagedConnectionMockForMasterMode,
+        connection_class=SentinelManagedConnectionMock,
     )
     return connection_pool
 

tests/test_sentinel_managed_connection.py

@@ -7,32 +7,24 @@
 
 
 class SentinelManagedConnectionMock(SentinelManagedConnection):
-    def connect_to_same_address(self) -> None:
-        pass
-
-    def can_read_destructive(self) -> bool:
-        # Mock this function to always return False.
-        # Trrue means there's still data to be read and hence we can't reconnect
-        # to this connection yet
-        return False
-
-
-class SentinelManagedConnectionMockForReplicaMode(SentinelManagedConnectionMock):
-    def connect(self) -> None:
+    def _connect_to_sentinel(self) -> None:
         """
-        This simulates the behavior of connect when
+        This simulates the behavior of _connect_to_sentinel when
         :py:class:`~redis.SentinelConnectionPool`
         is in replica mode.
         It'll call rotate_slaves and connect to the next replica.
         """
-        import random
-        import time
+        if self.connection_pool.is_master:
+            self.host, self.port = ("master", 1)
+        else:
+            import random
+            import time
 
-        self.host = f"host-{random.randint(0, 10)}"
-        self.port = time.time()
+            self.host = f"host-{random.randint(0, 10)}"
+            self.port = time.time()
 
 
-class SentinelManagedConnectionMockForMasterMode(SentinelManagedConnectionMock):
+class SentinelManagedConnectionMock(SentinelManagedConnectionMock):
     def connect_to(self, address: Tuple[str, int]) -> None:
         """
         This simulates the behavior of connect_to when
@@ -41,7 +33,6 @@ def connect_to(self, address: Tuple[str, int]) -> None:
         It'll try to connect to master. In this mock class,
         it'll just set the host and port without actually connecting.
         """
-        self.host, self.port = address
 
 
 @pytest.fixture()
@@ -54,7 +45,7 @@ def connection_pool_replica_mock() -> SentinelConnectionPool:
         "usasm",
         sentinel_manager,
         is_master=False,
-        connection_class=SentinelManagedConnectionMockForReplicaMode,
+        connection_class=SentinelManagedConnectionMock,
     )
     return connection_pool
 
@@ -69,7 +60,7 @@ def connection_pool_master_mock() -> SentinelConnectionPool:
         "usasm",
         sentinel_manager,
         is_master=True,
-        connection_class=SentinelManagedConnectionMockForMasterMode,
+        connection_class=SentinelManagedConnectionMock,
     )
     return connection_pool