fix tests

Author: kristjanvalur

tests/test_asyncio/test_cwe_404.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import sys
 
 import pytest
@@ -8,13 +9,19 @@
 
 
 async def pipe(
-    reader: asyncio.StreamReader, writer: asyncio.StreamWriter, delay: float, name=""
+    reader: asyncio.StreamReader,
+    writer: asyncio.StreamWriter,
+    proxy: "DelayProxy",
+    name="",
+    event: asyncio.Event = None,
 ):
     while True:
         data = await reader.read(1000)
         if not data:
             break
-        await asyncio.sleep(delay)
+        if event:
+            event.set()
+        await asyncio.sleep(proxy.delay)
         writer.write(data)
         await writer.drain()
 
@@ -24,18 +31,32 @@ def __init__(self, addr, redis_addr, delay: float):
         self.addr = addr
         self.redis_addr = redis_addr
         self.delay = delay
+        self.send_event = asyncio.Event()
 
     async def start(self):
         self.server = await asyncio.start_server(self.handle, *self.addr)
         self.ROUTINE = asyncio.create_task(self.server.serve_forever())
 
+    @contextlib.contextmanager
+    def override(self, delay: float = 0.0):
+        """
+        Allow to override the delay for parts of tests which aren't time dependent,
+        to speed up execution.
+        """
+        old = self.delay
+        self.delay = delay
+        try:
+            yield
+        finally:
+            self.delay = old
+
     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.delay, "to redis:"))
-        pipe2 = asyncio.create_task(
-            pipe(redis_reader, writer, self.delay, "from redis:")
+        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)
 
     async def stop(self):
@@ -60,23 +81,26 @@ async def test_standalone(delay):
         # note that we connect to proxy, rather than to Redis directly
         async with Redis(host="localhost", port=5380, single_connection_client=b) as r:
 
-            await r.set("foo", "foo")
-            await r.set("bar", "bar")
+            with dp.override():
+                await r.set("foo", "foo")
+                await r.set("bar", "bar")
 
+            dp.send_event.clear()
             t = asyncio.create_task(r.get("foo"))
-            await asyncio.sleep(delay)
+            # wait until the task has sent, and then some, to make sure it has settled on
+            # reading.
+            await dp.send_event.wait()
+            await asyncio.sleep(0.05)
             t.cancel()
-            try:
+            with pytest.raises(asyncio.CancelledError):
                 await t
-                sys.stderr.write("try again, we did not cancel the task in time\n")
-            except asyncio.CancelledError:
-                sys.stderr.write(
-                    "canceled task, connection is left open with unread response\n"
-                )
 
-            assert await r.get("bar") == b"bar"
-            assert await r.ping()
-            assert await r.get("foo") == b"foo"
+            # make sure that our previous request, cancelled while waiting for a repsponse,
+            # didn't leave the connection in a bad state
+            with dp.override():
+                assert await r.get("bar") == b"bar"
+                assert await r.ping()
+                assert await r.get("foo") == b"foo"
 
     await dp.stop()
 
@@ -90,8 +114,9 @@ async def test_standalone_pipeline(delay):
     await dp.start()
     for b in [True, False]:
         async with Redis(host="localhost", port=5380, single_connection_client=b) as r:
-            await r.set("foo", "foo")
-            await r.set("bar", "bar")
+            with dp.override():
+                await r.set("foo", "foo")
+                await r.set("bar", "bar")
 
             pipe = r.pipeline()
 
@@ -100,23 +125,32 @@ async def test_standalone_pipeline(delay):
             pipe2.ping()
             pipe2.get("foo")
 
+            dp.send_event.clear()
             t = asyncio.create_task(pipe.get("foo").execute())
-            await asyncio.sleep(delay)
+            # wait until task has settled on the read
+            await dp.send_event.wait()
+            await asyncio.sleep(0.05)
             t.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await t
 
-            pipe.get("bar")
-            pipe.ping()
-            pipe.get("foo")
-            pipe.reset()
+            # we have now cancelled the pieline in the middle of a request, make sure
+            # that the connection is still usable
+            with dp.override():
+                pipe.get("bar")
+                pipe.ping()
+                pipe.get("foo")
+                await pipe.reset()
 
-            assert await pipe.execute() is None
+                # check that the pipeline is empty after reset
+                assert await pipe.execute() == []
 
-            # validating that the pipeline can be used as it could previously
-            pipe.get("bar")
-            pipe.ping()
-            pipe.get("foo")
-            assert await pipe.execute() == [b"bar", True, b"foo"]
-            assert await pipe2.execute() == [b"bar", True, b"foo"]
+                # validating that the pipeline can be used as it could previously
+                pipe.get("bar")
+                pipe.ping()
+                pipe.get("foo")
+                assert await pipe.execute() == [b"bar", True, b"foo"]
+                assert await pipe2.execute() == [b"bar", True, b"foo"]
 
     await dp.stop()
 
@@ -129,19 +163,21 @@ async def test_cluster(request):
 
     r = RedisCluster.from_url("redis://localhost:5381")
     await r.initialize()
-    await r.set("foo", "foo")
-    await r.set("bar", "bar")
+    with dp.override():
+        await r.set("foo", "foo")
+        await r.set("bar", "bar")
 
+    dp.send_event.clear()
     t = asyncio.create_task(r.get("foo"))
-    await asyncio.sleep(0.050)
+    await dp.send_event.wait()
+    await asyncio.sleep(0.05)
     t.cancel()
-    try:
+    with pytest.raises(asyncio.CancelledError):
         await t
-    except asyncio.CancelledError:
-        pytest.fail("connection is left open with unread response")
 
-    assert await r.get("bar") == b"bar"
-    assert await r.ping()
-    assert await r.get("foo") == b"foo"
+    with dp.override():
+        assert await r.get("bar") == b"bar"
+        assert await r.ping()
+        assert await r.get("foo") == b"foo"
 
     await dp.stop()