Write Integration Tests For Your Spring WebSocket Endpoints

Last Updated:  July 13, 2022 | Published: July 22, 2020

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

Spring Boot Application Setup for WebSockets

As there's a Spring Boot Starter available for WebSockets, our application setup is as simple as the following:

When testing our application, the Spring Boot Starter Test (aka. testing swiss army knife) already includes a set of testing libraries like JUnit, AssertJ, JsonPath, etc.

In addition to this essential toolbox, we're adding Awaitlity to the mix (one of the many great testing tools and libraries of the Java testing ecosystem).

Compared to testing a regular controller endpoint with @WebMvcTest, where we get access to the response synchronously, our WebSocket responses are asynchronous. Awaitlity helps us test such async Java code.

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 can 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) we send to our WebSocket endpoints, we 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 use Awaitility for our test verification.

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 we 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. Within one second, we expect that the count of our latch is zero. Otherwise, the Awaitility verification fails the test.

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

Further Spring Boot testing-related content is available here:

The source code for this example is available on GitHub.

Joyful testing,

Philip

>