Close various objects created during asyncio tests (#3005)

* Close various objects created during asyncio tests

* Fix resource leake in test_cwe_404.py
Need to wait for individual handler tasks when shutting down server.

Author: kristjanvalur

tests/test_asyncio/test_commands.py

@@ -370,10 +370,12 @@ async def test_client_setinfo(self, r: redis.Redis):
         info = await r2.client_info()
         assert info["lib-name"] == "test2"
         assert info["lib-ver"] == "1234"
+        await r2.aclose()
         r3 = redis.asyncio.Redis(lib_name=None, lib_version=None)
         info = await r3.client_info()
         assert info["lib-name"] == ""
         assert info["lib-ver"] == ""
+        await r3.aclose()
 
     @skip_if_server_version_lt("2.6.9")
     @pytest.mark.onlynoncluster

tests/test_asyncio/test_connect.py

@@ -73,6 +73,8 @@ async def _handler(reader, writer):
         try:
             return await _redis_request_handler(reader, writer, stop_event)
         finally:
+            writer.close()
+            await writer.wait_closed()
             finished.set()
 
     if isinstance(server_address, str):

tests/test_asyncio/test_connection.py

@@ -85,6 +85,8 @@ async def get_conn(_):
 
     assert init_call_count == 1
     assert command_call_count == 2
+    r.connection = None  # it was a Mock
+    await r.aclose()
 
 
 @skip_if_server_version_lt("4.0.0")
@@ -143,6 +145,7 @@ async def mock_connect():
     conn._connect.side_effect = mock_connect
     await conn.connect()
     assert conn._connect.call_count == 3
+    await conn.disconnect()
 
 
 async def test_connect_without_retry_on_os_error():
@@ -194,6 +197,7 @@ async def test_connection_parse_response_resume(r: redis.Redis):
         pytest.fail("didn't receive a response")
     assert response
     assert i > 0
+    await conn.disconnect()
 
 
 @pytest.mark.onlynoncluster

tests/test_asyncio/test_cwe_404.py

@@ -15,6 +15,8 @@ def __init__(self, addr, redis_addr, delay: float = 0.0):
         self.send_event = asyncio.Event()
         self.server = None
         self.task = None
+        self.cond = asyncio.Condition()
+        self.running = 0
 
     async def __aenter__(self):
         await self.start()
@@ -63,10 +65,10 @@ async def stop(self):
         except asyncio.CancelledError:
             pass
         await self.server.wait_closed()
-        # do we need to close individual connections too?
-        # prudently close all async generators
-        loop = self.server.get_loop()
-        await loop.shutdown_asyncgens()
+        # Server does not wait for all spawned tasks.  We must do that also to ensure
+        # that all sockets are closed.
+        async with self.cond:
+            await self.cond.wait_for(lambda: self.running == 0)
 
     async def pipe(
         self,
@@ -75,6 +77,7 @@ async def pipe(
         name="",
         event: asyncio.Event = None,
     ):
+        self.running += 1
         try:
             while True:
                 data = await reader.read(1000)
@@ -94,6 +97,10 @@ async def pipe(
                 # ignore errors on close pertaining to no event loop. Don't want
                 # to clutter the test output with errors if being garbage collected
                 pass
+            async with self.cond:
+                self.running -= 1
+                if self.running == 0:
+                    self.cond.notify_all()
 
 
 @pytest.mark.onlynoncluster

tests/test_asyncio/test_retry.py

@@ -131,5 +131,8 @@ async def test_get_set_retry_object(self, request):
         assert r.get_retry()._retries == new_retry_policy._retries
         assert isinstance(r.get_retry()._backoff, ExponentialBackoff)
         assert exiting_conn.retry._retries == new_retry_policy._retries
+        await r.connection_pool.release(exiting_conn)
         new_conn = await r.connection_pool.get_connection("_")
         assert new_conn.retry._retries == new_retry_policy._retries
+        await r.connection_pool.release(new_conn)
+        await r.aclose()

tests/test_asyncio/test_sentinel.py

@@ -183,30 +183,30 @@ async def test_discover_slaves(cluster, sentinel):
 
 @pytest.mark.onlynoncluster
 async def test_master_for(cluster, sentinel, master_ip):
-    master = sentinel.master_for("mymaster", db=9)
-    assert await master.ping()
-    assert master.connection_pool.master_address == (master_ip, 6379)
+    async with sentinel.master_for("mymaster", db=9) as master:
+        assert await master.ping()
+        assert master.connection_pool.master_address == (master_ip, 6379)
 
     # Use internal connection check
-    master = sentinel.master_for("mymaster", db=9, check_connection=True)
-    assert await master.ping()
+    async with sentinel.master_for("mymaster", db=9, check_connection=True) as master:
+        assert await master.ping()
 
 
 @pytest.mark.onlynoncluster
 async def test_slave_for(cluster, sentinel):
     cluster.slaves = [
         {"ip": "127.0.0.1", "port": 6379, "is_odown": False, "is_sdown": False}
     ]
-    slave = sentinel.slave_for("mymaster", db=9)
-    assert await slave.ping()
+    async with sentinel.slave_for("mymaster", db=9) as slave:
+        assert await slave.ping()
 
 
 @pytest.mark.onlynoncluster
 async def test_slave_for_slave_not_found_error(cluster, sentinel):
     cluster.master["is_odown"] = True
-    slave = sentinel.slave_for("mymaster", db=9)
-    with pytest.raises(SlaveNotFoundError):
-        await slave.ping()
+    async with sentinel.slave_for("mymaster", db=9) as slave:
+        with pytest.raises(SlaveNotFoundError):
+            await slave.ping()
 
 
 @pytest.mark.onlynoncluster
@@ -260,7 +260,7 @@ async def mock_disconnect():
         calls += 1
 
     with mock.patch.object(pool, "disconnect", mock_disconnect):
-        await client.close()
+        await client.aclose()
 
     assert calls == 1
     await pool.disconnect()

tests/test_asyncio/test_sentinel_managed_connection.py

@@ -34,3 +34,4 @@ async def mock_connect():
     conn._connect.side_effect = mock_connect
     await conn.connect()
     assert conn._connect.call_count == 3
+    await conn.disconnect()