pytest中WebSocket调用receive_text()时出现冻结问题排查

问题背景 #

旧版本WebSocket私聊服务可正常通过pytest测试，更新代码后，测试调用receive_text()时出现冻结现象，且WebSocketDisconnect异常无法触发。旧版本中message = await self._receive()会返回"disconnecting"，新版本则卡在websocket.receive步骤。

相关代码文件 #

config.py #

engine_test = create_async_engine(TEST_POSTGRES_URI, poolclass=NullPool)
async_session_maker = async_sessionmaker(
    engine_test, class_=AsyncSession, expire_on_commit=False
)


async def override_session():
    async with async_session_maker() as session:
        yield session

topus.dependency_overrides[get_session] = override_session
async def defaults_user():
    try:
        async with async_session_maker() as session:
            test_user = UserDB(
                username="TestUserDB",
                password=hash_password("12345678"),
            )
            test_user2 = UserDB(
                username="TestUserDB2",
                password=hash_password("12345678"),
            )

            session.add_all([test_user, test_user2])
            await session.commit()
    except Exception as e:
        print(f"An error occurred while adding default user: {e}")

@pytest.fixture(autouse=True, scope='session')
async def lifespan():
    async with engine_test.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    await defaults_user()
    yield
    async with engine_test.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)

测试正常版本：old_chat.py #

class PrivateManager:
    """ 私聊消息管理器
    """

    def __init__(self):
        self.connections: Dict[str, List[WebSocket]] = {}

    def register_websocket(self, chat_id: str, websocket: WebSocket) -> None:
        if chat_id not in self.connections:
            self.connections[chat_id] = []
        self.connections[chat_id].append(websocket)

    async def broadcast(self, chat_id: str, message: str, sender_id: int, friend_id: int, add_to_db: bool, db: AsyncSession = None) -> None:
        if chat_id not in self.connections:
            return

        if add_to_db:
            await self.save_message_to_db(chat_id=chat_id, message=message, sender=sender_id, friend_id=friend_id, db=db)

        for websocket in self.connections[chat_id]:
            await websocket.send_text(message)

    def disconnect(self, chat_id, websocket: WebSocket):
        chat: list = self.connections.get(chat_id)
        chat.remove(websocket)

@chat.websocket('/private_chat/{friend_id}')
async def private_chat(friend_id: int, websocket: WebSocket, token: str = Depends(decode_token), db: AsyncSession = Depends(get_session)):
    if token['id'] == friend_id:
        return JSONResponse(status_code=status.HTTP_403_FORBIDDEN, content='不能给自己发消息')
    chat_id = chat_id_generator(token['id'], friend_id)

    private_manager.register_websocket(chat_id, websocket)
    await websocket.accept()

    try:
        while True:
            data = await websocket.receive_text()
            await private_manager.broadcast(chat_id=chat_id, message=data, sender_id=token['id'], friend_id=friend_id, db=db, add_to_db=True)
    except WebSocketDisconnect:
        logging.error(msg='da', exc_info=True)
        private_manager.disconnect(chat_id, websocket)

pytest测试代码 #

"""
用户发送私聊消息
"""
client = TestClient(topus)
login = {
    "username": "TestUserDB",
    "password": "12345678",
}
response = client.post('/auth/login', json=login)
assert response.status_code == 200

with client.websocket_connect('/chat/private_chat/2') as websocket:
    websocket.send_text('Hi, Private!')
    ans = websocket.receive_text()
    assert ans == 'Hi, Private!'
    websocket.close()

出现冻结版本：new_chat.py #

class PrivateManager:
    """ 私聊消息管理器
    """

    def __init__(self):
        self.connections: Dict[str, List[WebSocket]] = {}

    async def register_websocket(self, chat_id: str, user_id: int, friend_id: int,  websocket: WebSocket, db: AsyncSession) -> None:
        if chat_id not in self.connections:
            self.connections[chat_id] = []
        self.connections[chat_id].append(websocket)
        await get_or_create_chat(
            chat_id=chat_id, user_id=user_id, friend_id=friend_id, db=db)

    async def broadcast(self, chat_id: str, message: str, sender_id: int, friend_id: int, add_to_db: bool, userwebsocket: WebSocket, db: AsyncSession = None) -> None:
        if chat_id not in self.connections:
            return

        if add_to_db:
            await self.save_message_to_db(chat_id=chat_id, message=message, sender=sender_id, friend_id=friend_id, db=db)

        for websocket in self.connections[chat_id]:
            if websocket != userwebsocket:
                await websocket.send_text(message)

    def disconnect(self, chat_id, websocket: WebSocket):
        chat: list = self.connections.get(chat_id)
        if chat:
            chat.remove(websocket)

    @staticmethod
    async def save_message_to_db(chat_id: str, message: str, sender: int, friend_id, db: AsyncSession) -> None:
        await save_message(chat_id=chat_id, message=message, sender=sender, friend_id=friend_id, db=db)
   

async def private_chat(user_id: int, friend_id: int, websocket: WebSocket, db: AsyncSession = Depends(get_session)):
    if user_id == friend_id:
        raise WebSocketException(
            code=status.WS_1007_INVALID_FRAME_PAYLOAD_DATA, reason='不能给自己发消息')
    chat_id = chat_id_generator(user_id, friend_id)

    await private_manager.register_websocket(chat_id=chat_id, user_id=user_id, friend_id=friend_id, websocket=websocket, db=db)
    await websocket.accept()

    try:
        while True:
            data = await websocket.receive_text()
            await private_manager.broadcast(chat_id=chat_id, message=data, sender_id=user_id, friend_id=friend_id, db=db, add_to_db=True, userwebsocket=websocket)

    except WebSocketDisconnect:
        private_manager.disconnect(chat_id, websocket)

venv中WebSocket类相关代码 #

class WebSocket(HTTPConnection):
    def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:
        super().__init__(scope)
        assert scope["type"] == "websocket"
        self._receive = receive
        self._send = send
        self.client_state = WebSocketState.CONNECTING
        self.application_state = WebSocketState.CONNECTING

    async def receive(self) -> Message:
        """
        接收ASGI websocket消息，确保状态转换合法。
        """
        if self.client_state == WebSocketState.CONNECTING:
            message = await self._receive()
            message_type = message["type"]
            if message_type != "websocket.connect":
                raise RuntimeError(
                    'Expected ASGI message "websocket.connect", '
                    f"but got {message_type!r}"
                )
            self.client_state = WebSocketState.CONNECTED
            return message
        elif self.client_state == WebSocketState.CONNECTED:
            message = await self._receive()
            message_type = message["type"]
            if message_type not in {"websocket.receive", "websocket.disconnect"}:

问题分析与修复方案 #

核心问题 #

1. 测试冻结原因：新版本广播逻辑添加了if websocket != userwebsocket判断，仅给其他连接发送消息。但测试代码只建立了一个WebSocket连接，发送消息后测试端的receive_text()永远收不到返回值，导致阻塞冻结。
2. WebSocketDisconnect异常不触发：可能是register_websocket中的get_or_create_chat函数存在数据库操作阻塞/死锁，或者广播过程中未捕获异常导致接收循环中断，无法响应断开事件。

修复步骤 #

解决测试冻结

调整测试用例，模拟两个用户的连接，确保有接收方在线：

# 修改后的测试代码
with client.websocket_connect('/chat/private_chat/2') as ws1, client.websocket_connect('/chat/private_chat/1') as ws2:
    ws1.send_text('Hi, Private!')
    ans = ws2.receive_text()
    assert ans == 'Hi, Private!'
    ws1.close()
    ws2.close()

解决WebSocketDisconnect异常问题

• 检查get_or_create_chat函数，确保数据库操作无阻塞、无死锁，避免卡住WebSocket连接流程。
• 给广播逻辑添加异常捕获，防止单个连接发送失败导致整个接收循环中断：

async def broadcast(self, chat_id: str, message: str, sender_id: int, friend_id: int, add_to_db: bool, userwebsocket: WebSocket, db: AsyncSession = None) -> None:
    if chat_id not in self.connections:
        return

    if add_to_db:
        await self.save_message_to_db(chat_id=chat_id, message=message, sender=sender_id, friend_id=friend_id, db=db)

    for websocket in self.connections[chat_id]:
        if websocket != userwebsocket:
            try:
                await websocket.send_text(message)
            except Exception as e:
                logging.error(f"Failed to send message to websocket: {e}")
                self.disconnect(chat_id, websocket)

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