Release resources in test

Author: kristjanvalur

tests/test_asyncio/test_cwe_404.py

@@ -45,7 +45,7 @@ def __init__(self, addr, redis_addr, delay: float):
 
     async def start(self):
         # test that we can connect to redis
-        with async_timeout(2):
+        async with async_timeout(2):
             redis_reader, redis_writer = await asyncio.open_connection(*self.redis_addr)
         redis_writer.close()
         self.server = await asyncio.start_server(self.handle, *self.addr)
@@ -67,15 +67,23 @@ def override(self, delay: float = 0.0):
     async def handle(self, reader, writer):
         # establish connection to redis
         redis_reader, redis_writer = await asyncio.open_connection(*self.redis_addr)
-        pipe1 = asyncio.create_task(
-            pipe(reader, redis_writer, self, "to redis:", self.send_event)
-        )
-        pipe2 = asyncio.create_task(pipe(redis_reader, writer, self, "from redis:"))
-        await asyncio.gather(pipe1, pipe2)
+        try:
+            pipe1 = asyncio.create_task(
+                pipe(reader, redis_writer, self, "to redis:", self.send_event)
+            )
+            pipe2 = asyncio.create_task(pipe(redis_reader, writer, self, "from redis:"))
+            await asyncio.gather(pipe1, pipe2)
+        finally:
+            redis_writer.close()
+            redis_reader.close()
 
     async def stop(self):
         # clean up enough so that we can reuse the looper
         self.ROUTINE.cancel()
+        try:
+            await self.ROUTINE
+        except asyncio.CancelledError:
+            pass
         loop = self.server.get_loop()
         await loop.shutdown_asyncgens()
 
@@ -181,7 +189,7 @@ async def test_cluster(request, redis_addr):
         remap.append({"from_port": port, "to_port": remapped})
         forward_addr = redis_addr[0], port
         proxy = DelayProxy(
-            addr=("127.0.0.1", remapped), redis_addr=forward_addr, delay=delay
+            addr=("127.0.0.1", remapped), redis_addr=forward_addr, delay=0
         )
         proxies.append(proxy)
 
@@ -198,30 +206,29 @@ async def wait_for_send():
         )
 
     @contextlib.contextmanager
-    def override():
+    def override(delay: int = 0):
         with contextlib.ExitStack() as stack:
             for p in proxies:
-                stack.enter_context(p.override())
+                stack.enter_context(p.override(delay=delay))
             yield
 
-    with override():
-        r = RedisCluster.from_url(
-            f"redis://127.0.0.1:{remap_base}", host_port_remap=remap
-        )
+    with contextlib.closing(
+        RedisCluster.from_url(f"redis://127.0.0.1:{remap_base}", host_port_remap=remap)
+    ) as r:
         await r.initialize()
         await r.set("foo", "foo")
         await r.set("bar", "bar")
 
-    all_clear()
-    t = asyncio.create_task(r.get("foo"))
-    # cannot wait on the send event, we don't know which node will be used
-    await wait_for_send()
-    await asyncio.sleep(delay)
-    t.cancel()
-    with pytest.raises(asyncio.CancelledError):
-        await t
+        all_clear()
+        with override(delay=delay):
+            t = asyncio.create_task(r.get("foo"))
+            # cannot wait on the send event, we don't know which node will be used
+            await wait_for_send()
+            await asyncio.sleep(delay)
+            t.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await t
 
-    with override():
         # try a number of requests to excercise all the connections
         async def doit():
             assert await r.get("bar") == b"bar"