diff --git a/tests/conftest.py b/tests/conftest.py index 15cc647..e18ee08 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -76,7 +76,7 @@ async def _pipe(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter if writer not in self.frozen_connections: writer.close() if hasattr(writer, "wait_closed"): - await writer.wait_closed() + await asyncio.shield(writer.wait_closed()) async def _handle_client( self, diff --git a/tests/test_tcp_proxy.py b/tests/test_tcp_proxy.py new file mode 100644 index 0000000..757af12 --- /dev/null +++ b/tests/test_tcp_proxy.py @@ -0,0 +1,24 @@ +import asyncio +import contextlib + +from .conftest import TcpProxy + + +async def test_cancelled_pipe_does_not_break_connection_closure() -> None: + server = await asyncio.start_server(lambda _, __: None, host="127.0.0.1", port=0) + port = server.sockets[0].getsockname()[1] + try: + _, writer = await asyncio.open_connection(host="127.0.0.1", port=port) + reader = asyncio.StreamReader() + reader.feed_eof() + + proxy = TcpProxy(src_port=0, dst_port=port) + pipe = asyncio.create_task(proxy._pipe(reader, writer)) # pyright: ignore[reportPrivateUsage] + await asyncio.sleep(0) + pipe.cancel() + with contextlib.suppress(asyncio.CancelledError): + await pipe + + await writer.wait_closed() + finally: + server.close()