使用PyTest测试FastAPI WebSocket连接的关闭：一种可靠的方法

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

# ... (FastAPI应用和依赖覆盖设置) ...

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

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

# 假设你的应用结构如下，这里为了完整性提供示例
# src/game_manager.py
class GameManager:
    def __init__(self):
        self.games = {}

    def add_new_game(self, max_players, room_name, password):
        self.games[room_name] = {"max_players": max_players, "password": password, "clients": {}}

    async def connect(self, websocket: WebSocket, room_name: str, password: str | None):
        if room_name not in self.games:
            # 关键：如果房间不存在，立即抛出 WebSocketDisconnect
            raise WebSocketDisconnect(code=1008, reason="Room does not exist")

        # 假设这里会处理密码验证等，并最终接受连接
        await websocket.accept()
        # 假设 client_id 是从某个地方生成的
        client_id = f"client_{len(self.games[room_name]['clients'])}"
        websocket.scope["client_id"] = client_id
        self.games[room_name]["clients"][client_id] = websocket
        print(f"Client {client_id} connected to {room_name}")


    async def handle_message(self, room_name, client_id, data):
        print(f"Received message from {client_id} in {room_name}: {data}")

    async def remove(self, websocket: WebSocket):
        # 实际的移除逻辑
        print(f"Client {websocket.scope.get('client_id')} disconnected.")

# src/main.py
from fastapi import FastAPI
app = FastAPI()
router = APIRouter()

def get_manager():
    # 实际应用中可能是单例或依赖注入
    return GameManager()

@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)

app.include_router(router)


# tests/test_websockets.py
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

# 假设 get_manager 是你的依赖注入函数
from src.main import get_manager
app.dependency_overrides[get_manager] = override_manager

client = TestClient(app)

class TestWebsocketConnection:
    def test_connect_to_non_existing_room_solution(self):
        # 使用 pytest.raises 包裹，并在连接建立后尝试接收数据
        with pytest.raises(WebSocketDisconnect):
            with client.websocket_connect("/ws/non_existing_room") as ws:
                # 关键步骤：尝试从已关闭的连接接收数据
                ws.receive_json()