FastAPI WebSocket连接关闭测试：Pytest实战指南

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

# ... (省略导入和应用初始化代码)

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"})

from fastapi import APIRouter, WebSocket, Depends
from fastapi.websockets import WebSocketDisconnect
from typing import Annotated

# 假设 GameManager 和 get_manager 已定义
from src.game_manager import GameManager, get_manager 

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:
        # manager.connect 负责处理连接逻辑，如果条件不满足，它会抛出 WebSocketDisconnect
        await manager.connect(websocket, room_name, password) 
        # 获取客户端ID
        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)

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

# 假设 src.game_manager 和 src.main 已正确导入
from src.game_manager import GameManager, get_manager
from src.main import app

# 依赖覆盖，用于测试环境
async def override_manager() -> GameManager:
    try:
        yield override_manager.manager
    except AttributeError:
        manager = GameManager()
        manager.add_new_game(max_players=2, room_name="foo", password=None)
        manager.add_new_game(max_players=2, room_name="bar", password="123")
        override_manager.manager = manager
        yield override_manager.manager

app.dependency_overrides[get_manager] = override_manager

client = TestClient(app)

class TestWebsocketConnection:
    def test_connect_to_non_existing_room_correct(self):
        # 期望 WebSocketDisconnect 异常被抛出
        with pytest.raises(WebSocketDisconnect):
            # 建立 WebSocket 连接
            with client.websocket_connect("/ws/non_existing_room") as ws:
                # 尝试接收数据。由于服务器会立即关闭连接，此操作会抛出 WebSocketDisconnect
                ws.receive_json()