TestMain—What is it Good For?

Thu, Jan 8, 2015

As stated in the release notes, Go 1.4 focuses primarily on implementation work, but it also provides a few new tools for developers. The testing package provides one of these tools. Test code may now contain a TestMain function that provides more control over running tests than was available in prior releases of Go.

Justinas Stankevičius and Cory Jacobsen have written excellent posts about TestMain recently. Both posts demonstrate how a package can perform global setup and shutdown steps when running tests, but there are other uses for TestMain.

First a quick review of the mechanics. When we run go test the go tool generates a main program that runs the test functions for our package. If the go tool finds a TestMain function it instead generates code to call TestMain. A typical TestMain follows.

func TestMain(m *testing.M) {
    setup()
    code := m.Run() 
    shutdown()
    os.Exit(code)
}

From the above code we can see that writing a TestMain function allows us to control four aspects of test execution.

Setup.
How and when to run the tests.
Shutdown.
Exit behavior.

With appropriate use of these four items, we can satisfy several different use cases that were not well served prior to Go 1.4.

Use Cases

The following sections describe four uses for TestMain that I think are worth understanding. The first two cases motivated the Go team to add support for TestMain. The third case is drawn from my experience with a project at my day job. The last use case is the sole example of a TestMain in the Go 1.4 standard library.

Global Startup and Shutdown Hooks

As already mentioned, adding a TestMain to a package allows it to run arbitrary code before and after tests run. But only the second part is really new. Pre-test setup has always been possible by defining an init() function in a test file. The challenge—as described in issue #8159—has been bookending that with corresponding shutdown code when all tests have completed.

Testing from the Main OS Thread

Many graphics libraries are particular about which OS thread calls their API. Gustavo Niemeyer encountered this while working on his qml package for Go. In issue #8202 Russ Cox demonstrated how TestMain can address this need:

func init() {
    runtime.LockOSThread()
}

func TestMain(m *testing.Main) {
    go func() {
        os.Exit(m.Run())
    }()
    runGraphics()
}

The above code ensures that runGraphics() runs in the main goroutine locked to the main OS thread while the package tests run elsewhere. Presumably the tests exercise APIs that communicate with the runGraphics() goroutine.

Subprocess Tests

Sometimes we need to test the behavior of a process, rather than a function. For example, one of my projects follows the crash-only software design philosophy. We have tests that deliberately crash and restart pieces of the system in different combinations and permutations between key events.

Initially we wrote a separate package main that produced an executable the crash tests launched. This approach worked, but it was fragile. When we forgot to rebuild the test executable before running go test the behavior we just tried to fix would still appear broken.

We recently adopted a better approach demonstrated by Andrew Gerrand in the presentation on testing techniques he gave at Google I/O 2014. On slide 23, he showed how to reuse the test binary produced by go test as the child process. This technique eliminates the need for a separate build and ensures the testing and tested code are always in sync.

TestMain was not available at the time Andrew gave his presentation. In Go 1.3 our only choice was for the child process to call the tested code from within a TestXxx() function. This constraint came with some baggage.

The command line arguments passed to the child process must be compatible with the testing package.
The test framework will, at a minimum, write either ‘PASS’ or ‘FAIL’ to stdout after the test function returns. The child process must make sure to call os.Exit() rather than return from the TestXxx() function to avoid the extra output.

We end up writing pseudo-tests such as this example from the os/exec package:

// TestHelperProcess isn't a real test. It's used
// as a helper process for TestParameterRun.
func TestHelperProcess(*testing.T) {
    if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
        return
    }
    defer os.Exit(0)
    ...
}

These issues are admittedly not major. They do not—by themselves—make a convincing argument for adding the TestMain feature. But TestMain gives us a potentially less finicky solution. Adapting Andrew’s example:

func Crasher() {
    fmt.Println("Going down in flames!")
    os.Exit(1)
}

func TestMain(m *testing.M) {
    switch os.Getenv("TEST_MAIN") {
    case "crasher":
        Crasher()
    default:
        os.Exit(m.Run())
    }
}

func TestCrasher(t *testing.T) {
    cmd := exec.Command(os.Args[0])
    cmd.Env = append(os.Environ(), "TEST_MAIN=crasher")
    err := cmd.Run()
    if e, ok := err.(*exec.ExitError); ok && !e.Success() {
        return
    }
    t.Fatalf("process err %v, want exit status 1", err)
}

Global Resource Checks

Go programmers often look to the standard library for examples of good, idiomatic Go code. I found one use of TestMain in the Go 1.4 standard library. The net/http package contains the following interesting use of TestMain.

func TestMain(m *testing.M) {
    v := m.Run()
    if v == 0 && goroutineLeaked() {
        os.Exit(1)
    }
    os.Exit(v)
}

// Verify the other tests didn't leave
// any goroutines running.
func goroutineLeaked() bool {
    ...
}

As you can see, if all the tests succeed TestMain will still exit with a failure code if it finds any leaking goroutines. Ideally each test would check for leaked resources, but occasionally the approach taken above is the best way.

Conclusion

The TestMain feature added to Go’s testing framework in the latest release is a simple solution for several testing use cases. TestMain provides a global hook to perform setup and shutdown, control the testing environment, run different code in a child process, or check for resources leaked by test code. Most packages will not need a TestMain, but it is a welcome addition for those times when it is needed.