Every Go program’s entry point is the main function. Typically, the main function initializes the program with configuration values from different sources, sets up dependencies, and hands these values off to deeper components of the application.
As the main function is declared as func main(), it has neither input parameters nor return values. A function with such a declaration is hard to unit test. In this article, I’m going to share a method to unit test the main function. This method is inspired by the source code of HashiCorp’s Terraform.
Sample CLI main function
We start with a simple Go CLI application:
main.go
package main
import (
"flag"
"fmt"
"os"
)
func main() {
name := flag.String("name", "", "Your Name")
if *name == "" {
fmt.Printf("Missing flag -name\n")
os.Exit(1)
}
fmt.Printf("Hi %v", *name)
}
The application looks for the flag -name. If it is set it prints “Hi <name>” to the console otherwise it prints an error message to the console and exits the program with the exit code 1.
At first glance, this is a very simple program. But its main function depends on multiple other packages.
- The flag package to parse the CLI flags
- The os package to exit the program with a non-zero exit code
- The fmt package to print to console
Let’s see how we can get around those dependencies and unit test our application code.
Test CLI flags
Our first target is to make sure the application parses the flags provided correctly. The flag package uses the variable os.Args, exposed by the os package, to parse the flags given to the application. We can use this to set the flags in our test as follows:
main_test.go
package main
import (
"flag"
"os"
"testing"
)
func TestFlags(T *testing.T) {
// We manipuate the Args to set them up for the testcases
// after this test we restore the initial args
oldArgs := os.Args
defer func() { os.Args = oldArgs }()
cases := []struct {
Name string
Args []string
ExpectedExit int
ExpectedOutput string
}{
{"flags set", []string{"-name", "johannes"}, 0, "Hi johannes\n"},
{"flags not set", []string{"test"}, 1, "Missing flag -name\n"},
}
for _, tc := range cases {
// this call is required because otherwise flags panics, if args are set between flag.Parse calls
flag.CommandLine = flag.NewFlagSet(tc.Name, flag.ExitOnError)
// we need a value to set Args[0] to, cause flag begins parsing at Args[1]
os.Args = append([]string{tc.Name}, tc.Args...)
main()
}
}
Notice there is no assertion in this test we’re calling the main function but it doesn’t return a value. Additionally, if we run this test it exits after the first test case due to the os.Exit call in the main function. So how can we test, that the function exits with the expected code and make sure every test case is actually executed?
Test OS exit codes
The solution is to not test the actual os.Exit call, but to make the main function simple enough, so we don’t have the urge to test it. To do so, we define a function called realMain that returns an int. This int is the exit code. Then we call the realMain function inside the os.Exit call.
main.go
package main
import (
"flag"
"fmt"
"io"
"os"
)
func main() {
// I'm ok with not testing this call
os.Exit(realMain())
}
func realMain() int {
name := flag.String("name", "", "Your Name")
flag.Parse()
if *name == "" {
fmt.Printf("Missing flag -name\n")
return 1
}
fmt.Printf("Hi %v\n", *name)
return 0
}
Now let’s fix the test and add an actual assertion:
main_test.go
package main
import (
"flag"
"os"
"testing"
)
func TestFlags(T *testing.T) {
// We manipuate the Args to set them up for the testcases
// After this test we restore the initial args
oldArgs := os.Args
defer func() { os.Args = oldArgs }()
cases := []struct {
Name string
Args []string
ExpectedExit int
ExpectedOutput string
}{
{"flags set", []string{"-name", "johannes"}, 0, "Hi johannes\n"},
{"flags not set", []string{"test"}, 1, "Missing flag -name\n"},
}
for _, tc := range cases {
// this call is required because otherwise flags panics,
// if args are set between flag.Parse call
flag.CommandLine = flag.NewFlagSet(tc.Name, flag.ExitOnError)
// we need a value to set Args[0] to cause flag begins parsing at Args[1]
os.Args = append([]string{tc.Name}, tc.Args...)
actualExit := realMain()
if tc.ExpectedExit != actualExit {
T.Errorf("Wrong exit code for args: %v, expected: %v, got: %v",
tc.Args, tc.ExpectedExit, actualExit)
}
}
}
Great, the test is running and it passes all test cases in the table. The only thing that’s left is to test that the CLI is actually printing the expected outputs to the console. Let’s see how to test the fmt calls in the realMain function.
Test the fmt package
In the application, we use the function fmt.Printf to print the response of the CLI to stdout. To test that we can refactor the app to use fmt.Fprintf instead. This function is called with an additional io.Writer argument, which it writes to. We now have to pass os.Stdout as the io.Writer argument in the application code. As our main function logic is already in the realMain function this is straightforward:
main.go
package main
import (
"flag"
"fmt"
"io"
"os"
)
func main() {
// I'm ok with not testing this call
os.Exit(realMain(os.Stdout))
}
func realMain(out io.Writer) int {
name := flag.String("name", "", "Your Name")
flag.Parse()
if *name == "" {
fmt.Fprintf(out, "Missing flag -name\n")
return 1
}
fmt.Fprintf(out, "Hi %v\n", *name)
return 0
}
So the application is still working as expected, now we fix the test by calling realMain with a bytes.Buffer pointer. The fmt calls will write the output to the bytes.Buffer and we can use this to test the output against the ExpectedOutput field of our test cases.
main_test.go
package main
import (
"bytes"
"flag"
"os"
"testing"
)
func TestFlags(T *testing.T) {
// We manipuate the Args to set them up for the testcases
// After this test we restore the initial args
oldArgs := os.Args
defer func() { os.Args = oldArgs }()
cases := []struct {
Name string
Args []string
ExpectedExit int
ExpectedOutput string
}{
{"flags set", []string{"-name", "johannes"}, 0, "Hi johannes\n"},
{"flags not set", []string{"test"}, 1, "Missing flag -name\n"},
}
for _, tc := range cases {
// this call is required because otherwise flags panics,
// if args are set between flag.Parse call
flag.CommandLine = flag.NewFlagSet(tc.Name, flag.ExitOnError)
// we need a value to set Args[0] to cause flag begins parsing at Args[1]
os.Args = append([]string{tc.Name}, tc.Args...)
var buf bytes.Buffer
actualExit := realMain(&buf)
if tc.ExpectedExit != actualExit {
T.Errorf("Wrong exit code for args: %v, expected: %v, got: %v",
tc.Args, tc.ExpectedExit, actualExit)
}
actualOutput := buf.String()
if tc.ExpectedOutput != actualOutput {
T.Errorf("Wrong output for args: %v, expected %v, got: %v",
tc.Args, tc.ExpectedOutput, actualOutput)
}
}
}
Conclusion
It is hard to test the main function in Go, cause it doesn’t take any parameters nor returns any values. But if you have the urge to test your main logic anyways, it is a great approach to make the main function as trivial as possible and just call another function. With this, you can test a function that is declared by yourself and takes dependencies as interfaces, to enable mocking, or returns the values you want to test against.
This article was inspired by the tests of Terraform. Reading other people’s code is one of the best ways to improve your own coding skills. And with the amazing collection of successful open-source projects written in Go, you have a good opportunity to learn from some of the best Go developers in the world.
Great article. I am pondering using Bats (https://github.com/bats-core/bats-core), but I do not like the idea of splitting the test suite, but this is a great start to get basic tests going.