在Pytest中测试FastAPI WebSocket连接关闭的正确方法

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

# 假设 app 和 get_manager, override_manager 等已定义
# ... (省略依赖注入和TestClient初始化代码) ...

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

@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)
        # 如果 manager.connect 成功，客户端才会被接受并进入消息循环
        client_id = websocket.scope["client_id"]
        while True:
            data = await websocket.receive_json()
            await manager.handle_message(room_name, client_id, data)
    except WebSocketDisconnect:
        # 当 manager.connect 抛出 WebSocketDisconnect 或客户端断开时，这里会捕获
        await manager.remove(websocket)
        # FastAPI/Starlette 框架会确保连接被正确关闭

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

# 假设以下导入和定义已存在
from src.game_manager import GameManager
from src.main import app
from fastapi import Depends, WebSocket, APIRouter # 假设 router 是一个 APIRouter 实例

# 假设 get_manager 是一个依赖函数
def get_manager():
    # 这是一个占位符，实际应用中会从依赖注入系统获取 GameManager 实例
    pass

# 重写依赖以进行测试
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(self):
        """
        测试连接到不存在的房间时，WebSocket连接是否被正确关闭。
        """
        with pytest.raises(WebSocketDisconnect):
            with client.websocket_connect("/ws/non_existing_room") as ws:
                # 尝试从已关闭的连接接收数据，这将触发 WebSocketDisconnect
                ws.receive_json()

    def test_connect_to_existing_room(self):
        """
        测试连接到存在的房间时，连接是否成功建立且未断开。
        """
        with client.websocket_connect("/ws/foo") as ws:
            # 确保连接成功，可以发送和接收数据
            ws.send_json({"message": "Hello from foo"})
            # 接收一条消息，验证连接活跃
            response = ws.receive_json()
            assert "type" in response and response["type"] == "websocket.send" # 示例断言

        # 确保连接在退出 with 块后正常关闭
        with pytest.raises(WebSocketDisconnect):
            # 尝试在连接关闭后接收数据，验证其已关闭
            ws.receive_json()