Testing WebSocket in Go

It’s been a while since my last post in this blog. I also posted some posts in my Medium while I was in the Gojek Bootcamp. Feel free to check it out here.

A few weeks ago, I got this idea for making real-time chat server with WebSocket. I was curious because I had never been playing with WebSocket in Go previously.

Writing test for WebSocket

The idea behind why I wrote the test for WebSocket was that I wanted to try end-to-end test in Go. It was something that I had never been done before. I used to not writing test and tested the code end-to-end manually. And it was such a PITA.

For the test, I used this library simply because it is very straightforward and provides the sample code to get started. The next thing I wanted to achieve was to test concurrent connections. So, I wrote the test like this:

// ...
numberOfClients := 3

for i := 1; i <= numberOfClients; i++ {
  userID := i * 100
  t.Run(fmt.Sprintf("clients_%d", userID), func (t *testing.T) {
    t.Parallel()
    // connect user to websocket server
  })
}
// ...

However, when you are doing testing in parallel, you should be aware of the maximum number of tests running in parallel. The default limit is set to the value of GOMAXPROCS which is 4 in my laptop. To override this value, you can specify it by passing -parallels <N> flag when executing go test.

$ go test -parallel 10 ./...

I used subtest to evaluate each client for expected behavior. The subtests are running in parallel and grows as the number of clients used for testing increase (I manually set the number of clients). However, as previously I didn’t specify the -parallels <N> flag, the test was stucked because the clients and the server were not running simultaneously. Hence, it caused deadlock as most of the time the server goroutine was not running and the clients failed to connect to the server.

Closing connection gracefully

Later, I also found that the c.Close() to close the WebSocket connection didn’t work as I expected. Turned out it closes the connection directly without sending CLOSE packet to inform the server that the connection is going to be closed. Then I found the procedure to send CLOSE packet here. However, I used WriteControl instead of WriteMessage because WriteControl is more suitable for handling control packets.

// ...
d := wstest.NewDialer(e)
c, resp, err := d.Dial("ws://whatever/messages/listen", requestHeader)

// ...
defer func() {
  c.WriteControl(
    websocket.CloseMessage,
    websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
    time.Now().Add(time.Second),
  )
  c.Close()
}()
// ...

Summary

I learned a lot around testing in Go when playing with this project. Part of the learnings are how parallel testing works, how it is capped by default maximum number of test running simultaneously, and how to enable printing by passing -v flag to go test command.

However, this is not the end yet. Because I just found other problems in building WebSocket server. But it’s a story for another post.

Sample code

You can check the complete source-code here.


References