Sunday, October 16, 2016

Testing Golang code that uses exec.Command

At some point you may need to test code that calls exec.Command, but do not want the command to actually run. There are other posts that describe how to do this in Go such as Nate Finch's excellent write-up at https://npf.io/2015/06/testing-exec-command/. Before reading this post, I highly recommend you read Nate's post to get familiar with the mechanisms that allow the code in this post to work. His write-up is a simple example of how to get something to work, however I was looking for something more flexible that could be used in any test without having to reproduce a bunch of boilerplate code and to allow mocked responses for multiple calls to the same command within a test.  So I decided to write a utility to help with those very pieces.

First, to help keep things organized let's define a struct to hold some details about how we want the command that is executed to behave.

type ExecCmdTestResult struct {
 command  string
 exitCode int
 stdOut   string
 stdErr   string
}

The command parameter will hold the actual command that we want to mock a response for. The exitCode is the exit code we want the process to exit with. The stdOut and stdErr strings hold what we want to be mocked to stdout and stderr of the command process.

We will also create the following structure which allows us to hold the mocked responses and to know which test function to call when we want to execute the mocked command.

type ExecCmdTestHelper struct {
 testResults        map[string][]ExecCmdTestResult
 testHelperFuncName string
}

For convenience, we create a "New" function to create the helper.

func NewExecCmdTestHelper(testHelperFuncName string) *ExecCmdTestHelper {
 return &ExecCmdTestHelper{
  testResults:        make(map[string][]ExecCmdTestResult),
  testHelperFuncName: testHelperFuncName,
 }
}

Next, we need a way to add mocked results to the helper, so we define a function that let's us add a mocked results one-by-one.

func (e *ExecCmdTestHelper) AddExecResult(stdOut, stdErr string, exitCode int, command ...string) {
 fullCommand := strings.Join(command, " ")
 base64Command := base64.StdEncoding.EncodeToString([]byte(fullCommand))

 result := ExecCmdTestResult{
  stdOut:   stdOut,
  stdErr:   stdErr,
  exitCode: exitCode,
  command:  fullCommand,
 }

 if e.testResults[base64Command] == nil {
  e.testResults[base64Command] = make([]ExecCmdTestResult, 0)
 }

 e.testResults[base64Command] = append(e.testResults[base64Command], result)
}

The function above creates an ExecCmdTestResult instance and adds it to the list of results.

The helper will also provide a function that will be used as the stand-in for exec.Command. We'll call it ExecCommand which is shown below.

func (m *ExecCmdTestHelper) ExecCommand(command string, args ...string) *exec.Cmd {
 cs := []string{"-test.run=" + m.testHelperFuncName, "--", command}
 cs = append(cs, args...)
 cmd := exec.Command(os.Args[0], cs...)

 fullCommand := command

 if len(args) > 0 {
  fullCommand = command + " " + strings.Join(args, " ")
 }

 base64Command := base64.StdEncoding.EncodeToString([]byte(fullCommand))

 if len(m.testResults[base64Command]) == 0 {
  fmt.Println("No result was setup for command: ", fullCommand)
  return nil
 }

 // Retrieve next result
 mockResults := m.testResults[base64Command][0]

 // Remove current result so that next time it will use next result that was setup.  If no next result, re-use same result.
 if len(m.testResults[base64Command]) > 1 {
  m.testResults[base64Command] = m.testResults[base64Command][1:]
 }

 stdout := execTestStdOutputKey + "=" + mockResults.stdOut
 stderr := execTestStdErrorKey + "=" + mockResults.stdErr
 exitCode := execTestExitCodeKey + "=" + strconv.FormatInt(int64(mockResults.exitCode), 10)

 cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1", stdout, stderr, exitCode}

 return cmd
}

The above function will take the responses you set on the helper and cycle through them for each call to the specified command. The responses are stored in the command's environment so that they can be retrieved when the function that mocks the responses can use them to provide the desired response. If the command is called more times than the number of results that were mocked for that command, the last mocked result for all subsequent calls to that same command will be used. For example, if there was only one result mocked with an exit code of "0" and we call the same command twice, the second time the command is called, it will also receive an exit code of "0". If no mocked result was found for the given command, then it return nil.

Within your code you will need to replace the calls to exec.Command with a variable that contains a function that has the function signature like ExecCommand above. This way, in production you can set the variable to exec.Command, but for tests it can be set to the helper's ExecCommand function. In your tests, instead of exec.Command, the above function will be called which in turn will execute the test with the name given (testHelperFuncName). Within your test file you must add a function with that name which calls another new function named RunTestExecCmd which will mock stdout, stderr, and the exit code.

Add this to your test file:
func TestHelperProcess(t *testing.T) {
 RunTestExecCmd()
}

Function that mocks process response:

func RunTestExecCmd() {
 if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
  return
 }

 stdout := os.Getenv(execTestStdOutputKey)
 stderr := os.Getenv(execTestStdErrorKey)
 exitCode, err := strconv.ParseInt(os.Getenv(execTestExitCodeKey), 10, 64)

 if err != nil {
  os.Exit(1)
 }

 fmt.Fprintf(os.Stdout, stdout)
 fmt.Fprintf(os.Stderr, stderr)

 os.Exit(int(exitCode))
}

The function above retrieves the values that were set by the ExecCommand function from the command's environment and uses them to create the proper stdout and stderr strings and exit with the proper code.

All of the code above, along with an example of how to use it can be found at the following gist: https://gist.github.com/kglee79/db8f0bf3eafe962e0feddac8451387da

Hopefully you can now more conveniently test code that calls exec.Command!