How to Correctly Mock a WebSocket session?

Question

I have created a simple websocket application using springBoot. I am new to Mockito and I am trying to unit test the behaviour of following class with mockito and junit.

@Component
public class TextHandler extends TextWebSocketHandler {

WebSocketSession session;

@Override
public void handleTextMessage(WebSocketSession session, TextMessage message)
        throws InterruptedException, IOException {

        // send message
        if (session.isOpen()) {
            try {

                session.sendMessage(new TextMessage("Hello from the websocket"));
            } finally {
                session.close();
            }
        } else {
            System.out.println("no open session available");
        }
}

I created a stub for the session under test package as follows.

@Component
public class WebSocketSessionStub implements WebSocketSession{

@Override
public String getId() {

    return "SESSION1";
}

@Override
public URI getUri() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public HttpHeaders getHandshakeHeaders() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public Map<String, Object> getAttributes() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public Principal getPrincipal() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public InetSocketAddress getLocalAddress() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public InetSocketAddress getRemoteAddress() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public String getAcceptedProtocol() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public void setTextMessageSizeLimit(int messageSizeLimit) {
    // TODO Auto-generated method stub

}

@Override
public int getTextMessageSizeLimit() {
    // TODO Auto-generated method stub
    return 0;
}

@Override
public void setBinaryMessageSizeLimit(int messageSizeLimit) {
    // TODO Auto-generated method stub

}

@Override
public int getBinaryMessageSizeLimit() {
    // TODO Auto-generated method stub
    return 0;
}

@Override
public List<WebSocketExtension> getExtensions() {
    // TODO Auto-generated method stub
    return null;
}

@Override
public void sendMessage(WebSocketMessage<?> message) throws IOException {
    // TODO Auto-generated method stub

}

@Override
public boolean isOpen() {
    System.out.println("isOpen");
    return true;
}

@Override
public void close() throws IOException {
    // TODO Auto-generated method stub

}

@Override
public void close(CloseStatus status) throws IOException {
    // TODO Auto-generated method stub

   }

 }

Following is my unit test class.

@RunWith(SpringRunner.class)
@SpringBootTest
public class TextHandlerTest {

@Autowired
TextHandler textHandler;

@Mock
WebSocketSessionStub ws;

@Mock
WebSocketMessage<TextMessage> webSocketMessage;

TextMessage textMsg = new TextMessage("Test Message".getBytes());

@Before
public void setUp() throws Exception {
    MockitoAnnotations.initMocks(this);
    textHandler.handleMessage(ws, textMsg);     
}


@Test
public void verifyCallToIsOpenConnection() throws Exception {

    verify(ws, times(1)).isOpen();
    System.out.println("isOpen " + ws.isOpen());

 }

}

Above test passes, however isOpen evaluates to false. Therefore I cannot verify the sendMessage method call. How can I rectify it?

Answer 1

You can mock WebSocketSession and inject it into TextHandler#handleTextMessage.

/*  
 *  You don't need spring context while mocking. 
 *  Hence No need to have spring runner. 
 */
@RunWith(MockitoJunitRunner.class)
public class TextHandlerTest {

    @Test
    public void verifyCallToIsOpenConnection() {
         WebSocketSession session = mock(WebSocketSession.class);
         TextMessage textMsg = new TextMessage("Test Message".getBytes());

         when(session.isOpen()).thenReturn(true);

         TextHandler textHandler = new TextHandler();

         // Pass the mocked session object here
         textHandler. handleTextMessage(session, textMsg);

         // Now you can verify if session.sendMessage() was called or not
         verify(session, times(1)).sendMessage(textMsg);
    }
}