Spring MVC环境下如何使用JUnit4编写javax.websocket单元测试

可以沿用「Mock核心依赖+无需启动完整外部容器」的测试思路，不过你使用的是JSR 356标准的原生@ServerEndpoint注解实现WebSocket，没有用到Spring封装的WebSocket模块，所以无法直接复用MockMvc工具，有两种落地方式可选：

方案一：纯单元测试（无需启动容器，执行速度最快） #

直接Mock Session等WebSocket容器提供的对象，调用端点的生命周期方法验证业务逻辑即可，不需要启动任何服务。

依赖引入 #

如果项目没有引入Mockito，先添加测试依赖：

<dependency>
    <groupId>org.mockito</groupId>
    <artifactId>mockito-core</artifactId>
    <version>2.28.2</version>
    <scope>test</scope>
</dependency>

测试代码示例 #

package com.ws.socket;

import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;
import javax.websocket.RemoteEndpoint;
import javax.websocket.Session;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class SocketUnitTest {
    // Mock容器提供的Session对象
    @Mock
    private Session session;
    // Mock消息发送工具对象
    @Mock
    private RemoteEndpoint.Basic basicRemote;

    @Test
    public void testOnMessageLogic() throws Exception {
        // 构造测试参数
        String testSessionId = "test-socket-001";
        String testMsg = "hello websocket";

        // 绑定mock方法的返回值
        when(session.getId()).thenReturn(testSessionId);
        when(session.getBasicRemote()).thenReturn(basicRemote);

        // 直接调用WebSocket端点的生命周期方法
        Socket socket = new Socket();
        socket.onOpen(session);
        socket.onMessage(session, testMsg);

        // 验证逻辑是否符合预期：是否给客户端返回了指定格式的消息
        verify(basicRemote, times(1)).sendText(" already receive msg about your say " + testMsg);
    }
}

该方案适合验证核心业务逻辑，执行速度最快，优先选择该方案测试功能。

方案二：集成测试（启动嵌入式Tomcat，模拟真实连接全流程） #

如果需要覆盖连接建立、握手、消息收发的完整流程，可以在测试中启动嵌入式Tomcat，注册WebSocket端点后用客户端发起真实请求测试，和你现有的Spring测试逻辑可以对齐。

依赖引入 #

<dependency>
    <groupId>org.apache.tomcat.embed</groupId>
    <artifactId>tomcat-embed-core</artifactId>
    <version>8.5.87</version>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.apache.tomcat</groupId>
    <artifactId>tomcat-websocket</artifactId>
    <version>8.5.87</version>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>javax.websocket</groupId>
    <artifactId>javax.websocket-client-api</artifactId>
    <version>1.1</version>
    <scope>test</scope>
</dependency>

测试代码示例 #

package com.ws.socket;

import com.ws.config.ServletConfig;
import com.ws.config.SpringConfig;
import org.apache.catalina.Context;
import org.apache.catalina.startup.Tomcat;
import org.apache.tomcat.websocket.server.WsContextListener;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
import org.springframework.test.context.web.WebAppConfiguration;
import javax.websocket.*;
import java.net.URI;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertEquals;

@WebAppConfiguration
@ContextConfiguration(classes = {ServletConfig.class, SpringConfig.class})
@RunWith(SpringJUnit4ClassRunner.class)
public class SocketIntegrationTest {
    private Tomcat tomcat;
    private final int port = 8082;
    private final String wsEndpointUrl = "ws://localhost:" + port + "/MyEndpoint";

    @Before
    public void startTomcat() throws Exception {
        // 启动嵌入式Tomcat
        tomcat = new Tomcat();
        tomcat.setPort(port);
        Context context = tomcat.addContext("", System.getProperty("java.io.tmpdir"));
        // 注册WebSocket支持
        context.addApplicationListener(WsContextListener.class.getName());
        // 注册自定义WebSocket端点
        ServerEndpointConfig endpointConfig = ServerEndpointConfig.Builder.create(Socket.class, "/MyEndpoint")
                .configurator(new ServerEndpointConfig.Configurator() {})
                .build();
        context.getServletContext().setAttribute("javax.websocket.server.ServerContainer", endpointConfig);
        tomcat.start();
    }

    @Test
    public void testFullWebSocketFlow() throws Exception {
        CountDownLatch latch = new CountDownLatch(1);
        String testMsg = "test integration msg";
        String[] receivedMsg = new String[1];

        // 构造WebSocket客户端发起连接
        ClientEndpointConfig clientConfig = ClientEndpointConfig.Builder.create().build();
        WebSocketContainer container = ContainerProvider.getWebSocketContainer();
        Session session = container.connectToServer(new Endpoint() {
            @Override
            public void onOpen(Session session, EndpointConfig config) {
                session.addMessageHandler(String.class, message -> {
                    receivedMsg[0] = message;
                    latch.countDown();
                });
                try {
                    session.getBasicRemote().sendText(testMsg);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }, clientConfig, URI.create(wsEndpointUrl));

        // 等待接收返回结果，超时时间3秒
        boolean received = latch.await(3, TimeUnit.SECONDS);
        session.close();

        assertEquals(true, received);
        assertEquals(" already receive msg about your say " + testMsg, receivedMsg[0]);
    }

    @After
    public void stopTomcat() throws Exception {
        if (tomcat != null) {
            tomcat.stop();
            tomcat.destroy();
        }
    }
}

该方案适合测试协议层、容器适配相关的逻辑，执行速度稍慢。

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