Improve test

Author: fizzi01

tests/test_tcp.py

@@ -1264,14 +1264,20 @@ class _TestSSL(tb.SSLTestCase):
     TIMEOUT = 60
 
     def test_start_tls_buffer_transfer(self):
+        TIMEOUT = 10
+
         if (
             self.implementation == 'asyncio'
-            and sys.version_info[:2] <= (3, 11)
+            or sys.version_info[:2] < (3, 11)
         ):
             # StreamWriter.start_tls() introduced in Python 3.11
             raise unittest.SkipTest(
                 'StreamWriter.start_tls() not supported'
             )
+        self.loop.set_exception_handler(lambda loop, ctx: None)
+
+        client_read_buffered = asyncio.Event()
+        server_sent_ok = asyncio.Event()
 
         HELLO_MSG = b'1' * self.PAYLOAD_SIZE
         BUFFERED_MSG = b'buffered data before TLS'
@@ -1281,80 +1287,89 @@ def test_start_tls_buffer_transfer(self):
         client_context = self._create_client_ssl_context()
 
         async def handle_client(reader, writer):
-            # Send data before TLS upgrade
-            writer.write(BUFFERED_MSG)
-            await writer.drain()
-            await asyncio.sleep(0.2)
+            try:
+                # Send data before TLS upgrade
+                writer.write(BUFFERED_MSG)
+                await writer.drain()
 
-            # Read pre-TLS data
-            data = await reader.readexactly(len(HELLO_MSG))
-            self.assertEqual(len(data), len(HELLO_MSG))
+                await asyncio.wait_for(
+                    client_read_buffered.wait(),
+                    timeout=TIMEOUT
+                )
 
-            # Upgrade to TLS (server side)
-            try:
+                # Read pre-TLS data
+                data = await asyncio.wait_for(
+                    reader.readexactly(len(HELLO_MSG)),
+                    timeout=TIMEOUT,
+                )
+                self.assertEqual(len(data), len(HELLO_MSG))
+
+                # Upgrade to TLS (server side)
                 # We need the wait_for because the broken version hangs here
                 await asyncio.wait_for(
                     writer.start_tls(server_context),
-                    timeout=2)
-                self.assertIsNotNone(writer.get_extra_info('sslcontext'))
-            except asyncio.TimeoutError:
+                    timeout=TIMEOUT,)
                 self.assertIsNotNone(writer.get_extra_info('sslcontext'))
 
-            # Send/receive over TLS
-            writer.write(b'OK')
-            await writer.drain()
-
-            data = await reader.readexactly(len(HELLO_MSG))
-            self.assertEqual(len(data), len(HELLO_MSG))
+                # Send/receive over TLS
+                writer.write(b'OK')
+                await writer.drain()
+                server_sent_ok.set()
 
-            writer.close()
-            await self.wait_closed(writer)
+                data = await asyncio.wait_for(
+                    reader.readexactly(len(HELLO_MSG)),
+                    timeout=TIMEOUT,
+                )
+                self.assertEqual(len(data), len(HELLO_MSG))
+            finally:
+                if not writer.is_closing():
+                    writer.close()
+                    await self.wait_closed(writer)
 
         async def client(addr):
             # Use open_connection for StreamReader/StreamWriter
             reader, writer = await asyncio.open_connection(*addr)
 
-            # Read buffered data before TLS
-            buffered = await reader.readexactly(len(BUFFERED_MSG))
-            self.assertEqual(buffered, BUFFERED_MSG,
-                             "Wrong pre-TLS buffered data from server")
-
-            # Write before TLS upgrade
-            writer.write(HELLO_MSG)
-            await writer.drain()
-
-            # Upgrade to TLS
             try:
-                # We need the wait_for because the broken version hangs here
-                await asyncio.wait_for(
-                    writer.start_tls(client_context),
-                    timeout=2)
+                # Read buffered data before TLS
+                buffered = await reader.readexactly(len(BUFFERED_MSG))
+                self.assertEqual(buffered, BUFFERED_MSG,
+                                 "Wrong pre-TLS buffered data from server")
+                client_read_buffered.set()
+
+                # Write before TLS upgrade
+                writer.write(HELLO_MSG)
+                await writer.drain()
+
+                # Upgrade to TLS
+                await writer.start_tls(client_context)
                 self.assertIsNotNone(writer.get_extra_info('sslcontext'))
-            except asyncio.TimeoutError:
-                self.assertIsNotNone(writer.get_extra_info('sslcontext'))
-
-            # Verify communication over TLS
-            tls_data = await reader.readexactly(2)
-            self.assertEqual(tls_data, b'OK',
-                             "Wrong data from server after TLS upgrade")
 
-            # Continue over TLS
-            writer.write(HELLO_MSG)
-            await writer.drain()
+                # Verify communication over TLS
+                await server_sent_ok.wait()
+                tls_data = await reader.readexactly(2)
+                self.assertEqual(tls_data, b'OK',
+                                 "Wrong data from server after TLS upgrade")
 
-            writer.close()
-            await self.wait_closed(writer)
+                # Continue over TLS
+                writer.write(HELLO_MSG)
+                await writer.drain()
+            finally:
+                if not writer.is_closing():
+                    writer.close()
+                    await self.wait_closed(writer)
 
         async def run_test():
             srv = await asyncio.start_server(
                 handle_client, '127.0.0.1', 0, family=socket.AF_INET)
 
-            addr = srv.sockets[0].getsockname()
-
-            await asyncio.wait_for(client(addr), timeout=10)
+            try:
+                addr = srv.sockets[0].getsockname()
 
-            srv.close()
-            await srv.wait_closed()
+                await asyncio.wait_for(client(addr), timeout=self.TIMEOUT)
+            finally:
+                srv.close()
+                await srv.wait_closed()
 
         self.loop.run_until_complete(run_test())