Go Patterns: Return Cleanup Functions

#golang

This post attempts to document and justify the rationale behind a simple pattern we have been using in many of our tests: returning cleanup functions.

Problem

If you have written a non-trivial amount of Python code before, you are likely familiar with Context Managers. If that is not you, the quick version is that they allow you to write something similar to the following.

with db.connect() as conn:
  # ...

The idiomatic way of doing the same in Go is with the use of defer statements.

conn, err := db.Connect()
if err != nil {
  // ...
}
defer conn.Close()

This is enough for many use cases except when you are unable to add the cleanup method (Close) to the retuned object (conn).

An example of this that comes up often while writing tests is the need to create a temporary directory and cleaning up afterwards. Let us look at the documentation for "io/ioutil".TempDir.

TempDir creates a new temporary directory [...] and returns the path of the new directory. [...] It is the caller's responsibility to remove the directory when no longer needed.

Since TempDir returns a string, adding a Close or Delete method to it is not an option. So one might be inclined to fall back to a Python-like pattern.

// WithTempDir creates a temporary directory and calls the given
// function with the path to it. The directory is cleaned up when
// the function returns.
//
// The test is marked as failed if the temporary directory could
// not be created.
func WithTempDir(t *testing.T, f func(dir string)) {
  dir, err := ioutil.TempDir("", "test")
  if err != nil {
    t.Fatalf("failed to create temporary directory: %v", err)
  }
  defer os.RemoveAll(dir)
  f(dir)
})

The above function would be used in a test like so,

func TestFoo(t *testing.T) {
  WithTempDir(t, func(dir string) {
    // ...
  })
}

Why With* is Bad

If you already see why this is not the best way to do this, skip on over to the Solution section. Otherwise, read on.

This is obviously subjective but here are a handful of obvious issues with it.

Solution

Ideally, we want to be able to defer cleanup functions but adding a Close() or Delete() method to string is not possible. How about returning an anonymous function that does the cleanup?

// CreateTempDir creates a temporary directory and returns the path
// to it, along with a function to clean up when the directory is
// no longer needed.
//
// The test is marked as failed if the temporary directory could
// not be created.
func CreateTempDir(t *testing.T) (dir string, cleanup func()) {
  dir, err := ioutil.TempDir("", "test")
  if err != nil {
    t.Fatalf("failed to create temporary directory: %v", err)
  }
  return dir, func() {
    os.RemoveAll(dir)
  }
}

That looks like it would work!

Here is how you would use it:

func TestFoo(t *testing.T) {
  dir, cleanup := CreateTempDir(t)
  defer cleanup()
  
  // ...
}

We can call any number of similar functions without affecting readability of the core logic.

dir, cleanup := CreateTempDir(t)
defer cleanup()

file, cleanup := CreateTempFile(t, dir)
defer cleanup()

// ...

Plus it is easy to call it dynamically and defer cleanup in a loop.

dirs := make([]string, N)
for i := 0; i < N; i++ {
    dir, cleanup := CreateTempDir(t)
    defer cleanup()
    dirs[i] = dir
}

(Note that deferred calls are executed when the surrounding function returns.)

This is a fairly simple pattern and as demonstrated, it is more readable, flexible, and idiomatic than Context Manager-style With* functions. It has found frequent use in our codebase for testing setup and teardown, and in certain cases, outside tests as well.