Go: Testing WebSockets

2023-07-10

How can you test your Golang websocket server?

The net/http/httptest provides utilities for testing your http server. It can also test upgraded connections such as a websocket connection.

I am a big proponent of using tests as a part of the development process, even in my spare time projects. I like to know ensure that my code does what it think it does. And especially when building a server tests are necessary to remove the need for building a client application to call your server API with.

The Server to test

Lets start by setting up a simple websocket server using the gorilla/websocket package. The server is as simple as it gets with only one route 'localhost:8080/' which is tied to the handler sckt which opens a websocket connection and reads messages from the client and sends it back to the client.

1 2package main 3 4import ( 5"log" 6"net/http" 7"github.com/gorilla/websocket" 8) 9 10//gorilla/websockets upgrader. used to upgrade the connection to websocket connection. 11var upgrader = &websocket.Upgrader{} 12 13func sckt(w http.ResponseWriter, req *http.Request){ 14 // returns an upgraded conn obj. a socket connection. 15 socket, err := upgrader.Upgrade(w,req,nil) 16 if err != nil { 17 log.Fatal("in function sckt: ", err) 18 } 19 // make sure to close the connection 20 defer socket.Close() 21 22 //Read and Write messages forever. Not very useful but serves the purpose for testing. 23 for { 24 //socket.ReadMessage is blocking 25 _,message,err := socket.ReadMessage() 26 if err != nil { 27 break 28 } 29 30 err = socket.WriteMessage(websocket.TextMessage,[]byte(message)) 31 if err != nil{ 32 break 33 } 34 } 35} 36 37func main(){ 38 // http server with one route. The websocket connection NOT USED IN TESTS 39 http.HandleFunc("/", sckt) 40 http.ListenAndServe("localhost:8080",nil) 41} 42

Testing our server

So how do you go about testing this? First of all we will use the testing package to write a test function. Inside the test function we will use the net/http/httptest to setup a testserver which takes the sckt handler function and then makes websocket request to with the websocket.Dialer.Dial function. Everything here is provided by the standard library except the gorilla websocket connection itself.

Lets First look at the code and then go through what is does

1 2 3package main 4 5import ( 6 "net/http" 7 "net/http/httptest" 8 "net/url" 9 "strings" 10 "testing" 11 12 "github.com/gorilla/websocket" 13) 14 15func TestWebsck(t *testing.T) { 16 17 server := httptest.NewServer(http.HandlerFunc(sckt)) 18 defer server.Close() 19 20 scktRouteUrl, _ := url.Parse(server.URL) 21 scktRouteUrl.Scheme = "ws" 22 23 header := http.Header{} 24 25 conn, _, err := websocket.DefaultDialer.Dial(scktRouteUrl.String(), header) 26 if err != nil { 27 t.Fatalf("%v", err) 28 } 29 30 defer conn.Close() 31 32 myMessage := "testing!" 33 34 err = conn.WriteMessage(websocket.TextMessage, []byte(myMessage)) 35 if err != nil { 36 t.Fatalf("%v",err) 37 } 38 39 _,message, err := conn.ReadMessage() 40 if err != nil { 41 t.Fatalf("%v", err) 42 } 43 44 if strings.Compare(string(message), myMessage) != 0 { 45 t.Fatalf("Expected something, but got %s", string(message)) 46 } 47} 48

The error checks make the code look slightly cluttered but there really is not much going on. We start by creating a test server server := httptest.NewServer(http.HandlerFunc(sckt)). Next we create a client request to the url of the test server. We make sure that we make a ws request and not http request by changing the Scheme scktRouteUrl.Scheme = "ws". The Dial returns a *websocket.conn object, a *http.Request and an error. We are interested in the *websocket.conn which we will use to send messages to the our server with.

Since we specified the conn, _, err := websocket.DefaultDialer.Dial(scktRouteUrl.String(), header) with the url of our scktRouteUrl. The conn object is tied the the route of our sckt handler function. The conn acts if it were a web client and we can write and read on it to test our sckt route.

To test the route we do one write to the server and then one read. We know that read calls are blocking, so we need to make sure that we do things the way the server expects. It first expects a write by the client and then a read from the client. This is done with the lines

1 2 myMessage := "testing!" 3 4 err = conn.WriteMessage(websocket.TextMessage, []byte(myMessage)) 5 if err != nil { 6 t.Fatalf("%v",err) 7 } 8 9 _,message, err := conn.ReadMessage() 10 if err != nil { 11 t.Fatalf("%v", err) 12 } 13

Finally we compare the data sent and the data recieved to make sure they are equal.

1 if strings.Compare(string(message), myMessage) != 0 { 2 t.Fatalf("Expected something, but got %s", string(message)) 3 }

And that concludes our unit test!

This example was very simple and for this particular route the test does not provide much value but we can imagine how to use this testing strategy for more complicated connections. The general strategy is to first create a test server, create a client connection and then send messages to the route as specified by the server. It is a good first step for testing our routes

That said this approach is not perfect and i can see weaknesses. One of which is if our server has routes that depend on some state on the server and on the client. Then we would need some sort of mocking to create tests. Another flaw in the example is that we create a new server in the beginning of the test. Creating a new test would require creaing a new testing server. However this can easily be fixed with a refactor, either by using a "before every" setup function for the tests or store a test server in some global state either by wrapping the testing.T in our own struct or just creating a global variable in our test file might be enough for simple use cases.