Running "go test" command on a remote host using SSH

README.md

To use this snippet:

1. Add the following to the package you would like to be able to test remotely:

   func TestMain(m *testing.M) {
     testutil.MainWithRemoteExecutionSupport(m)
   }

2. Run your tests as GO_TEST_REMOTE_EXECUTION_SSH_HOST=<SSH host> go test ./<path to your package>.

   In case your hosts use different glibc versions and you get an error executing a binary, add CGO_ENABLED=0 to the command above.

test.go

package testutil

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"testing"
)

// GoTestRemoteExecutionSSHHost defines environment variable name used for specifying remote SSH host for
// executing Go tests.
const GoTestRemoteExecutionSSHHost = "GO_TEST_REMOTE_EXECUTION_SSH_HOST"

// MainWithRemoteExecutionSupport is indented to be called from TestMain(m *testing.M) function in a package, which
// contains integration tests which expects to be executed on a remote host specified by
// GO_TEST_REMOTE_EXECUTION_SSH_HOST environment variable.
//
// When GO_TEST_REMOTE_EXECUTION_SSH_HOST environment variable is set, instead of executing tests locally, test binary
// will copy itself to a specified remote host over SSH and execute there, which enables easy integration tests for
// code which e.g. requires root or modifies system files.
func MainWithRemoteExecutionSupport(m *testing.M) {
	host := os.Getenv(GoTestRemoteExecutionSSHHost)
	if host == "" {
		os.Exit(m.Run())
	}

	cmdWithRedirectedOutput := func(command string, args ...string) (*exec.Cmd, io.Closer) {
		cmd := exec.Command(command, args...)
		// Redirect stdin to enable responding to possible SSH prompts.
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout

		stderrPipeReader, stderrPipeWriter := io.Pipe()
		cmd.Stderr = stderrPipeWriter

		scanner := bufio.NewScanner(stderrPipeReader)

		go func() {
			for scanner.Scan() {
				text := scanner.Text()
				// Filter this annoying SSH message which is quite difficult to get rid of via configuration.
				if strings.HasPrefix(text, "Warning: Permanently added") {
					continue
				}

				fmt.Fprintf(os.Stderr, text+"\n")
			}
		}()

		return cmd, stderrPipeWriter
	}

	runOnTestHost := func(command string, args ...string) error {
		cmd, closer := cmdWithRedirectedOutput("ssh", append([]string{host, command}, args...)...)

		// Make sure to close stderr pipe so scanner can terminate.
		defer func() {
			if err := closer.Close(); err != nil {
				fmt.Fprintf(os.Stderr, "Failed closing process stderr pipe: %v", err)
			}
		}()

		return cmd.Run()
	}

	runSelfRemotelyOnTestHost := func() error {
		srcPath := os.Args[0]
		dstPath := "/tmp/" + base64.StdEncoding.EncodeToString([]byte(os.Args[0]))
		command := "sudo"

		// Append original args to be able to use -count, -v or -run remotely.
		// Though skip first argument as it's a test binary name and second as it's -test.testlogfile which will not exist
		// on the remote host.
		args := append([]string{dstPath}, os.Args[2:]...)

		commandRaw := strings.Join(append([]string{command}, args...), " ")
		log.Printf("Running %q remotely at %q as %q\n", srcPath, host, commandRaw)

		cmd, closer := cmdWithRedirectedOutput("scp", srcPath, fmt.Sprintf("%s:%s", host, dstPath))

		// Make sure to close stderr pipe so scanner can terminate.
		defer func() {
			if err := closer.Close(); err != nil {
				fmt.Fprintf(os.Stderr, "Failed closing process stderr pipe: %v", err)
			}
		}()

		if err := cmd.Run(); err != nil {
			return fmt.Errorf("copying test binary: %w", err)
		}

		// As we upload test binaries with names based on Go build ID which changes on every run, by default remove
		// the binary after a run to avoid filling up test host.
		defer func() {
			if err := runOnTestHost("rm", dstPath); err != nil {
				log.Printf("Failed removing test binary: %v", err)
			}
		}()

		if err := runOnTestHost("chmod", "+x", dstPath); err != nil {
			return fmt.Errorf("making test binary executable: %w", err)
		}

		if err := runOnTestHost(command, args...); err != nil {
			return fmt.Errorf("running test binary: %w", err)
		}

		return nil
	}

	if err := runSelfRemotelyOnTestHost(); err != nil {
		fmt.Fprintf(os.Stderr, "Failed running test binary remotely on test host: %v", err)
		os.Exit(1)
	}
}