Why use TestMain for testing in Go?

The TestMain function was released with Go 1.4 about 2 years ago and I feel like it is still not used enough.

What is TestMain function exactly?

No, it is not to test the main function on your program. Basically the TestMain function provides more control over running tests than was available in the prior releases of Go. So if the test code contains a function:

func TestMain(m *testing.M)

that function will be called instead of running the tests directly. The M struct contains methods to access and run the tests. But there are some rules to remember when using TestMain function:

1. You can define this function in your test file: func TestMain(m *testing.M).
2. Since function names need to be unique in a given package, you can only define the TestMain once for each package. If your package has multiple test files, choose a logic place for your single TestMain function.
3. The testing.M struct has a single defined function named Run(). As you can imagine, it runs all the tests within the package. Run() returns an exit code that can be passed to os.Exit.
4. A minimal implementation looks like this: func TestMain(m *testing.M) { os.Exit(m.Run()) }.
5. If you don't call os.Exit with the return code, your test command will return 0(zero). Yes, even if a test fails.

Here is an example:

package mypackagename

import (
    "testing"
    "os"
)

func TestMain(m *testing.M) {
    log.Println("Do stuff BEFORE the tests!")
    exitVal := m.Run()
    log.Println("Do stuff AFTER the tests!")

    os.Exit(exitVal)
}

func TestA(t *testing.T) {
    log.Println("TestA running")
}

func TestB(t *testing.T) {
    log.Println("TestB running")
}

This would be the output:

2017/12/29 00:32:17 Do stuff BEFORE the tests!
2017/12/29 00:32:17 TestA running
2017/12/29 00:32:17 TestB running
PASS
2017/12/29 00:32:17 Do stuff AFTER the tests!

So this may come in handy when you need to do some global set-up/tear-down for your tests, but remember, TestMain will only run once, so if you need to clear a database or a tmp folder every each test you need to do it yourself.

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. Running global test setup has always been possible by defining an init() function in a test file. The challenge has been aligning that with corresponding shutdown code when all tests have completed.

How about we do a more complex example then that?

Suppose we need to do some database testing where we had to do some setup before each test and after everything drop what I did and restore the original one. Without using TestMain or any other framework we had to do something like this:

func TestDBFeatureA(t *testing.T) {
    models.TestDBManager.Setup()
    defer models.TestDBManager.Exit()

    // Do the tests
}
func TestDBFeatureB(t *testing.T) {
    models.TestDBManager.Setup()
    defer models.TestDBManager.Exit()

    // Do the tests
}

As you can see in every single test which utilized the DB had to go through the same process of setup and exiting and as the testing suite grow, the time it took to run tests grew linearly.

Now it comesTestMain() to the rescue making it possible to run these migrations only once. So now the code would look more like this:

func TestDBFeatureA(t *testing.T) {
    defer models.TestDBManager.Reset()

    // Do the tests
}
func TestDBFeatureB(t *testing.T) {
    defer models.TestDBManager.Reset()

    // Do the tests
}
func TestMain(m *testing.M) {
    models.TestDBManager.Setup()
    // os.Exit() does not respect defer statements
    code := m.Run()
    models.TestDBManager.Exit()
    os.Exit(code)
}

While each test must still clean up after itself, that would only involve restoring the initial data, which is way faster than doing the schema migrations.

This approach also reduces code duplication since we now only have one line for database management in each test instead of two and would run a lot faster without a doubt.

Simple is usually better!

The TestMain function is a simple solution that it brings a lot to the table. Off course that not all packages will need to implement the TestMain but it is a nice addition for those that need. Specially because I feel like a lot of people use third party testing frameworks just because of their setup and tear down features, specially those used to deal with a lot of third party frameworks in other languages. I just think that the Go Standard Library has a more to offer then usually people realize and sometimes they bring third party framework without the actual need to. Not only related to testing.

Did you like the post? :)

How about clicking on the little "heart" on the left or below the post so more people can read it too?

Cheers!