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:
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | <?xml version="1.0" encoding="UTF-8"?> <project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd"> <modelVersion>4.0.0</modelVersion> <parent> <groupId>org.springframework.boot</groupId> <artifactId>spring-boot-starter-parent</artifactId> <version>2.7.0</version> <relativePath/> <!-- lookup parent from repository --> </parent> <groupId>de.rieckpil.blog</groupId> <artifactId>spring-websocket-integration-tests</artifactId> <version>0.0.1-SNAPSHOT</version> <name>spring-websocket-integration-tests</name> <description>Demo project for Spring Boot</description> <properties> <java.version>17</java.version> <awaitility.version>4.2.0</awaitility.version> </properties> <dependencies> <dependency> <groupId>org.springframework.boot</groupId> <artifactId>spring-boot-starter-websocket</artifactId> </dependency> <dependency> <groupId>org.springframework.boot</groupId> <artifactId>spring-boot-starter-test</artifactId> <scope>test</scope> </dependency> <dependency> <groupId>org.awaitility</groupId> <artifactId>awaitility</artifactId> <version>${awaitility.version}</version> <scope>test</scope> </dependency> </dependencies> <build> <plugins> <plugin> <groupId>org.springframework.boot</groupId> <artifactId>spring-boot-maven-plugin</artifactId> </plugin> </plugins> </build> </project> |
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:
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | @Configuration @EnableWebSocketMessageBroker public class WebSocketConfig implements WebSocketMessageBrokerConfigurer { @Override public void configureMessageBroker(MessageBrokerRegistry config) { config.enableSimpleBroker("/topic"); config.setApplicationDestinationPrefixes("/app"); } @Override public void registerStompEndpoints(StompEndpointRegistry registry) { registry.addEndpoint("/ws-endpoint") .withSockJS(); } } |
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
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | @Controller public class GreetingController { @MessageMapping("/welcome") @SendTo("/topic/greetings") public String greeting(String payload) { System.out.println("Generating new greeting message for " + payload); return "Hello, " + payload + "!"; } @SubscribeMapping("/chat") public Message sendWelcomeMessageOnSubscription() { Message welcomeMessage = new Message(); welcomeMessage.setMessage("Hello World!"); return welcomeMessage; } } |
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:
| 1 2 3 4 5 6 7 8 | @SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT) class GreetingControllerTest { @LocalServerPort private Integer port; // further setup to come } |
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:
| 1 2 3 4 5 | @BeforeEach void setup() { this.webSocketStompClient = new WebSocketStompClient(new SockJsClient( List.of(new WebSocketTransport(new StandardWebSocketClient())))); } |
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:
| 1 2 3 4 | StompSession session = webSocketStompClient .connect(String.format("ws://localhost:%d/ws-endpoint", port), new StompSessionHandlerAdapter() { }) .get(1, SECONDS); |
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.
| 1 2 3 4 | // pick one or use the default SimpleMessageConverter webSocketStompClient.setMessageConverter(new StringMessageConverter()); webSocketStompClient.setMessageConverter(new ByteArrayMessageConverter()); webSocketStompClient.setMessageConverter(new MappingJackson2MessageConverter()); |
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:
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | @Test void verifyGreetingIsReceived() throws Exception { BlockingQueue<String> blockingQueue = new ArrayBlockingQueue<>(1); webSocketStompClient.setMessageConverter(new StringMessageConverter()); StompSession session = webSocketStompClient .connect(getWsPath(), new StompSessionHandlerAdapter() {}) .get(1, SECONDS); session.subscribe("/topic/greetings", new StompFrameHandler() { @Override public Type getPayloadType(StompHeaders headers) { return String.class; } @Override public void handleFrame(StompHeaders headers, Object payload) { blockingQueue.add((String) payload); } }); session.send("/app/welcome", "Mike"); await() .atMost(1, SECONDS) .untilAsserted(() -> assertEquals("Hello, Mike!", blockingQueue.poll())); } |
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:
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | @Test void verifyWelcomeMessageIsSent() throws Exception { CountDownLatch latch = new CountDownLatch(1); webSocketStompClient.setMessageConverter(new MappingJackson2MessageConverter()); StompSession session = webSocketStompClient .connect(getWsPath(), new StompSessionHandlerAdapter() { }) .get(1, SECONDS); session.subscribe("/app/chat", new StompFrameHandler() { @Override public Type getPayloadType(StompHeaders headers) { return Message.class; } @Override public void handleFrame(StompHeaders headers, Object payload) { latch.countDown(); } }); await() .atMost(1, SECONDS) .untilAsserted(() -> assertEquals(0, latch.getCount())); } |
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:
- Spring Boot Unit and Integration Testing Overview
- Spring Boot Test Slices: Overview and Usage
- Testing Spring Boot Applications: Five Common Pitfalls
The source code for this example is available on GitHub.
Joyful testing,
Philip
