Testing in Go: Clean Tests Using t.Cleanup

Table of Contents

Go v1.14 ships with improvements across different aspects of the language. Two of them are brand new t.Cleanup, and b.Cleanup methods, added to the testing package.

The introduction of these methods will make it easier for our tests to clean up after themselves. This was always achievable through the careful composition of (sub)tests and helper functions, but since Go 1.14, the testing package will ship with one right way to do that.

Let’s explore through a scenario, one very common in web development, how these methods work, and how to put them in use.

In this article, we will focus on the t.Cleanup function, but the points made here apply to the b.Cleanup function too.

Vulnerable Authentication #

A web service that we own has an authentication middleware that uses HTTP Basic Authentication. The middleware takes the Authorization header values of an HTTP request using the BasicAuth helper method and authenticates the request. Depending on the authentication result, it will return an error response or let the request go through.

The AuthMiddleware authorization middleware:

type AuthMiddleware struct {
	db *gorm.DB
}

func (am *AuthMiddleware) Validate(username, password string) bool {
	var u User

	if err := am.db.Where("username = ? AND password = ?", username, password).First(&u).Error; err != nil {
		return false
	}

	return true
}

func (am *AuthMiddleware) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if u, p, ok := r.BasicAuth(); ok && am.Validate(u, p) {
			next.ServeHTTP(w, r)
		} else {
			http.Error(w, "Forbidden", http.StatusForbidden)
		}
	})
}

The AuthMiddleware is a struct that has a db attribute of the type *gorm.DB, which, in fact, is a Gorm database connection pool. The Validate method of the middleware will check the validity of the authorization credentials sent by the client. (Yes, this is a very dumb and vulnerable service - stores plain text passwords in the database. No, you should never keep your customers’ passwords in plain text.)

To do that, it uses the Request.BasicAuth method that returns the username and password provided in the request’s Authorization header. Then, we invoke the Validate method with u (the username) and p (the password) as arguments, that returns a bool.

When the request successfully authorized, the response of the endpoint will be returned. If the request does not supply proper HTTP Basic Authentication credentials, or the supplied credentials are invalid, it will return an HTTP 403 Forbidden error.

The simplistic server that builds the router and mounts the handler and the middleware:

func main() {
	db, err := gorm.Open("sqlite3", "./cleanups.db")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	aum := &AuthMiddleware{db}
	r := mux.NewRouter()
	r.Use(aum.Middleware)
	r.HandleFunc("/foo", func(w http.ResponseWriter, r *http.Request) {
                io.WriteString(w, "6 x 9 = 42\n")
	})

	log.Fatal(http.ListenAndServe(":8888", r))
}

The main function opens a Gorm connection to an SQLite3 database, builds a mux router, and mounts a handler at the /foo URI. Then, it mounts the router on a server and starts listening on 127.0.0.1:8888.

So, a sensible question would be, how can we test such a middleware? Since there is a database connection in play here, we have to make sure the database contains the actual username and password we want to send in the test.

Testing the Vulnerable Authentication #

Given that the authorization middleware expects a database connection, testing it can be a messy exercise. To test it, we will have to open a database connection (to a test database) and then use it to initialize a test server that mounts the router. Then, we have to send requests to that same server and run our assertions on its response body and status.

Let’s take a stab at it:

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/jinzhu/gorm"
)

func TestServer(t *testing.T) {
	db, err := gorm.Open("sqlite3", "./cleanups_test.db")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	r := Router(db)

	tcs := []struct {
		name           string
		username       string
		password       string
		responseBody   string
		responseStatus int
	}{
		{
			name:           "with invalid username-password combo",
			username:       "jane",
			password:       "doe",
			responseBody:   "Forbidden",
			responseStatus: http.StatusForbidden,
		},
		{
			name:           "with valid username-password combo",
			username:       "jane",
			password:       "doe123",
			responseBody:   "6 x 9 = 42",
			responseStatus: http.StatusOK,
		},
	}

	ts := httptest.NewServer(r)
	defer ts.Close()
	client := ts.Client()

	req, err := http.NewRequest("GET", fmt.Sprintf("%s/foo", ts.URL), nil)
	if err != nil {
		t.Fatal(err)
	}

	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			req.SetBasicAuth(tc.username, tc.password)

			res, err := client.Do(req)
			if err != nil {
				t.Fatal(err)
			}
			defer res.Body.Close()

			response, err := ioutil.ReadAll(res.Body)
			if err != nil {
				t.Errorf(err)
			}

			if res.StatusCode != tc.responseStatus {
				t.Errorf("Want '%d', got '%d'", tc.responseStatus, res.StatusCode)
			}

			if strings.TrimSpace(string(response)) != tc.responseBody {
				t.Errorf("Want '%s', got '%s'", tc.responseBody, string(response))
			}
		})
	}
}

We first open a database connection to a test database and pass the connection as an argument to the Router function. In the second highlighted block, we mount the returned router from the Router function on the test HTTP server. Also, we create a client for the server that we will use to send requests to the server.

We then build a new request for the server, and in the next highlighted block, we attach the username and the password of the test case to the request using the Request.SetBasicAuth method.

We send the request to the HTTP server using the client.Do method, and get a response back. We then parse the response and do our assertions on the status code of the response and the contents of the body.

If we ran this test, it would fail. Why? Although the test setup is correct, the test database is empty. When we send a request to the HTTP server, it will try to validate the credentials our test sends, and it will find an empty database:

$ go test ./... -v -count=1
=== RUN   TestServer
    TestServer: server_test.go:17: no such table: users

(/app/server_test.go:16)
[2020-02-16 17:29:26]  no such table: users
--- FAIL: TestServer (0.01s)
FAIL
FAIL	github.com/fteem/go-playground/testing-in-go-cleanup	0.015s
FAIL

This means our test is missing an initial seed of test data, where we create the users table and put some records in it so we can use them in the test.

Adding data and cleaning it up, the old way #

The widely-approved approach to seeding the test database is to insert some records before running the test and then remove them once the test is done. The benefit being that the database will always be empty after the tests are done running, without the need to recreate the database and its tables every time we run the tests. Precisely what our failing tests need.

Achieving this in the most idiomatic way is by using a cleanup closure that is returned by the function that inserts the data. Usually, the returning cleanup function “pattern” looks like this:

func createUser(t *testing.T, db *gorm.DB) func() {
	user := User{Username: "jane", Password: "doe123"}
	if err := db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}

	return func() {
		db.Delete(&user)
	}
}

func TestServer(t *testing.T) {
	db, err := gorm.Open("sqlite3", "./cleanups_test.db")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	cleanup := createUser(t, db)
	defer cleanup()

	r := Router(db)

	tcs := []struct {
		name           string
		username       string
		password       string
		responseBody   string responseStatus int
	}{
		{
			name:           "with invalid username-password combo",
			username:       "jane",
			password:       "doe",
			responseBody:   "Forbidden",
			responseStatus: http.StatusForbidden,
		},
		{
			name:           "with valid username-password combo",
			username:       "jane",
			password:       "doe123",
			responseBody:   "6 x 9 = 42",
			responseStatus: http.StatusOK,
		},
	}

	ts := httptest.NewServer(r)
	defer ts.Close()
	client := ts.Client()

	req, err := http.NewRequest("GET", fmt.Sprintf("%s/foo", ts.URL), nil)
	if err != nil {
		t.Fatal(err)
	}

	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			req.SetBasicAuth(tc.username, tc.password)

			res, err := client.Do(req)
			if err != nil {
				t.Fatal(err)
			}
			defer res.Body.Close()

			response, err := ioutil.ReadAll(res.Body)
			if err != nil {
				t.Error(err)
			}

			if res.StatusCode != tc.responseStatus {
				t.Errorf("Want '%d', got '%d'", tc.responseStatus, res.StatusCode)
			}

			if strings.TrimSpace(string(response)) != tc.responseBody {
				t.Errorf("Want '%s', got '%s'", tc.responseBody, string(response))
			}
		})
	}
}

The createUser function takes a database connection as an argument and uses it to insert a new User in the users table. After that, it will return a closure that, when invoked, will remove the added user.

The reason we pass the testing.T pointer to the createUser function is to be able to fail the test if the database cannot be opened for writing. We check the db.Error field on the Gorm database connection pool for any errors and invoke the t.Fatal function, if we hit any errors.

(You can read more about the behavior of t.Fatal in my article on the topic.)

In the test itself, we invoke createUser with the test database connection and the testing.T pointer. The return value of createUser is a function which we store in the cleanup variable. Then, we defer the invocation of cleanup to happen at the exit of the test function. Using cleanup, we can remove the data from the test database before exiting the tests.

While the approach above works and it’s idiomatic Go, it has a few downsides worth mentioning:

It works fine for one test case, but if the test file has many tests that modify the state of the test database, they will get contaminated with cleanup-like functions.
Adding to the above point, imagine tests where we have to clean up other things, such as processes, files, or open sockets. That will lead to a substantial increase in cleanup-closures pollution.
It adds cognitive overhead when reading the code, which otherwise is a straightforward testing code
The definition of the cleanup function is buried in the createUser function, visually separated from the test function that uses/invokes it. A reader of the test file can experience difficulty navigating the functions and putting the pieces together

Because of these reasons, the Go authors decided to add a t.Cleanup function to the testing package. The change was merged on November 4, 2019 – right on time to get into the v1.14 release.

Let’s explore how t.Cleanup can make our life easier.

Adding data and cleaning it up, using `t.Cleanup` #

With the t.Cleanup, and b.Cleanup methods, we get better control to cleaning up after our tests. t.Cleanup registers a function to be called when the test and all its subtests complete. This means that even if our tests run as subtests, which is the case in our example, the cleanup will only happen after all the subtests are done.

Here’s the reworked version of our tests:

func createUser(t *testing.T, db *gorm.DB) {
	user := User{Username: "jane", Password: "doe123"}
	if err := db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}

	t.Cleanup(func() {
		db.Delete(&user)
	})
}

func TestServer(t *testing.T) {
	db, err := gorm.Open("sqlite3", "./cleanups_test.db")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	createUser(t, db)

	r := Router(db)

    // Snipped for brewity...
}

While the code change is small, there is a substantial difference between the new and the old versions.

The most significant difference lies in the fact that the createUser function takes care of its clean up now. The invoker function (TestServer) does not need to worry about the cleanup of the data added in the database – we know that the createUser will remove the data. Most importantly, it will remove the data just in time – once all of the subtests are done.

Having t.Cleanup, and b.Cleanup in the standard library does not stop us from using the old way using composition and returning cleanup callbacks – it is still a valid way to clean up any state or files that our tests might create. But, using the new t.Cleanup & b.Cleanup functions things get a bit easier: there’s now a way to do just that, and the standard library supports it.

And having the Go v1 compatibility promise in mind – it is here to stay.

Vulnerable Authentication #

Testing the Vulnerable Authentication #

Adding data and cleaning it up, the old way #

Adding data and cleaning it up, using t.Cleanup #

Adding data and cleaning it up, using `t.Cleanup` #