Clean up the asyncio and sync connection tests.

Author: aaugustin

tests/asyncio/test_connection.py

@@ -112,8 +112,8 @@ async def test_aexit(self):
             await self.assertNoFrameSent()
         await self.assertFrameSent(Frame(Opcode.CLOSE, b"\x03\xe8"))
 
-    async def test_exit_with_exception(self):
-        """__exit__ with an exception closes the connection with code 1011."""
+    async def test_aexit_with_exception(self):
+        """__aexit__ with an exception closes the connection with code 1011."""
         with self.assertRaises(RuntimeError):
             async with self.connection:
                 raise RuntimeError
@@ -248,10 +248,9 @@ async def test_recv_during_recv_streaming(self):
         )
 
     async def test_recv_cancellation_before_receiving(self):
-        """recv can be canceled before receiving a frame."""
+        """recv can be canceled before receiving a message."""
         recv_task = asyncio.create_task(self.connection.recv())
         await asyncio.sleep(0)  # let the event loop start recv_task
-
         recv_task.cancel()
         await asyncio.sleep(0)  # let the event loop cancel recv_task
 
@@ -260,25 +259,25 @@ async def test_recv_cancellation_before_receiving(self):
         self.assertEqual(await self.connection.recv(), "😀")
 
     async def test_recv_cancellation_while_receiving(self):
-        """recv cannot be canceled after receiving a frame."""
-        recv_task = asyncio.create_task(self.connection.recv())
-        await asyncio.sleep(0)  # let the event loop start recv_task
-
-        gate = asyncio.get_running_loop().create_future()
+        """recv can be canceled while receiving a fragmented message."""
+        gate = asyncio.Event()
 
         async def fragments():
             yield "⏳"
-            await gate
+            await gate.wait()
             yield "⌛️"
 
         asyncio.create_task(self.remote_connection.send(fragments()))
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)
 
+        recv_task = asyncio.create_task(self.connection.recv())
+        await asyncio.sleep(0)  # let the event loop start recv_task
         recv_task.cancel()
         await asyncio.sleep(0)  # let the event loop cancel recv_task
 
+        gate.set()
+
         # Running recv again receives the complete message.
-        gate.set_result(None)
         self.assertEqual(await self.connection.recv(), "⏳⌛️")
 
     # Test recv_streaming.
@@ -407,28 +406,31 @@ async def test_recv_streaming_cancellation_before_receiving(self):
 
     async def test_recv_streaming_cancellation_while_receiving(self):
         """recv_streaming cannot be canceled after receiving a frame."""
-        recv_streaming_task = asyncio.create_task(
-            alist(self.connection.recv_streaming())
-        )
-        await asyncio.sleep(0)  # let the event loop start recv_streaming_task
-
-        gate = asyncio.get_running_loop().create_future()
+        gate = asyncio.Event()
 
         async def fragments():
             yield "⏳"
-            await gate
+            await gate.wait()
             yield "⌛️"
 
         asyncio.create_task(self.remote_connection.send(fragments()))
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)
 
+        recv_streaming_task = asyncio.create_task(
+            alist(self.connection.recv_streaming())
+        )
+        await asyncio.sleep(0)  # let the event loop start recv_streaming_task
+        await asyncio.sleep(0)  # experimentally, two runs of the event loop
+        await asyncio.sleep(0)  # are needed to receive the first fragment
         recv_streaming_task.cancel()
         await asyncio.sleep(0)  # let the event loop cancel recv_streaming_task
 
-        gate.set_result(None)
+        gate.set()
+
         # Running recv_streaming again fails.
         with self.assertRaises(ConcurrencyError):
-            await alist(self.connection.recv_streaming())
+            async for _ in self.connection.recv_streaming():
+                self.fail("did not raise")
 
     # Test send.
 
@@ -556,23 +558,31 @@ async def test_send_connection_closed_error(self):
         with self.assertRaises(ConnectionClosedError):
             await self.connection.send("😀")
 
-    async def test_send_while_send_blocked(self):
+    async def test_send_during_send(self):
         """send waits for a previous call to send to complete."""
         # This test fails if the guard with send_in_progress is removed
-        # from send() in the case when message is an Iterable.
-        self.connection.pause_writing()
-        asyncio.create_task(self.connection.send(["⏳", "⌛️"]))
-        await asyncio.sleep(MS)
+        # from send() in the case when message is an AsyncIterable.
+        gate = asyncio.Event()
+
+        async def fragments():
+            yield "⏳"
+            await gate.wait()
+            yield "⌛️"
+
+        asyncio.create_task(self.connection.send(fragments()))
+        await asyncio.sleep(0)  # let the event loop start the task
         await self.assertFrameSent(
             Frame(Opcode.TEXT, "⏳".encode(), fin=False),
         )
 
         asyncio.create_task(self.connection.send("✅"))
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)  # let the event loop start the task
         await self.assertNoFrameSent()
 
-        self.connection.resume_writing()
-        await asyncio.sleep(MS)
+        gate.set()
+        await asyncio.sleep(0)  # run the event loop
+        await asyncio.sleep(0)  # three times in order
+        await asyncio.sleep(0)  # to send three frames
         await self.assertFramesSent(
             [
                 Frame(Opcode.CONT, "⌛️".encode(), fin=False),
@@ -581,28 +591,26 @@ async def test_send_while_send_blocked(self):
             ]
         )
 
-    async def test_send_while_send_async_blocked(self):
-        """send waits for a previous call to send to complete."""
+    async def test_send_while_send_blocked(self):
+        """send waits for a blocked call to send to complete."""
         # This test fails if the guard with send_in_progress is removed
-        # from send() in the case when message is an AsyncIterable.
+        # from send() in the case when message is an Iterable.
         self.connection.pause_writing()
 
-        async def fragments():
-            yield "⏳"
-            yield "⌛️"
-
-        asyncio.create_task(self.connection.send(fragments()))
-        await asyncio.sleep(MS)
+        asyncio.create_task(self.connection.send(["⏳", "⌛️"]))
+        await asyncio.sleep(0)  # let the event loop start the task
         await self.assertFrameSent(
             Frame(Opcode.TEXT, "⏳".encode(), fin=False),
         )
 
         asyncio.create_task(self.connection.send("✅"))
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)  # let the event loop start the task
         await self.assertNoFrameSent()
 
         self.connection.resume_writing()
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)  # run the event loop
+        await asyncio.sleep(0)  # three times in order
+        await asyncio.sleep(0)  # to send three frames
         await self.assertFramesSent(
             [
                 Frame(Opcode.CONT, "⌛️".encode(), fin=False),
@@ -611,29 +619,30 @@ async def fragments():
             ]
         )
 
-    async def test_send_during_send_async(self):
-        """send waits for a previous call to send to complete."""
+    async def test_send_while_send_async_blocked(self):
+        """send waits for a blocked call to send to complete."""
         # This test fails if the guard with send_in_progress is removed
         # from send() in the case when message is an AsyncIterable.
-        gate = asyncio.get_running_loop().create_future()
+        self.connection.pause_writing()
 
         async def fragments():
             yield "⏳"
-            await gate
             yield "⌛️"
 
         asyncio.create_task(self.connection.send(fragments()))
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)  # let the event loop start the task
         await self.assertFrameSent(
             Frame(Opcode.TEXT, "⏳".encode(), fin=False),
         )
 
         asyncio.create_task(self.connection.send("✅"))
-        await asyncio.sleep(MS)
+        await asyncio.sleep(0)  # let the event loop start the task
         await self.assertNoFrameSent()
 
-        gate.set_result(None)
-        await asyncio.sleep(MS)
+        self.connection.resume_writing()
+        await asyncio.sleep(0)  # run the event loop
+        await asyncio.sleep(0)  # three times in order
+        await asyncio.sleep(0)  # to send three frames
         await self.assertFramesSent(
             [
                 Frame(Opcode.CONT, "⌛️".encode(), fin=False),
@@ -837,13 +846,9 @@ async def test_close_preserves_queued_messages(self):
         await self.connection.close()
 
         self.assertEqual(await self.connection.recv(), "😀")
-        with self.assertRaises(ConnectionClosedOK) as raised:
+        with self.assertRaises(ConnectionClosedOK):
             await self.connection.recv()
 
-        exc = raised.exception
-        self.assertEqual(str(exc), "sent 1000 (OK); then received 1000 (OK)")
-        self.assertIsNone(exc.__cause__)
-
     async def test_close_idempotency(self):
         """close does nothing if the connection is already closed."""
         await self.connection.close()
@@ -854,35 +859,42 @@ async def test_close_idempotency(self):
 
     async def test_close_during_recv(self):
         """close aborts recv when called concurrently with recv."""
-        recv_task = asyncio.create_task(self.connection.recv())
-        await asyncio.sleep(MS)
-        await self.connection.close()
+
+        async def closer():
+            await asyncio.sleep(MS)
+            await self.connection.close()
+
+        asyncio.create_task(closer())
+
         with self.assertRaises(ConnectionClosedOK) as raised:
-            await recv_task
+            await self.connection.recv()
 
         exc = raised.exception
         self.assertEqual(str(exc), "sent 1000 (OK); then received 1000 (OK)")
         self.assertIsNone(exc.__cause__)
 
     async def test_close_during_send(self):
         """close fails the connection when called concurrently with send."""
-        gate = asyncio.get_running_loop().create_future()
+        close_gate = asyncio.Event()
+        exit_gate = asyncio.Event()
+
+        async def closer():
+            await close_gate.wait()
+            await self.connection.close()
+            exit_gate.set()
 
         async def fragments():
             yield "⏳"
-            await gate
+            close_gate.set()
+            await exit_gate.wait()
             yield "⌛️"
 
-        send_task = asyncio.create_task(self.connection.send(fragments()))
-        await asyncio.sleep(MS)
-
-        asyncio.create_task(self.connection.close())
-        await asyncio.sleep(MS)
-
-        gate.set_result(None)
+        asyncio.create_task(closer())
 
-        with self.assertRaises(ConnectionClosedError) as raised:
-            await send_task
+        iterator = fragments()
+        async with contextlib.aclosing(iterator):
+            with self.assertRaises(ConnectionClosedError) as raised:
+                await self.connection.send(iterator)
 
         exc = raised.exception
         self.assertEqual(
@@ -1050,7 +1062,7 @@ async def test_keepalive_times_out(self, getrandbits):
             # 4.x ms: a pong frame is dropped.
             await asyncio.sleep(5 * MS)
         # 6 ms: no pong frame is received; the connection is closed.
-        await asyncio.sleep(2 * MS)
+        await asyncio.sleep(3 * MS)
         # 7 ms: check that the connection is closed.
         self.assertEqual(self.connection.state, State.CLOSED)
 
@@ -1066,14 +1078,15 @@ async def test_keepalive_ignores_timeout(self, getrandbits):
             # 4.x ms: a pong frame is dropped.
             await asyncio.sleep(5 * MS)
         # 6 ms: no pong frame is received; the connection remains open.
-        await asyncio.sleep(2 * MS)
+        await asyncio.sleep(3 * MS)
         # 7 ms: check that the connection is still open.
         self.assertEqual(self.connection.state, State.OPEN)
 
     async def test_keepalive_terminates_while_sleeping(self):
         """keepalive task terminates while waiting to send a ping."""
         self.connection.ping_interval = 3 * MS
         self.connection.start_keepalive()
+        self.assertFalse(self.connection.keepalive_task.done())
         await asyncio.sleep(MS)
         self.assertFalse(self.connection.keepalive_task.done())
         await self.connection.close()
@@ -1231,9 +1244,7 @@ async def test_writing_in_data_received_fails(self):
         # The connection closed exception reports the injected fault.
         with self.assertRaises(ConnectionClosedError) as raised:
             await self.connection.recv()
-        cause = raised.exception.__cause__
-        self.assertEqual(str(cause), "Cannot call write() after write_eof()")
-        self.assertIsInstance(cause, RuntimeError)
+        self.assertIsInstance(raised.exception.__cause__, RuntimeError)
 
     async def test_writing_in_send_context_fails(self):
         """Error when sending outgoing frame is correctly reported."""
@@ -1244,9 +1255,7 @@ async def test_writing_in_send_context_fails(self):
         # The connection closed exception reports the injected fault.
         with self.assertRaises(ConnectionClosedError) as raised:
             await self.connection.pong()
-        cause = raised.exception.__cause__
-        self.assertEqual(str(cause), "Cannot call write() after write_eof()")
-        self.assertIsInstance(cause, RuntimeError)
+        self.assertIsInstance(raised.exception.__cause__, RuntimeError)
 
     # Test safety nets — catching all exceptions in case of bugs.
 
@@ -1336,11 +1345,11 @@ async def test_broadcast_skips_closing_connection(self):
 
     async def test_broadcast_skips_connection_with_send_blocked(self):
         """broadcast logs a warning when a connection is blocked in send."""
-        gate = asyncio.get_running_loop().create_future()
+        gate = asyncio.Event()
 
         async def fragments():
             yield "⏳"
-            await gate
+            await gate.wait()
 
         send_task = asyncio.create_task(self.connection.send(fragments()))
         await asyncio.sleep(MS)
@@ -1354,7 +1363,7 @@ async def fragments():
             ["skipped broadcast: sending a fragmented message"],
         )
 
-        gate.set_result(None)
+        gate.set()
         await send_task
 
     @unittest.skipIf(
@@ -1363,11 +1372,11 @@ async def fragments():
     )
     async def test_broadcast_reports_connection_with_send_blocked(self):
         """broadcast raises exceptions for connections blocked in send."""
-        gate = asyncio.get_running_loop().create_future()
+        gate = asyncio.Event()
 
         async def fragments():
             yield "⏳"
-            await gate
+            await gate.wait()
 
         send_task = asyncio.create_task(self.connection.send(fragments()))
         await asyncio.sleep(MS)
@@ -1381,7 +1390,7 @@ async def fragments():
         self.assertEqual(str(exc), "sending a fragmented message")
         self.assertIsInstance(exc, ConcurrencyError)
 
-        gate.set_result(None)
+        gate.set()
         await send_task
 
     async def test_broadcast_skips_connection_failing_to_send(self):