FastAPI WebSocket连接关闭的PyTest测试实践

import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect

# 假设app和get_manager以及override_manager已正确定义
# ... (省略了app和manager的依赖覆盖代码)
client = TestClient(app)

class TestWebsocketConnection:
    def test_connect_to_non_existing_room_initial_attempt(self):
        with pytest.raises(WebSocketDisconnect) as e_info:
            with client.websocket_connect("/ws/non_existing_room") as ws:
                # 尝试发送数据，但如果连接已关闭，可能不会立即触发异常
                ws.send_json({"message": "Hello world"})

# 运行时可能返回：
# FAILED tests/test_websockets.py::TestWebsocketConnection::test_connect_to_non_existing_room - Failed: DID NOT RAISE <class 'starlette.websockets.WebSocketDisconnect'>

import pytest
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocketDisconnect
from typing import Annotated

# 假设你的FastAPI应用和GameManager的定义如下
# src/game_manager.py
class GameManager:
    def __init__(self):
        self.games = {} # 存储游戏房间信息

    async def connect(self, websocket, room_name, password):
        if room_name not in self.games:
            # 如果房间不存在，则抛出WebSocketDisconnect
            raise WebSocketDisconnect(code=1008, reason="Room does not exist")
        # 实际连接逻辑...
        await websocket.accept()
        print(f"Client connected to room: {room_name}")
        # 这里为了测试，假设连接成功后不会立即发送数据

    async def remove(self, websocket):
        # 清理连接逻辑
        print("Client disconnected.")

    async def handle_message(self, room_name, client_id, data):
        # 处理消息逻辑
        pass

# src/main.py
from fastapi import FastAPI, APIRouter, Depends, WebSocket
from fastapi.routing import APIRoute

# 为了演示，这里简化get_manager
def get_manager() -> GameManager:
    return GameManager()

app = FastAPI()
router = APIRouter()

@router.websocket("/ws/{room_name}")
@router.websocket("/ws/{room_name}/{password}")
async def websocket_endpoint(
    websocket: WebSocket,
    manager: Annotated[GameManager, Depends(get_manager)],
):
    room_name = websocket.path_params["room_name"]
    password = websocket.path_params.get("password", None)

    try:
        await manager.connect(websocket, room_name, password)
        # client_id = websocket.scope["client_id"] # 实际应用中会获取
        while True:
            data = await websocket.receive_json()
            # await manager.handle_message(room_name, client_id, data) # 实际应用中会处理
    except WebSocketDisconnect:
        await manager.remove(websocket)
    except Exception as e:
        print(f"Unexpected error: {e}")
        await manager.remove(websocket)


app.include_router(router)

# tests/test_websockets.py
# 依赖覆盖，确保测试环境隔离且可控
async def override_get_manager() -> GameManager:
    try:
        # 尝试使用已存在的manager实例
        yield override_get_manager.manager
    except AttributeError:
        # 如果不存在，则创建并初始化一个新的manager
        manager = GameManager()
        manager.games["foo"] = {} # 添加一个存在的房间用于其他测试
        override_get_manager.manager = manager
        yield override_get_manager.manager

# 将依赖覆盖应用到FastAPI应用
app.dependency_overrides[get_manager] = override_get_manager

client = TestClient(app)

class TestWebsocketConnection:
    def test_connect_to_non_existing_room_correctly_closed(self):
        """
        测试连接到不存在的房间时，连接是否被正确关闭。
        通过尝试接收数据来触发WebSocketDisconnect异常。
        """
        with pytest.raises(WebSocketDisconnect) as excinfo:
            with client.websocket_connect("/ws/non_existing_room") as ws:
                # 关键步骤：尝试从已关闭的连接接收数据
                # 这将触发并捕获WebSocketDisconnect异常
                ws.receive_json()

        # 可选：进一步断言异常的详细信息，例如错误码或原因
        assert excinfo.type is WebSocketDisconnect
        assert excinfo.value.code == 1008
        assert "Room does not exist" in excinfo.value.reason