如何在同一进程中用pytest测试asyncio Socket服务器与客户端

我用asyncio写了一个包含Server和Client类的Socket服务端客户端模块，逻辑是检测本地指定端口是否被占用，没被占用就启动Server，被占用就用Client，用来做本地测试。用pytest写单元测试时，把服务端和客户端分开在不同进程跑测试能通过，但在同一进程里同时测两者时，测试永远完成不了。想知道这是asyncio的问题、pytest的问题，还是我自己代码的问题？

WebSockets.websocket 原代码 #

import asyncio


class Server:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.server = None
        self.connections = {}

    async def start(self):
        self.server = server = await asyncio.start_server(self.handle_client, self.host, self.port)
        addr = server.sockets[0].getsockname()
        print(f'Serving on {addr}')
        async with server:
            await server.serve_forever()

    async def handle_client(self, reader, writer):
        addr = writer.get_extra_info('peername')
        print(f'New connection from {addr}')
        self.connections[addr] = writer
        try:
            while not reader.at_eof():
                data = await reader.read(100)
                message = data.decode()
                if message:
                    print(f'Received {message!r} from {addr}')
                    await self.send(message, addr)
        finally:
            del self.connections[addr]
            writer.close()

    async def send(self, message, addr):
        writer = self.connections.get(addr)
        if not writer:
            return
        writer.write(message.encode())
        await writer.drain()

    async def recv(self, addr):
        reader, _ = await asyncio.open_connection(addr[0], addr[1])
        data = await reader.read(100)
        return data.decode()

    async def stop(self):
        await self.server.close()


class Client:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.reader = None
        self.writer = None

    async def connect(self):
        self.reader, self.writer = await asyncio.open_connection(self.host, self.port)

    async def send(self, message):
        self.writer.write(message.encode())

    async def recv(self):
        data = await self.reader.read(100)
        return data.decode()

    async def close(self):
        self.writer.close()

原测试代码 #

import socket
import asyncio
from WebSockets.websocket import Server, Client
import pytest


def check_port(address, port):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind((address, port))
        s.close()
        return "server"
    except OSError:
        return "client"


async def serve():
    server = Server('localhost', 8910)
    await server.start()


async def client_connect():
    client = Client('localhost', 8910)
    await client.connect()
    await client.send("Hello world!")

    message = await client.recv()
    if message:
        print(f"Received {message!r} from server")
        assert message == "Hello world!"
        await client.close()


@pytest.mark.asyncio
async def test_sockets():
    if check_port("localhost", 8910) == "server":
        await serve()
        # for reference on how to use in non pytest module
        # loop = asyncio.new_event_loop()
        # loop.run_until_complete(serve())
        # loop.close()
    else:
        await client_connect()
        # for reference on how to use in non pytest module
        # loop = asyncio.new_event_loop()
        # loop.run_until_complete(client_connect())
        # loop.close()

问题原因与解决方法 #

问题根源 #

这是你的代码逻辑问题，和asyncio、pytest无关。当测试进入serve()函数后，await server.start()会执行到await server.serve_forever()——这个调用会永久阻塞事件循环，不会主动返回，导致测试进程卡在这，永远无法结束。

你原测试逻辑是在同一个函数里要么跑服务端要么跑客户端，但跑服务端时没有停止服务的逻辑，自然无法完成测试。

修复方案 #

1. 修改Server类：让服务端启动后不阻塞主线程，改用异步任务运行，同时完善停止逻辑
2. 修改Client类：修复send方法的数据发送问题
3. 调整测试函数：让服务端和客户端在同一事件循环内协作，完成测试后主动停止服务端

修改后的Server类关键代码

class Server:
    # ... 其他方法保持不变 ...

    async def start(self):
        self.server = await asyncio.start_server(self.handle_client, self.host, self.port)
        addr = self.server.sockets[0].getsockname()
        print(f'Serving on {addr}')
        # 用异步任务运行服务，避免阻塞主线程
        self.server_task = asyncio.create_task(self.server.serve_forever())

    async def stop(self):
        self.server.close()
        await self.server.wait_closed()
        # 取消服务任务并处理取消异常
        if hasattr(self, 'server_task'):
            self.server_task.cancel()
            try:
                await self.server_task
            except asyncio.CancelledError:
                pass

修改后的Client类send方法

class Client:
    # ... 其他方法保持不变 ...

    async def send(self, message):
        self.writer.write(message.encode())
        await self.writer.drain()  # 确保数据被发送到服务端

修改后的测试函数

@pytest.mark.asyncio
async def test_sockets():
    port = 8910
    role = check_port("localhost", port)
    
    if role == "server":
        # 启动服务端
        server = Server('localhost', port)
        await server.start()
        try:
            # 执行客户端测试逻辑
            await client_connect()
        finally:
            # 测试完成后停止服务端
            await server.stop()
    else:
        await client_connect()

修改后，服务端会以异步任务的方式运行，不会阻塞测试主线程，客户端测试完成后主动停止服务端，整个测试流程就能正常结束了。

内容的提问来源于stack exchange，提问作者Dshiz