FastAPI WebSocket测试阻塞：线程进程导致测试无限运行

为FastAPI WebSocket应用编写测试时，单个测试会无限运行无法停止，导致后续测试无法启动。原因是FastAPI启动时调用了一个线程，该线程又生成了进程，线程与进程的组合在TestClient中引发问题，致使测试无法终止。相关代码如下：

import uvicorn as uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import threading
import asyncio
from fastapi.testclient import TestClient
from multiprocessing import Process

app = FastAPI()

class ProcessHeavyCPU:
    def __init__(self): pass

    def run(self):  
        while True:
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within process")

class ThreadHandler(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        processes = []
        for _ in range(2):
            p = Process(target=ProcessHeavyCPU().run)
            processes.append(p)

        [x.start() for x in processes]

        while True:                    
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within thread")

@app.on_event("startup")
async def startup_event():
    thread1 = ThreadHandler()
    thread1.daemon = True
    thread1.start()

async def handle_msg(websocket):
    await websocket.send_json({"msg": "New msg received"})

background_tasks = set()

@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            msg = await websocket.receive_text()
            task = asyncio.create_task(handle_msg(websocket))
            background_tasks.add(task)
            task.add_done_callback(background_tasks.discard)
    except WebSocketDisconnect:
        await websocket.close()

def test_example():         
    with TestClient(app) as client:
        with client.websocket_connect("/") as websocket:
            websocket.send_json({
                "new_msg": "xyz"
            })
            resp = websocket.receive_json()
            assert resp["msg"] == "New msg received"
            print("finished test")

if __name__ == '__main__':
    uvicorn.run("main:app", host="0.0.0.0", port=8081, reload=True, access_log=False)

执行测试命令：python -m pytest .\main.py
使用Python版本：3.10.5

核心修复思路 #

解决测试阻塞的关键是隔离测试环境与后台线程/进程，同时给后台任务添加可终止机制，避免无限循环占用资源。

具体修改方案 #

1. 给子进程设置守护属性
   创建Process对象时添加daemon=True，确保主进程退出时自动终止子进程，防止子进程脱离控制无限运行。

2. 添加测试环境判断
   通过环境变量控制是否启动后台线程，测试时跳过启动逻辑，直接测试WebSocket核心功能。

3. 修复WebSocket消息类型不匹配问题
   原测试中用send_json发送消息，但接口定义的是receive_text，会导致接收超时，需改为对应类型的调用。

4. 添加后台任务主动终止机制
   给线程和进程添加停止标志，在应用关闭时主动清理资源。

修改后的完整代码 #

import os
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import threading
import asyncio
from fastapi.testclient import TestClient
from multiprocessing import Process

app = FastAPI()

class ProcessHeavyCPU:
    def __init__(self):
        self.running = True

    def run(self):  
        while self.running:
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within process")

class ThreadHandler(threading.Thread):
    def __init__(self):
        super().__init__()
        self.running = True
        self.processes = []

    def run(self):
        for _ in range(2):
            p = Process(target=ProcessHeavyCPU().run)
            p.daemon = True  # 设置子进程为守护进程
            self.processes.append(p)

        [x.start() for x in self.processes]

        while self.running:                    
            print(f"Thread ID: [{str(threading.get_ident())}] | endless loop within thread")

    def stop(self):
        self.running = False
        # 主动终止子进程
        for p in self.processes:
            if p.is_alive():
                p.terminate()
                p.join()

# 全局保存线程实例，用于关闭时清理
thread_handler = None

@app.on_event("startup")
async def startup_event():
    global thread_handler
    # 测试环境下不启动后台线程
    if os.getenv("TESTING") != "1":
        thread_handler = ThreadHandler()
        thread_handler.daemon = True
        thread_handler.start()

@app.on_event("shutdown")
async def shutdown_event():
    global thread_handler
    if thread_handler:
        thread_handler.stop()

async def handle_msg(websocket):
    await websocket.send_json({"msg": "New msg received"})

background_tasks = set()

@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            msg = await websocket.receive_text()
            task = asyncio.create_task(handle_msg(websocket))
            background_tasks.add(task)
            task.add_done_callback(background_tasks.discard)
    except WebSocketDisconnect:
        await websocket.close()

def test_example():
    # 设置测试环境变量
    os.environ["TESTING"] = "1"
    with TestClient(app) as client:
        with client.websocket_connect("/") as websocket:
            websocket.send_text("xyz")  # 与接口receive_text对应
            resp = websocket.receive_json()
            assert resp["msg"] == "New msg received"
            print("finished test")
    # 清理环境变量
    del os.environ["TESTING"]

if __name__ == '__main__':
    uvicorn.run("main:app", host="0.0.0.0", port=8081, reload=True, access_log=False)

测试验证 #

执行原测试命令python -m pytest .\main.py，测试会正常执行并退出，不会出现无限阻塞的情况。

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