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:
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 | <?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.3.0.RELEASE</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>11</java.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> <exclusions> <exclusion> <groupId>org.junit.vintage</groupId> <artifactId>junit-vintage-engine</artifactId> </exclusion> </exclusions> </dependency> </dependencies> <build> <plugins> <plugin> <groupId>org.springframework.boot</groupId> <artifactId>spring-boot-maven-plugin</artifactId> </plugin> </plugins> </build> </project> |
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:
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 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
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 public 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) 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.
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 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:
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 | @Test public 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) { System.out.println("Received message: " + payload); blockingQueue.add((String) payload); } }); session.send("/app/welcome", "Mike"); assertEquals("Hello, Mike!", blockingQueue.poll(1, SECONDS)); } |
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.
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 public 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(); } }); if (!latch.await(1, TimeUnit.SECONDS)) { fail("Message not received"); } } |
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 is 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
