Add aclose method to asyncio.RedisCluster

Author: kristjanvalur

redis/asyncio/cluster.py

@@ -395,27 +395,31 @@ async def initialize(self) -> "RedisCluster":
                         )
                         self._initialize = False
                     except BaseException:
-                        await self.nodes_manager.close()
-                        await self.nodes_manager.close("startup_nodes")
+                        await self.nodes_manager.aclose()
+                        await self.nodes_manager.aclose("startup_nodes")
                         raise
         return self
 
-    async def close(self) -> None:
+    async def aclose(self) -> None:
         """Close all connections & client if initialized."""
         if not self._initialize:
             if not self._lock:
                 self._lock = asyncio.Lock()
             async with self._lock:
                 if not self._initialize:
                     self._initialize = True
-                    await self.nodes_manager.close()
-                    await self.nodes_manager.close("startup_nodes")
+                    await self.nodes_manager.aclose()
+                    await self.nodes_manager.aclose("startup_nodes")
+
+    async def close(self) -> None:
+        """alias for aclose() for backwards compatibility"""
+        await self.aclose()
 
     async def __aenter__(self) -> "RedisCluster":
         return await self.initialize()
 
     async def __aexit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
-        await self.close()
+        await self.aclose()
 
     def __await__(self) -> Generator[Any, None, "RedisCluster"]:
         return self.initialize().__await__()
@@ -767,13 +771,13 @@ async def _execute_command(
                 self.nodes_manager.startup_nodes.pop(target_node.name, None)
                 # Hard force of reinitialize of the node/slots setup
                 # and try again with the new setup
-                await self.close()
+                await self.aclose()
                 raise
             except ClusterDownError:
                 # ClusterDownError can occur during a failover and to get
                 # self-healed, we will try to reinitialize the cluster layout
                 # and retry executing the command
-                await self.close()
+                await self.aclose()
                 await asyncio.sleep(0.25)
                 raise
             except MovedError as e:
@@ -790,7 +794,7 @@ async def _execute_command(
                     self.reinitialize_steps
                     and self.reinitialize_counter % self.reinitialize_steps == 0
                 ):
-                    await self.close()
+                    await self.aclose()
                     # Reset the counter
                     self.reinitialize_counter = 0
                 else:
@@ -1323,7 +1327,7 @@ async def initialize(self) -> None:
         # If initialize was called after a MovedError, clear it
         self._moved_exception = None
 
-    async def close(self, attr: str = "nodes_cache") -> None:
+    async def aclose(self, attr: str = "nodes_cache") -> None:
         self.default_node = None
         await asyncio.gather(
             *(
@@ -1471,7 +1475,7 @@ async def execute(
                     if type(e) in self.__class__.ERRORS_ALLOW_RETRY:
                         # Try again with the new cluster setup.
                         exception = e
-                        await self._client.close()
+                        await self._client.aclose()
                         await asyncio.sleep(0.25)
                     else:
                         # All other errors should be raised.

tests/test_asyncio/test_cluster.py

@@ -38,7 +38,7 @@
 )
 
 from ..ssl_utils import get_ssl_filename
-from .compat import mock
+from .compat import aclosing, mock
 
 pytestmark = pytest.mark.onlycluster
 
@@ -270,7 +270,39 @@ async def test_host_port_startup_node(self) -> None:
         cluster = await get_mocked_redis_client(host=default_host, port=default_port)
         assert cluster.get_node(host=default_host, port=default_port) is not None
 
-        await cluster.close()
+        await cluster.aclose()
+
+    async def test_aclosing(self) -> None:
+        cluster = await get_mocked_redis_client(host=default_host, port=default_port)
+        called = 1
+
+        async def mock_aclose(_):
+            nonlocal called
+            called += 1
+
+        with mock.patch.object(cluster, "aclose", mock_aclose):
+            async with aclosing(cluster):
+                pass
+            assert called == 1
+        await cluster.aclose()
+
+    async def test_close_is_aclose(self) -> None:
+        """
+        Test that it is possible to use host & port arguments as startup node
+        args
+        """
+        cluster = await get_mocked_redis_client(host=default_host, port=default_port)
+
+        called = 1
+
+        async def mock_aclose(_):
+            nonlocal called
+            called += 1
+
+        with mock.patch.object(cluster, "aclose", mock_aclose):
+            await cluster.close()
+            assert called == 1
+        await cluster.aclose()
 
     async def test_startup_nodes(self) -> None:
         """
@@ -289,7 +321,7 @@ async def test_startup_nodes(self) -> None:
             and cluster.get_node(host=default_host, port=port_2) is not None
         )
 
-        await cluster.close()
+        await cluster.aclose()
 
         startup_node = ClusterNode("127.0.0.1", 16379)
         async with RedisCluster(startup_nodes=[startup_node], client_name="test") as rc:
@@ -417,7 +449,7 @@ async def read_response_mocked(*args: Any, **kwargs: Any) -> None:
                     )
                 )
 
-        await rc.close()
+        await rc.aclose()
 
     async def test_execute_command_errors(self, r: RedisCluster) -> None:
         """
@@ -461,7 +493,7 @@ async def test_execute_command_node_flag_replicas(self, r: RedisCluster) -> None
             conn = primary._free.pop()
             assert conn.read_response.called is not True
 
-        await r.close()
+        await r.aclose()
 
     async def test_execute_command_node_flag_all_nodes(self, r: RedisCluster) -> None:
         """
@@ -690,7 +722,7 @@ def execute_command_mock_third(self, *args, **options):
                 await read_cluster.get("foo")
                 mocks["send_command"].assert_has_calls([mock.call("READONLY")])
 
-                await read_cluster.close()
+                await read_cluster.aclose()
 
     async def test_keyslot(self, r: RedisCluster) -> None:
         """
@@ -762,7 +794,7 @@ def raise_error(target_node, *args, **kwargs):
                 await rc.get("bar")
                 assert execute_command.failed_calls == rc.cluster_error_retry_attempts
 
-            await rc.close()
+            await rc.aclose()
 
     async def test_set_default_node_success(self, r: RedisCluster) -> None:
         """
@@ -843,7 +875,7 @@ async def test_can_run_concurrent_commands(self, request: FixtureRequest) -> Non
                 *(rc.echo("i", target_nodes=RedisCluster.ALL_NODES) for i in range(100))
             )
         )
-        await rc.close()
+        await rc.aclose()
 
     def test_replace_cluster_node(self, r: RedisCluster) -> None:
         prev_default_node = r.get_default_node()
@@ -901,7 +933,7 @@ def address_remap(address):
                 assert await r.set("byte_string", b"giraffe")
                 assert await r.get("byte_string") == b"giraffe"
             finally:
-                await r.close()
+                await r.aclose()
         finally:
             await asyncio.gather(*[p.aclose() for p in proxies])
 
@@ -1002,7 +1034,7 @@ async def test_initialize_before_execute_multi_key_command(
         url = request.config.getoption("--redis-url")
         r = RedisCluster.from_url(url)
         assert 0 == await r.exists("a", "b", "c")
-        await r.close()
+        await r.aclose()
 
     @skip_if_redis_enterprise()
     async def test_cluster_myid(self, r: RedisCluster) -> None:
@@ -1065,7 +1097,7 @@ async def test_cluster_delslots(self) -> None:
         assert node0._free.pop().read_response.called
         assert node1._free.pop().read_response.called
 
-        await r.close()
+        await r.aclose()
 
     @skip_if_server_version_lt("7.0.0")
     @skip_if_redis_enterprise()
@@ -1076,7 +1108,7 @@ async def test_cluster_delslotsrange(self):
         await r.cluster_addslots(node, 1, 2, 3, 4, 5)
         assert await r.cluster_delslotsrange(1, 5)
         assert node._free.pop().read_response.called
-        await r.close()
+        await r.aclose()
 
     @skip_if_redis_enterprise()
     async def test_cluster_failover(self, r: RedisCluster) -> None:
@@ -1286,7 +1318,7 @@ async def test_readonly(self) -> None:
         for replica in r.get_replicas():
             assert replica._free.pop().read_response.called
 
-        await r.close()
+        await r.aclose()
 
     @skip_if_redis_enterprise()
     async def test_readwrite(self) -> None:
@@ -1299,7 +1331,7 @@ async def test_readwrite(self) -> None:
         for replica in r.get_replicas():
             assert replica._free.pop().read_response.called
 
-        await r.close()
+        await r.aclose()
 
     @skip_if_redis_enterprise()
     async def test_bgsave(self, r: RedisCluster) -> None:
@@ -1524,7 +1556,7 @@ async def test_client_kill(
         ]
         assert len(clients) == 1
         assert clients[0].get("name") == "redis-py-c1"
-        await r2.close()
+        await r2.aclose()
 
     @skip_if_server_version_lt("2.6.0")
     async def test_cluster_bitop_not_empty_string(self, r: RedisCluster) -> None:
@@ -2302,7 +2334,7 @@ async def test_acl_log(
 
         await r.acl_deluser(username, target_nodes="primaries")
 
-        await user_client.close()
+        await user_client.aclose()
 
 
 class TestNodesManager:
@@ -2359,7 +2391,7 @@ async def test_init_slots_cache_not_all_slots_covered(self) -> None:
                 cluster_slots=cluster_slots,
                 require_full_coverage=True,
             )
-            await rc.close()
+            await rc.aclose()
         assert str(ex.value).startswith(
             "All slots are not covered after query all startup_nodes."
         )
@@ -2385,7 +2417,7 @@ async def test_init_slots_cache_not_require_full_coverage_success(self) -> None:
 
         assert 5460 not in rc.nodes_manager.slots_cache
 
-        await rc.close()
+        await rc.aclose()
 
     async def test_init_slots_cache(self) -> None:
         """
@@ -2416,7 +2448,7 @@ async def test_init_slots_cache(self) -> None:
 
         assert len(n_manager.nodes_cache) == 6
 
-        await rc.close()
+        await rc.aclose()
 
     async def test_init_slots_cache_cluster_mode_disabled(self) -> None:
         """
@@ -2427,7 +2459,7 @@ async def test_init_slots_cache_cluster_mode_disabled(self) -> None:
             rc = await get_mocked_redis_client(
                 host=default_host, port=default_port, cluster_enabled=False
             )
-            await rc.close()
+            await rc.aclose()
         assert "Cluster mode is not enabled on this node" in str(e.value)
 
     async def test_empty_startup_nodes(self) -> None:
@@ -2514,7 +2546,7 @@ async def test_cluster_one_instance(self) -> None:
         for i in range(0, REDIS_CLUSTER_HASH_SLOTS):
             assert n.slots_cache[i] == [n_node]
 
-        await rc.close()
+        await rc.aclose()
 
     async def test_init_with_down_node(self) -> None:
         """