This post is an introduction to a rather underrated feature introduced to Go in version 1.14 - test cleanup. Or maybe I'm just getting too excited every time "testing" is mentioned in any release notes? Well, I do like testing, and I do appreciate what this particular change brings to my (our?) life.
Background
For the sake of this blog post, we assume that we have to test a small "database" that has one feature - the ability to calculate a square of any provided integer value. The calculation is a bit time consuming, which we explicitly define in the code:
func (d *Database) Square(n int) int {
if d.connection == nil {
return -2
}
time.Sleep(100 * time.Millisecond)
if n > 100 {
return -1
}
return n * n
}
As you can see from the snippet, we also want the database to mock a state of being open (having a connection ready for calculations) and closed (when we're done with it). While closing operation is rather quick, starting the connection really takes some time:
func (d *Database) Open(cp ConnPool) {
time.Sleep(2 * time.Second)
d.connection = cp.GetConn()
d.values = make(map[int]string)
}
func (d *Database) Close(cp ConnPool) error {
time.Sleep(100 * time.Millisecond)
cp.FreeConn(d.connection)
d.connection = nil
return nil
}
ConnPool represents a wrapper on a buffered channel that mocks the mechanism for limiting the number of concurrent connections.
Last, but not least, since we are well-organized developers, we absolutely cannot leave the "connection" to our database open after running all of our tests.
It seems simple enough, but the problem is, our code has about 100 edge cases that we need to check, and we want to have our feedback loop as short as possible...
Sequential run
Let's start with the simplest approach. We open the connection, run one test case after the other, and finally close the database. Since we are clever, we use defer not to forget about that last part. Our test looks like this:
func TestDatabase_Sequential(t *testing.T) {
pool := initializeConnectionPool(10)
db := &Database{}
db.Open(pool)
defer db.Close(pool)
for i := 0; i < 100; i++ {
func(i int) {
t.Run(fmt.Sprintf("test for %d", i), func(t *testing.T) {
res := db.Square(i)
if res != i*i {
t.Errorf("Expected %v, got %v", i*i, res)
}
})
}(i)
}
}
// time go test . -count 1 -run TestDatabase_Sequential
// ok github.com/mycodesmells/golang-examples/testing/cleanup 12.423s
// go test . -count 1 -run TestDatabase_Sequential 0.47s user 0.40s system 6% cpu 12.852 total
We are happy to say that all of our cases passed, and it took just under 13 seconds. We are not happy with the result, because we know that each calculation is independent of the other, so we may as well parallelize it, right?
Parallel run
Fortunately, we know that *testing.T has a function called Parallel() that informs our test runner that this particular run can be executed at the same time as the others. Let's try that then:
func TestDatabase_Parallel(t *testing.T) {
pool := initializeConnectionPool(10)
db := &Database{}
db.Open(pool)
defer db.Close(pool)
for i := 0; i < 100; i++ {
t.Run(fmt.Sprintf("test for %d", i), func(t *testing.T) {
t.Parallel()
res := db.Square(i)
if res != i*i {
t.Errorf("Expected %v, got %v", i*i, res)
}
})
}
}
// time go test . -count 1 -run TestDatabase_Parallel
// ...
// FAIL github.com/mycodesmells/golang-examples/testing/cleanup 2.546s
// go test . -count 1 -run TestDatabase_Parallel 0.45s user 0.65s system 35% cpu 3.156 total
We managed to shorten the test run to just over three seconds, but all of the tests failed! Why is that? It's because when we ran each test in a separate goroutine, then the main one got straight to its end and called deferred db.Close(). By the time each spawned goroutine gets to its calculations, the connection is closed and an appropriate "error" value is returned.
Duplicate setup code
Let's not complicate things too much, maybe we just create a separate database for each test case? This can't be a bad idea, right?
func TestDatabase_DuplicateSetup(t *testing.T) {
pool := initializeConnectionPool(10)
for i := 0; i < 100; i++ {
func(i int) {
t.Run(fmt.Sprintf("test for %d", i), func(t *testing.T) {
t.Parallel()
db := &Database{}
db.Open(pool)
defer db.Close(pool)
res := db.Square(i)
if res != i*i {
t.Errorf("Expected %v, got %v", i*i, res)
}
})
}(i)
}
}
// time go test . -count 1 -run TestDatabase_DuplicateSetup
// ok github.com/mycodesmells/golang-examples/testing/cleanup 28.998s
// go test . -count 1 -run TestDatabase_DuplicateSetup 0.49s user 0.74s system 4% cpu 29.748 total
It turns out that this approach is even worse than our initial, sequential run! But we should expect this to be this way since we are opening a connection so many times, have a limitation on the concurrent connections count, etc...
Cleanup instead of defer
Is there another way to go here? Yes, as long as you are using Go 1.14 already! There is a new function available on *testing.T object (and on a few others in this package) that can really help here:
func TestDatabase_Cleanup(t *testing.T) {
pool := initializeConnectionPool(10)
db := &Database{}
db.Open(pool)
t.Cleanup(func() {
db.Close(pool)
})
for i := 0; i < 100; i++ {
func(i int) {
t.Run(fmt.Sprintf("test for %d", i), func(t *testing.T) {
t.Parallel()
res := db.Square(i)
if res != i*i {
t.Errorf("Expected %v, got %v", i*i, res)
}
})
}(i)
}
}
// time go test . -count 1 -run TestDatabase_Cleanup
// ok github.com/mycodesmells/golang-examples/testing/cleanup 3.743s
// go test . -count 1 -run TestDatabase_Cleanup 0.50s user 0.78s system 27% cpu 4.634 total
Ladies and gentlemen, we have a winner! Less than five seconds! What is the trick here? What t.Cleanup(..) does is actually waits for all the spawned goroutines to finish, then calls whatever is defined inside! This is a massive improvement, since you really don't want to handle parallelism yourself in the tests (readability takes a huge hit if you do that manually), but you still want to get the (correct) test results as soon as possible.
Summary
While this may seem like an imaginary problem, there are already a few real-life scenarios where I can see myself using t.Cleanup. I don't want to say that you will not be able to survive without it, but I'm sure it may make your life just a bit easier in some more difficult testing situations.