如何在Rocket框架的异步测试中连接WebSocket路由？

这个问题我之前也碰到过，Rocket的本地测试客户端确实没绑定真实端口，没法直接用Tungstenite这类外部客户端连接，但其实我们完全可以在进程内完成WebSocket的测试，不用依赖外部网络。下面给你具体的解决方法和代码示例：

首先，Rocket的异步本地客户端支持处理WebSocket的升级请求，我们只需要在测试中添加标准的WebSocket握手头，然后通过.upgrade()方法获取到WebSocket会话，就能直接和你的路由进行消息交互了。

完整测试代码示例 #

#[cfg(test)]
mod tests {
    use rocket::local::asynchronous::Client;
    use rocket_ws::{WebSocket, Message};
    use super::*;

    #[rocket::async_test]
    async fn test_websocket_echo() {
        // 1. 构建并启动Rocket测试实例
        let rocket = rocket::build()
            .mount("/", routes![super::ws_route]);
        let client = Client::tracked(rocket).await
            .expect("Failed to initialize test client");

        // 2. 发起带有WebSocket握手头的升级请求
        let mut ws_session = client.get("/some_route")
            .header("Connection", "Upgrade")
            .header("Upgrade", "websocket")
            .header("Sec-WebSocket-Version", "13")
            .header("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") // 标准测试用密钥
            .upgrade()
            .await
            .expect("Failed to upgrade to WebSocket connection");

        // 3. 测试消息发送与接收（这里假设你的路由是回声服务）
        let test_msg = "Hello from WebSocket test!";
        ws_session.send(Message::Text(test_msg.into())).await
            .expect("Failed to send test message");

        // 接收服务器返回的消息
        if let Some(received) = ws_session.recv().await {
            match received {
                Ok(Message::Text(echoed)) => assert_eq!(echoed, format!("Echo: {}", test_msg)),
                Ok(_) => panic!("Expected text message, got other type"),
                Err(e) => panic!("Error receiving message: {}", e),
            }
        } else {
            panic!("WebSocket connection closed unexpectedly");
        }

        // 4. 主动关闭连接
        ws_session.close().await.expect("Failed to close WebSocket");
    }
}

关键细节说明 #

• 握手头的必要性：这些头是WebSocket协议要求的标准升级头，缺少任何一个都会导致服务器拒绝升级请求。其中Sec-WebSocket-Key只是用于验证握手的随机字符串，测试里用标准示例值即可，服务器会自动处理验证。
• .upgrade()方法：和普通请求用.dispatch()不同，.upgrade()会完成WebSocket的握手流程，直接返回可用的WebSocket会话实例，让你可以和路由中的WebSocket句柄直接交互。
• 异步测试宏：一定要用#[rocket::async_test]标记测试函数，它会自动处理Rocket异步测试所需的运行时环境。

对应路由示例（供参考） #

如果你的WebSocket路由是回声服务，大概是这样的：

use rocket_ws::{WebSocket, Channel, Message};

#[get("/some_route")]
async fn ws_route(mut ws: WebSocket) -> Channel<'static> {
    while let Some(msg) = ws.recv().await {
        match msg {
            Ok(Message::Text(text)) => {
                let _ = ws.send(Message::Text(format!("Echo: {}", text))).await;
            }
            Ok(_) => continue,
            Err(_) => break,
        }
    }
    ws.channel()
}

这种方式完全在进程内完成测试，没有网络开销，速度很快，而且能完整覆盖WebSocket路由的业务逻辑。

内容来源于stack exchange