Write Integration Tests For Your Spring WebSocket Endpoints

WebSockets allow establishing a full-duplex, two-way communication between the client and the server. With Spring WebSocket you can bootstrap your WebSocket application with ease. While there is great test support available for verifying Spring MVC endpoints (e.g. @RestController) using MockMvc, the support for testing WebSockets is rather limited. Continue reading to get an introduction to writing integration tests for Spring WebSocket endpoints.

Spring Boot application setup for WebSockets

As the Spring Boot team provides a Spring Boot Starter for WebSockets, our application setup is as simple as the following:

When it comes to testing the application, the swiss army knife Spring Boot Starter Test already provides all test dependencies we need.

Defining WebSocket endpoints with Spring WebSocket

The application defines a somewhat standard configuration for using WebSockets:

We are making use of an in-memory broker. Subscriptions to /topic pass through the clientInboundChannel and are forwarded to this broker. Apart from this, subscriptions to /app are application destinations and reach our WebSocket controllers. There is a great visualization available at the Spring WebSocket documentation that explains the difference between these two endpoints.

Our clients are able to establish the WebSocket connection at /ws-endpoint. In addition, we enable SockJS support as a fallback option.

For this demo, the application has two use cases:

• greet new clients when they send a message to /app/welcome
• send a welcome message to clients subscribing to /app/chat

With the @SendTo annotation we tell Spring to send the return value after passing the brokerChannel to /topic/greetings (in-memory broker).

Integration test setup for testing Spring WebSocket endpoints

We can make use of the @SpringBootTest annotation to populate the whole application context and start the embedded servlet container. This is important as we actually want to establish a real connection to our locally running Spring Boot application:

While we can use MockMvc or TestRestTemplate/WebTestClient for accessing regular Spring MVC endpoints, we have to manually create a client for our WebSocket test.

As our application makes use of the STOMP protocol (more information here ), we can use the WebSocketStompClient from Spring for this. The constructor of this WebSocket client expects an WebSocketClient instance. Here we can use a SockJsClient from Spring, as our application setup allows the SockJs fallback option:

This setup is all we need to create our WebSocket client. The actual connection to our running Tomcat is established once we call .connect() on the client. Here we have to pass the URL and can add a handler that is notified once the CONNECTED frame is received.

For this example, the handler is doing nothing as we don't have logic to test right after establishing the connection:

We'll use this session in the next chapter to verify the behavior of our application.

Depending on what payload (bytes, plain Strings, JSON) you send to your WebSocket endpoints, you have to configure the message converter of the WebSocketStompClient. The default is SimpleMessageConverter which might not fit for every use case.

Integration tests for the sample Spring WebSocket application

Let's start verifying the WebSocket endpoints of our Spring Boot application.

For the first use case, we can expect a greeting coming from /topic/greetings, whenever a client sends a message to the destination /app/welcome. As this is an asynchronous process, we can either use Awaitility or a data type that supports waiting on elements.

What's left is to subscribe to the topic and implement a basic StompFrameHandler that is capable of handling the payload we receive:

The BlockingQueue allows us to poll an element from the queue while waiting up to one second. As this endpoint expects a String, I'm configuring the correct message converter at the beginning of the test.

Next, let's use a different technique to verify our second use case. As we expect a welcome message to be sent to our client whenever a subscription happens for /app/chat, we can use a CountDownLatch for this. While you usually use this class when working with threads as a synchronization aid, it can also help us here.

We initialize the latch with a count of one. When we receive a STOMP frame, we use countDown to reduce the count of the latch. The verification happens with latch.await().  This returns true if the latch has counted down to zero. If that does not happen in a timeframe of one second, we fail the test.

For further examples on how to write unit and integration tests for Spring WebSocket, take a look at the Spring WebSocket Portfolio application.

You can find the source code for this example on GitHub.

Have fun writing integration tests for your Spring WebSocket endpoints,

Phil