pytest测试FastAPI永久WebSocket连接后挂起，如何终止TestClient？

1. 测试永久WebSocket连接的正确方式 #

主动关闭客户端连接 #

测试场景下无需严格遵循“仅客户端关闭”的设计，测试结束时显式关闭WebSocket连接，避免资源占用：

async def test_comedy_tom_hanks_async(websocket):
    # 执行测试逻辑
    # ...
    # 测试结束主动关闭连接
    await websocket.close()

修复线程泄漏问题 #

从报错来看，asyncio.to_thread创建的线程没有正常退出，导致pytest挂起。需要：

• 在read_from_ws函数中加入中断检查，比如用threading.Event控制线程退出：

def read_from_ws(websocket, stop_event):
    while not stop_event.is_set():
        # 读取WebSocket数据逻辑
        # ...

测试时创建事件，结束时触发：

async def test_comedy_tom_hanks_async(websocket):
    stop_event = threading.Event()
    try:
        coro = asyncio.to_thread(read_from_ws, websocket, stop_event)
        coro_waited = asyncio.wait_for(coro, timeout)
        results = await asyncio.gather(coro_waited)
    finally:
        stop_event.set()  # 触发线程退出

显式管理事件循环与线程 #

在fixture的teardown阶段清理线程，确保所有非守护线程退出：

import threading

@pytest.fixture(scope="function", name="websocket")
async def fixture_ws_audio():
    client = TestClient(app)
    with client.websocket_connect("/ws/audio") as websocket:
        yield websocket
    # 清理遗留线程
    for thread in threading.enumerate():
        if thread.is_alive() and thread != threading.main_thread():
            thread.join(timeout=2)  # 等待线程退出，超时则放弃

2. 强制终止Fixture的方法 #

在Fixture中添加强制清理逻辑 #

修改fixture，在yield之后主动关闭连接并标记线程为守护线程（主进程退出时自动终止）：

import threading

@pytest.fixture(scope="function", name="websocket")
async def fixture_ws_audio():
    client = TestClient(app)
    with client.websocket_connect("/ws/audio") as websocket:
        yield websocket
    # Teardown：强制关闭连接
    try:
        await websocket.close(code=1000)
    except Exception:
        pass
    # 标记非主线程为守护线程，让主进程退出时自动终止
    for thread in threading.enumerate():
        if thread != threading.main_thread():
            thread.daemon = True

使用pytest终结函数 #

通过request.addfinalizer添加终结函数，确保无论测试成功与否都会执行清理：

@pytest.fixture(scope="function", name="websocket")
async def fixture_ws_audio(request):
    client = TestClient(app)
    ws_context = client.websocket_connect("/ws/audio")
    websocket = ws_context.__enter__()
    
    def finalize():
        # 关闭WebSocket连接
        ws_context.__exit__(None, None, None)
        # 清理线程逻辑
        for thread in threading.enumerate():
            if thread.is_alive() and thread != threading.main_thread():
                thread.daemon = True
    request.addfinalizer(finalize)
    
    yield websocket

自定义线程池避免泄漏 #

使用自定义线程池执行read_from_ws，测试结束后线程池自动关闭，避免线程残留：

import asyncio
import concurrent.futures

async def test_comedy_tom_hanks_async(websocket):
    loop = asyncio.get_running_loop()
    # 使用自定义线程池，测试结束自动销毁
    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
        coro = loop.run_in_executor(executor, read_from_ws, websocket)
        coro_waited = asyncio.wait_for(coro, timeout)
        results = await asyncio.gather(coro_waited)

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