如何在Pytest测试套件中保持WebSocket连接跨用例持续活跃？

我正在设计一套自动化测试套件，模拟客户端通过REST API登录后端后建立WebSocket通信，需测试REST与WebSocket的不同功能。当前每个WebSocket测试流程为：

• 客户端登录→启动WS通信→发送WS消息并等待响应→校验响应结构并断言结果→关闭WS连接→测试结束

遇到的问题：运行多个此类WebSocket测试时，因频繁启停WS连接，测试客户端被后端判定为异常行为拉黑，导致长时间无法重新建立连接。

技术栈：Pytest框架 + requests模块（处理API调用） + websockets模块（实现WS通信）

曾尝试用multiprocessing拆分进程，但无法实现Pytest进程与WebSocket进程间的消息发送和响应接收，陷入困境。

现有WebSocket连接逻辑代码：

async def websocket_connection(device: Device, cmd_list: list[WebsocketMsg] = None):
    init_cmd = WsInitCommand(device)
    cmd_list.insert(0, init_cmd)

    async def wait_for_correct_response(ws_connection, msj_id: str) -> dict:
        response_received = False
        ws_response: dict = {}
        while not response_received:
            ws_response = json.loads(await ws_connection.recv())
            if 'id' in ws_response and ws_response['id'] == msj_id:
                response_received = True
        return ws_response


    async with websockets.connect(init_cmd.url, subprotocols=init_cmd.sub_protocols) as websocket:
        for cmd in cmd_list:
            await websocket.send(str(cmd.message))
            msg_response: dict = await wait_for_correct_response(websocket, cmd.msg_id)
        return True

1. 复用WebSocket连接（最优方案） #

用Pytest的会话级/模块级fixture创建一次WebSocket连接，供所有测试用例复用，彻底避免频繁启停连接的问题。

示例代码： #

import pytest
import asyncio
import json
from websockets import WebSocketClientProtocol

@pytest.fixture(scope="session")
async def websocket_session(device: Device):
    # 先执行REST登录逻辑（此处省略具体登录代码）
    # login_response = requests.post(login_url, data=device.login_data)
    
    # 初始化WebSocket连接
    init_cmd = WsInitCommand(device)
    websocket = await websockets.connect(init_cmd.url, subprotocols=init_cmd.sub_protocols)
    
    # 发送初始化命令并等待响应
    await websocket.send(str(init_cmd.message))
    async def wait_for_correct_response(ws, msg_id):
        while True:
            response = json.loads(await ws.recv())
            if response.get('id') == msg_id:
                return response
    await wait_for_correct_response(websocket, init_cmd.msg_id)
    
    yield websocket  # 将连接提供给测试用例使用
    
    # 所有测试完成后优雅关闭连接
    await websocket.close(code=1000, reason="测试会话结束")

# 测试用例示例
async def test_websocket_get_status(websocket_session):
    test_cmd = WebsocketMsg(msg_id="test_status", message={"action": "get_status"})
    await websocket_session.send(str(test_cmd.message))
    response = await wait_for_correct_response(websocket_session, test_cmd.msg_id)
    assert response.get('code') == 200
    assert 'status' in response

async def test_websocket_update_config(websocket_session):
    test_cmd = WebsocketMsg(msg_id="test_config", message={"action": "update", "data": {"interval": 30}})
    await websocket_session.send(str(test_cmd.message))
    response = await wait_for_correct_response(websocket_session, test_cmd.msg_id)
    assert response.get('result') == 'success'

关键说明：

• fixture的scope设为session（所有测试用例共享）或module（单个模块共享），可根据测试隔离需求灵活调整。
• 确保连接在会话结束时正常关闭，避免后端残留无效连接。

2. 优化单测试连接的关闭逻辑 #

如果必须为每个测试单独建立连接，确保每次测试后优雅关闭WebSocket连接，避免后端将异常断开判定为恶意行为。

修改现有代码，添加异常处理和标准关闭流程：

async def websocket_connection(device: Device, cmd_list: list[WebsocketMsg] = None):
    init_cmd = WsInitCommand(device)
    cmd_list.insert(0, init_cmd)
    websocket = None
    
    async def wait_for_correct_response(ws_connection, msj_id: str) -> dict:
        while True:
            ws_response = json.loads(await ws_connection.recv())
            if ws_response.get('id') == msj_id:
                return ws_response

    try:
        websocket = await websockets.connect(init_cmd.url, subprotocols=init_cmd.sub_protocols)
        for cmd in cmd_list:
            await websocket.send(str(cmd.message))
            await wait_for_correct_response(websocket, cmd.msg_id)
        return True
    except Exception as e:
        print(f"WS通信异常: {str(e)}")
        return False
    finally:
        # 确保连接正常关闭
        if websocket and websocket.open:
            await websocket.close(code=1000, reason="测试完成正常关闭")

3. 多进程通信的正确实现 #

如果必须用多进程拆分逻辑，使用multiprocessing.Queue实现进程间消息传递：WebSocket进程负责维护连接和处理消息，Pytest进程负责发送测试命令和接收响应。

示例代码： #

import multiprocessing
import asyncio
import json

def websocket_worker(device, cmd_queue, response_queue):
    async def wait_for_correct_response(ws, msg_id):
        while True:
            response = json.loads(await ws.recv())
            if response.get('id') == msg_id:
                return response

    async def run():
        init_cmd = WsInitCommand(device)
        async with websockets.connect(init_cmd.url, subprotocols=init_cmd.sub_protocols) as websocket:
            # 完成初始化
            await websocket.send(str(init_cmd.message))
            await wait_for_correct_response(websocket, init_cmd.msg_id)
            
            # 循环处理命令队列
            while True:
                cmd = cmd_queue.get()
                if cmd is None:  # 接收退出信号
                    break
                await websocket.send(str(cmd.message))
                response = await wait_for_correct_response(websocket, cmd.msg_id)
                response_queue.put(response)
    
    asyncio.run(run())

# Pytest fixture初始化多进程
@pytest.fixture(scope="session")
def websocket_process(device):
    cmd_queue = multiprocessing.Queue()
    response_queue = multiprocessing.Queue()
    process = multiprocessing.Process(target=websocket_worker, args=(device, cmd_queue, response_queue))
    process.start()
    yield (cmd_queue, response_queue)
    # 发送退出信号并回收进程
    cmd_queue.put(None)
    process.join()

# 测试用例示例
def test_websocket_with_process(websocket_process):
    cmd_queue, response_queue = websocket_process
    test_cmd = WebsocketMsg(msg_id="test_proc", message={"action": "get_data"})
    cmd_queue.put(test_cmd)
    # 设置超时避免无限阻塞
    response = response_queue.get(timeout=10)
    assert response.get('code') == 200

关键说明：

• 用Queue传递命令和响应，解决进程间无法直接调用异步函数的问题。
• 测试结束后必须发送退出信号，确保WebSocket进程正常关闭连接并退出。

4. 限流测试执行速度 #

通过控制测试用例的执行间隔，避免短时间内发起大量连接请求，降低被拉黑的概率：

方法1：自定义Pytest钩子 #

在conftest.py中添加钩子：

import time
import pytest

@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_makereport(item, call):
    outcome = yield
    # 每个测试用例执行后等待1秒
    time.sleep(1)

方法2：使用pytest-slowdown插件 #

pip install pytest-slowdown
# 执行测试时设置每个测试间隔1秒
pytest --slowdown 1

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