如何正确Mock WebSocket会话？Spring Boot应用单元测试求助

嘿，作为Mockito新手来测试WebSocket Handler确实得摸清楚门道，我来一步步带你搞定这个单元测试！首先咱们明确要测试的核心行为：当WebSocket会话处于打开状态时，handleTextMessage方法是否会发送指定的响应消息；当会话关闭时，是否不会执行发送操作。

第一步：准备依赖 #

确保你的pom.xml（或build.gradle）里包含了JUnit 5和Mockito的必要依赖：

<!-- JUnit 5 -->
<dependency>
    <groupId>org.junit.jupiter</groupId>
    <artifactId>junit-jupiter-api</artifactId>
    <version>5.9.2</version>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.junit.jupiter</groupId>
    <artifactId>junit-jupiter-engine</artifactId>
    <version>5.9.2</version>
    <scope>test</scope>
</dependency>

<!-- Mockito -->
<dependency>
    <groupId>org.mockito</groupId>
    <artifactId>mockito-core</artifactId>
    <version>4.11.0</version>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.mockito</groupId>
    <artifactId>mockito-junit-jupiter</artifactId>
    <version>4.11.0</version>
    <scope>test</scope>
</dependency>

第二步：编写单元测试类 #

咱们不需要启动Spring上下文，纯单元测试就够了，用Mockito模拟WebSocketSession对象，验证它的方法调用行为：

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;

import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.*;

@ExtendWith(MockitoExtension.class)
class TextHandlerTest {

    // 用Mockito模拟WebSocketSession对象
    @Mock
    private WebSocketSession session;

    // 实例化要测试的TextHandler
    private TextHandler textHandler = new TextHandler();

    @Test
    void handleTextMessage_WhenSessionIsOpen_SendsHelloMessage() throws InterruptedException, IOException {
        // 1. 指定模拟的session.isOpen()返回true
        when(session.isOpen()).thenReturn(true);

        // 2. 调用要测试的方法，传入模拟的session和任意测试消息
        textHandler.handleTextMessage(session, new TextMessage("test message"));

        // 3. 验证session的sendMessage方法被调用，且参数是预期的TextMessage
        verify(session).sendMessage(
                argThat(message -> ((TextMessage) message).getPayload().equals("Hello from the websocket..."))
        );
    }

    @Test
    void handleTextMessage_WhenSessionIsClosed_DoesNotSendMessage() throws InterruptedException, IOException {
        // 1. 指定模拟的session.isOpen()返回false
        when(session.isOpen()).thenReturn(false);

        // 2. 调用要测试的方法
        textHandler.handleTextMessage(session, new TextMessage("test message"));

        // 3. 验证session的sendMessage方法从未被调用
        verify(session, never()).sendMessage(any(TextMessage.class));
    }
}

关键步骤解释 #

• @ExtendWith(MockitoExtension.class)：让JUnit 5支持Mockito的注解，比如@Mock。
• @Mock：创建一个WebSocketSession的模拟对象，我们可以控制它的返回值，验证它的方法调用。
• when(session.isOpen()).thenReturn(true)：指定当调用session的isOpen()方法时，返回true（模拟会话打开的场景）。
• verify(session).sendMessage(...)：验证session的sendMessage方法是否被调用，并且传入了我们预期的消息内容。
• verify(session, never()).sendMessage(...)：验证当会话关闭时，sendMessage方法完全没有被调用过。

小提示 #

顺便提一句，你当前的handleTextMessage方法里的try块没有对应的catch分支，如果sendMessage抛出IOException会直接向上传播，实际项目中可以考虑添加异常处理逻辑，但这并不影响咱们的单元测试哦。

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