Function tcp_server

test/mitmproxy/addons/test_clientplayback.py:19–54

TCP server context manager that... 1. Exits only after all handlers have returned. 2. Ensures that all handlers are closed properly. If we don't do that, we get ghost errors in others tests from StreamWriter.__del__. Spawning a TCP server is relatively slow. Consider using in-me

(handle_conn, **server_args)

@asynccontextmanager
async def tcp_server(handle_conn, **server_args) -> Address:
    """TCP server context manager that...

    1. Exits only after all handlers have returned.
    2. Ensures that all handlers are closed properly. If we don't do that,
       we get ghost errors in others tests from StreamWriter.__del__.

    Spawning a TCP server is relatively slow. Consider using in-memory networking for faster tests.
    """
    if not hasattr(asyncio, "TaskGroup"):
        pytest.skip("Skipped because asyncio.TaskGroup is unavailable.")

    tasks = asyncio.TaskGroup()

    async def handle_conn_wrapper(
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        try:
            await handle_conn(reader, writer)
        except Exception as e:
            print(f"!!! TCP handler failed: {e}")
            raise
        finally:
            if not writer.is_closing():
                writer.close()
            await writer.wait_closed()

    async def _handle(r, w):
        tasks.create_task(handle_conn_wrapper(r, w))

    server = await asyncio.start_server(_handle, "127.0.0.1", 0, **server_args)
    await server.start_serving()
    async with server:
        async with tasks:
            yield server.sockets[0].getsockname()


@pytest.mark.parametrize("mode", ["http", "https", "upstream", "err"])