如何为Spring TextWebSocketHandler编写连接及消息发送集成测试？

前置依赖 #

确保项目测试环境已引入WebSocket和Spring Boot测试依赖，Maven配置示例如下：

<dependencies>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-websocket</artifactId>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-test</artifactId>
        <scope>test</scope>
    </dependency>
</dependencies>

集成测试代码 #

import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.web.server.LocalServerPort;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import static org.junit.jupiter.api.Assertions.*;

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
public class WebSocketIntegrationTest {

    @LocalServerPort
    private Integer port;

    private final CountDownLatch messageLatch = new CountDownLatch(1);
    private String receivedPayload;

    @Test
    public void verifyWebSocketConnectAndMessageSend() throws Exception {
        StandardWebSocketClient webSocketClient = new StandardWebSocketClient();
        
        // 发起WebSocket连接请求
        WebSocketSession session = webSocketClient.doHandshake(new TextWebSocketHandler() {
            @Override
            protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
                receivedPayload = message.getPayload();
                messageLatch.countDown();
                session.close();
            }
        }, "ws://localhost:" + port + "/data").get(3, TimeUnit.SECONDS);

        // 验证连接成功建立
        assertTrue(session.isOpen());
        // 验证3秒内收到服务端推送消息
        assertTrue(messageLatch.await(3, TimeUnit.SECONDS), "未接收到服务端推送的消息");
        // 验证消息内容与Message.sendMessage()返回值匹配
        assertEquals("text", receivedPayload);
    }
}

测试逻辑说明 #

• 测试启动时会启动完整的Spring Boot Web服务，使用随机端口避免端口冲突
• 初始化标准WebSocket客户端向/data端点发起连接，验证连接是否可以正常建立
• 客户端接收到服务端推送的第一条消息后，校验内容是否符合预期，校验完成自动关闭连接

现有代码优化建议 #

你当前SocketHandler的afterConnectionEstablished方法中使用了死循环直接推送消息，会占用WebSocket核心处理线程，导致后续连接无法正常接入，建议将循环推送逻辑提交到独立的线程池执行，避免阻塞核心线程。

内容的提问来源于stack exchange，提问作者jjakubos