How to Receive Response from Websocket Unit Test in Springboot

Question

I am new to websockets and I am trying to write a unit test.

My unit test runs fine but it has following two issue

1. Idk why but it forces me to expect same object that is being sent as an input(i.e WebSocketRequestData) to the websocket instead of the actual response from the websocket which is WebSocketData

2. And it returns an empty object as result so it passes NotNull assertion.

Can anyone please clear out this confusion for me!

And also what is the right way to get response from the my websocket in unit test?

here is the code for my websocketTest Class

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

@LocalServerPort
private Integer port;

static final String WEBSOCKET_TOPIC = "/user/locationrealtimedata/item" ;

BlockingQueue<WebSocketRequestData> blockingQueue;
WebSocketStompClient stompClient;

@BeforeEach
public void setup() {

    blockingQueue = new LinkedBlockingDeque<>();
    stompClient = new WebSocketStompClient(new SockJsClient(
            asList(new WebSocketTransport(new StandardWebSocketClient()))));
    
    stompClient.setMessageConverter(new MappingJackson2MessageConverter());

}

@Test
public void shouldReceiveAMessageFromTheServer() throws Exception {
    StompSession session = stompClient
            .connect(getWsPath(), new DefaultStompFrameHandler() {
            })
            .get(1, TimeUnit.SECONDS);
    session.subscribe(WEBSOCKET_TOPIC, new DefaultStompFrameHandler());

    WebSocketRequestData webSocketRequestData = new WebSocketRequestData();
    webSocketRequestData.setUserId("usr-1");
    webSocketRequestData.setAccountId("acc-1");
    webSocketRequestData.setGroupId("grp-1");
    session.send("/wsconn/start", webSocketRequestData);
    WebSocketRequestData responseObj = blockingQueue.poll(15, TimeUnit.SECONDS);
    Assertions.assertNotNull(responseObj);

}

class DefaultStompFrameHandler extends StompSessionHandlerAdapter{   
    @Override
    public Type getPayloadType(StompHeaders stompHeaders) {
        return WebSocketRequestData.class;
    }

    @Override
    public void handleFrame(StompHeaders stompHeaders, Object o) {
        blockingQueue.offer((WebSocketRequestData) o); // instead of **WebSocketData** it forces me to add casting for **WebSocketRequestData**
    }


    @Override
    public void handleException(StompSession session, StompCommand command, StompHeaders headers, byte[] payload, Throwable exception) {
        exception.printStackTrace();
    }


}

private String getWsPath() {
    return String.format("ws://localhost:%d/location_services/locationrealtimedata", port);
}

}

Thanks in advance

Answer 1

You are not forced to use the same Java class for the input and response type.

The request type is what you use within session.send("/endpoint", payload); in your case that's WebSocketRequestData:

WebSocketRequestData webSocketRequestData = new WebSocketRequestData();
webSocketRequestData.setUserId("usr-1");
webSocketRequestData.setAccountId("acc-1");
webSocketRequestData.setGroupId("grp-1");
session.send("/wsconn/start", webSocketRequestData);

When it comes to consuming messages you specify the actual response type you expect when implementing StompFrameHandler and overriding getPayloadType.

So instead of implementing StompSessionHandlerAdapter, use the StompFrameHandler interface and implement it as the following:

class DefaultStompFrameHandler extends StompSessionHandlerAdapter{   
    @Override
    public Type getPayloadType(StompHeaders stompHeaders) {
        return WebSocketData.class; // or any other class your expect
    }

    @Override
    public void handleFrame(StompHeaders stompHeaders, Object o) {
        blockingQueue.offer((WebSocketData) o); 
    }

    @Override
    public void handleException(StompSession session, StompCommand command, StompHeaders headers, byte[] payload, Throwable exception) {
        exception.printStackTrace();
    }

}

Also make sure your BlockingQueue is using the correct type BlockingQueue<WebSocketData> blockingQueue