add test_datagram_write_error_close_from_callback

Author: graingert

Lib/test/test_asyncio/test_events.py

@@ -1809,6 +1809,47 @@ def connection_lost(self, exc):
         loop.run_until_complete(asyncio.wait_for(protocol.lost, 10))
         self.assertTrue(protocol.errors)
 
+    def test_datagram_write_error_close_from_callback(self):
+        # See https://github.com/python/cpython/issues/156920: an
+        # error_received() callback that closes the transport must still
+        # result in connection_lost() being called eventually, instead of
+        # leaving the transport hanging forever. Two failing writes are
+        # used so that the first failure's error_received() call closes
+        # the transport while the second is still queued (close() defers
+        # to the write loop), and the second failure then empties the
+        # buffer with self._closing already True and no write in flight
+        # -- exercising the `self._closing` half of the
+        # `if self._buffer or self._closing:` condition in _loop_writing.
+        loop = self.loop
+
+        class Protocol(asyncio.DatagramProtocol):
+            def connection_made(self, transport):
+                self.transport = transport
+                self.errors = []
+                self.lost = loop.create_future()
+
+            def error_received(self, exc):
+                self.errors.append(exc)
+                self.transport.close()
+
+            def connection_lost(self, exc):
+                if not self.lost.done():
+                    self.lost.set_result(exc)
+
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sock.setblocking(False)
+        sock.bind(('127.0.0.1', 0))
+        transport, protocol = loop.run_until_complete(
+            loop.create_datagram_endpoint(Protocol, sock=sock))
+        addr = sock.getsockname()
+
+        oversized = b'\x00' * 70000
+        transport.sendto(oversized, addr)
+        transport.sendto(oversized, addr)
+
+        loop.run_until_complete(asyncio.wait_for(protocol.lost, 10))
+        self.assertEqual(len(protocol.errors), 2)
+
     def test_datagram_recvfrom_connection_reset_recovers(self):
         # gh-127057: a UDP socket that sent a datagram to an address that
         # wasn't listening can raise ConnectionResetError on a later