Painless WebSocket Tests With Spock Framework

Many Spring Boot web applications require you to go beyond the usual REST endpoints and start accepting WebSocket connections. Meanwhile, plain JUnit tests have fallen out of fashion, and suddenly everyone is talking about the potential of Spock.

All these changes sound thrilling, and you’d like to try the novel approach.

Yet, if not applied properly, these techniques could backfire on you, resulting in flaky and hard-to-debug tests — the usual price for asynchronous operations. So, how can you prevent it?

Read on to see how easy and gratifying your WebSocket coverage can be once you learn how to target the usual suspects, like race conditions and false positives. To make the process even smoother, make sure to help yourself to the source code available on GitHub.

The Echo: If Only All Protocols Were That Simple

Start with setting up Spring to handle incoming connections. This task requires a handler bean and a configuration to associate the bean with an address.

The easiest way to go about a handler is extending one of the helper classes like TextWebSocketHandler. For the sake of argument, we will implement RFC 862: the Echo Protocol. One caveat is that we will ignore empty messages to make the service more interesting to test.

This is what such a handler can look like:

public class EchoSocket extends TextWebSocketHandler {

    @Override
    public void handleTextMessage(WebSocketSession session,
                                  TextMessage message) throws IOException {

        if (message.getPayload().isEmpty()) return;

        session.sendMessage(message);
    }
}

Then, you need a WebSocketConfigurer implementation to register your handler bean under the specific address. The above handler gets declared as a bean and attached to /echo:

@Configuration
@EnableWebSocket
public class SocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(
            WebSocketHandlerRegistry registry) {
        registry.addHandler(echo(), "/echo");
    }

    @Bean
    public EchoSocket echo() {
        return new EchoSocket();
    }
}

With these two classes in place, the app should accept WebSocket sessions at localhost:8080/echo and always reply with the received text. The next thing to do is start adding tests to see if it works.

First Contact: Connecting to the Socket

Spock does not offer an easy way to open WebSocket connections, so we need a client library for that:

<dependency>
    <groupId>com.neovisionaries</groupId>
    <artifactId>nv-websocket-client</artifactId>
    <version>2.8</version>
</dependency>

The first thing to validate is that a connection gets accepted. If we let Spring pick a random port to start on, we need to know the value — so we inject it with @LocalServerPort.

At this point, remember to force the immediate disconnect of the socket we’ve used in each of the tests. Otherwise, as the number and complexity of tests increase, so will the number of unclosed sockets interfering in unexpected ways:

@SpringBootTest(
        classes = SpringWsApplication.class, 
        webEnvironment = RANDOM_PORT)
class EchoSocketTest extends Specification {

    @LocalServerPort int port

    def "socket opens"() {
        when:
        def socket = new WebSocketFactory()
                .createSocket("http://localhost:$/echo")
                .connect()

        then:
        socket.isOpen()

        cleanup:
        socket.disconnect(WebSocketCloseCode.NORMAL, null, 0)
    }
}

So far, so good, the test goes through, so we know we can open and close the socket.

The matter becomes less obvious when we start sending and receiving messages.

Unlike plain HTTP calls, WebSocket communication is not conducted in a request-response fashion. You shouldn't expect a synchronous response when a text message gets sent to a socket.

Spock Interaction-Based Testing: Here Come the Mocks

Before we move forward, let’s take a look at two methods that will be used throughout the rest of the article:

def openSocket() {
    new WebSocketFactory()
            .createSocket("http://localhost:$/echo")
            .connect()
}

static def closeSocket(socket) {
    socket.disconnect(WebSocketCloseCode.NORMAL, null, 0)
}

Taking into account what we already know and looking at the WebSocket class from the client library, submitting a text message seems easy — sendText(String message). Yet, we still need to figure out a way to receive responses.

There is no receive() method, so it is essential to register a WebSocketListener instead.

And how do we validate that the listener was called? Spock documentation has a whole section on interaction-based testing, by which they mean using mocks:

def "responds with original message"() {
    given:
    def socket = openSocket()

    and:
    def listenerMock = Mock(WebSocketListener)
    socket.addListener(listenerMock)

    when:
    socket.sendText("Hello")

    then:
    1 * listenerMock.onTextMessage(_, "Hello")

    cleanup:
    closeSocket(socket)
}

So we create a mock, add it as a socket listener, submit a text message, and expect the mock to be called exactly once because this is how the Echo Protocol works. Then, we try to run the test, and surprisingly it fails:

Too few invocations for:

1 * listenerMock.onTextMessage(_, "Hello")   (0 invocations)

Unmatched invocations (ordered by similarity):

None

What could have possibly gone wrong? We are sending a text message and waiting for a response — are we really, though?

Immediately after socket.sendText(), the test goes to the next line to execute the assertion. It expects the mock to have already been called exactly once rather than wait for EchoSocket to receive the message and respond. Great — we got ourselves a race condition.

If you add a short sleep() right before the assertion, it is likely to pass — but then you are doing two things wrong:

1. You are making the test run longer than necessary; i.e., 100 milliseconds is most likely more than it needs.
2. You cannot be sure that 100 milliseconds will always be enough — every 10th or 1000th execution can still fail, and you and your team need to understand that this is the nature of this test. For that matter, it’s the same with all your WebSocket tests if you follow down this path.

BlockingVariable: The Ultimate Weapon Against Race Conditions

The issues and principles mentioned above prove a need for a synchronization mechanism. One that would coordinate the thread that executes the test with the one that calls our WebSocketListener after the response arrives.

Spock provides a construct called BlockingVariable, which is not explained in their documentation — possibly because it is fairly straightforward.

To keep it simple, we can describe this variable as a wrapper around the value of a given type. The value is set like with any other setter, but it can produce two different results when get() is called — it either returns the value immediately, if it has already been set, or waits for a given time for it to be set and then fails.

You can consider using BlockingVariable in the following way:

1. Instantiate it with the expected type and maximum time to wait for get() to return.
2. Implement a WebSocketListener that forwards any received text to the BlockingVariable.
3. Try to get the value in the assertions block. Once it arrives, you can also assert what the text is, and if it never arrives — the call to get() has all the right to fail the test.

def "responds with original message"() {
    given:
    def latch = new BlockingVariable<String>(0.5)

    and:
    def socket = openSocket()
    socket.addListener(new WebSocketAdapter() {
        @Override
        void onTextMessage(WebSocket websocket, 
                           String text) {
            latch.set(text)
        }
    })

    when:
    socket.sendText("Hello")

    then:
    with(latch.get()) {
        it == "Hello"
    }

    cleanup:
    closeSocket(socket)
}

This should do it — no race conditions anymore. The test is forced to wait for the thread that receives the message to respond. It also takes exactly as much time as it needs to. The moment the value is set, the method get() returns, and all remaining assertions get executed.

As your test suite grows, consider replacing the above and: block with a utility method that opens a socket and forwards it to a latch in one go. You don’t want boilerplate code like the one below distracting you from the logic behind the test:

def openSocket(latch) {
    openSocket().addListener(new WebSocketAdapter() {
        @Override
        void onTextMessage(WebSocket websocket, 
                           String text) {
            latch.set(text)
        }
    })
}

So there you have it — with your first meaningful and reliable test in place, you and your teammates have no more excuses. Now is the time to continue working on your automated test coverage and keep rolling out all the amazing WebSocket-based features!

Once again, feel free to use the source code available on GitHub, and happy testing!