cmd/testwrapper: parse args better

Previously we were just smushing together args and not trying
to parse the values at all. This resulted in the args to testwrapper
being limited and confusing.

This makes it so that testwrapper parses flags in the exact format as `go test`
command and passes them down in the provided order. It uses tesing.Init to
register flags that `go test` understands, however those are not the only
flags understood by `go test` (such as `-exec`) so we register these separately.

Updates tailscale/corp#14975

Signed-off-by: Maisem Ali <maisem@tailscale.com>

.github/workflows/test.yml

@@ -47,7 +47,7 @@ jobs:
     - name: build test wrapper
       run: ./tool/go build -o /tmp/testwrapper ./cmd/testwrapper
     - name: integration tests as root
-      run: PATH=$PWD/tool:$PATH /tmp/testwrapper --sudo ./tstest/integration/ -race
+      run: PATH=$PWD/tool:$PATH /tmp/testwrapper -exec "sudo -E" -race ./tstest/integration/
 
   test:
     strategy:

cmd/testwrapper/args.go

@@ -0,0 +1,130 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package main
+
+import (
+	"flag"
+	"io"
+	"os"
+	"slices"
+	"strings"
+	"testing"
+)
+
+// defaultTestArgs contains the default values for all flags in the testing
+// package. It is used to reset the flag values in testwrapper tests to allow
+// parsing the flags again.
+var defaultTestArgs map[string]string
+
+// initDefaultTestArgs initializes defaultTestArgs.
+func initDefaultTestArgs() {
+	if defaultTestArgs != nil {
+		return
+	}
+	defaultTestArgs = make(map[string]string)
+	flag.CommandLine.VisitAll(func(f *flag.Flag) {
+		defaultTestArgs[f.Name] = f.DefValue
+	})
+}
+
+// registerTestFlags registers all flags from the testing package with the
+// provided flag set. It does so by calling testing.Init() and then iterating
+// over all flags registered on flag.CommandLine.
+func registerTestFlags(fs *flag.FlagSet) {
+	testing.Init()
+	type bv interface {
+		IsBoolFlag() bool
+	}
+
+	flag.CommandLine.VisitAll(func(f *flag.Flag) {
+		if b, ok := f.Value.(bv); ok && b.IsBoolFlag() {
+			fs.Bool(f.Name, f.DefValue == "true", f.Usage)
+			if name, ok := strings.CutPrefix(f.Name, "test."); ok {
+				fs.Bool(name, f.DefValue == "true", f.Usage)
+			}
+			return
+		}
+
+		// We don't actually care about the value of the flag, so we just
+		// register it as a string. The values will be passed to `go test` which
+		// will parse and validate them anyway.
+		fs.String(f.Name, f.DefValue, f.Usage)
+		if name, ok := strings.CutPrefix(f.Name, "test."); ok {
+			fs.String(name, f.DefValue, f.Usage)
+		}
+	})
+}
+
+// splitArgs splits args into three parts as consumed by go test.
+//
+//	go test [build/test flags] [packages] [build/test flags & test binary flags]
+//
+// We return these as three slices of strings [pre] [pkgs] [post].
+//
+// It is used to split the arguments passed to testwrapper into the arguments
+// passed to go test and the arguments passed to the tests.
+func splitArgs(args []string) (pre, pkgs, post []string, _ error) {
+	if len(args) == 0 {
+		return nil, nil, nil, nil
+	}
+
+	fs := newTestFlagSet()
+	// Parse stops at the first non-flag argument, so this allows us
+	// to parse those as values and then reconstruct them as args.
+	if err := fs.Parse(args); err != nil {
+		return nil, nil, nil, err
+	}
+	fs.Visit(func(f *flag.Flag) {
+		if f.Value.String() != f.DefValue && f.DefValue != "false" {
+			pre = append(pre, "-"+f.Name, f.Value.String())
+		} else {
+			pre = append(pre, "-"+f.Name)
+		}
+	})
+
+	// fs.Args() now contains [packages]+[build/test flags & test binary flags],
+	// to split it we need to find the first non-flag argument.
+	rem := fs.Args()
+	ix := slices.IndexFunc(rem, func(s string) bool { return strings.HasPrefix(s, "-") })
+	if ix == -1 {
+		return pre, rem, nil, nil
+	}
+	pkgs = rem[:ix]
+	post = rem[ix:]
+	return pre, pkgs, post, nil
+}
+
+func newTestFlagSet() *flag.FlagSet {
+	fs := flag.NewFlagSet("testwrapper", flag.ContinueOnError)
+	fs.SetOutput(io.Discard)
+
+	// Register all flags from the testing package.
+	registerTestFlags(fs)
+	// Also register the -exec flag, which is not part of the testing package.
+	// TODO(maisem): figure out what other flags we need to register explicitly.
+	fs.String("exec", "", "Command to run tests with")
+	fs.Bool("race", false, "build with race detector")
+	return fs
+}
+
+// testingVerbose reports whether the test is being run with verbose logging.
+var testingVerbose = func() bool {
+	verbose := false
+
+	// Likely doesn't matter, but to be correct follow the go flag parsing logic
+	// of overriding previous values.
+	for _, arg := range os.Args[1:] {
+		switch arg {
+		case "-test.v", "--test.v",
+			"-test.v=true", "--test.v=true",
+			"-v", "--v",
+			"-v=true", "--v=true":
+			verbose = true
+		case "-test.v=false", "--test.v=false",
+			"-v=false", "--v=false":
+			verbose = false
+		}
+	}
+	return verbose
+}()

cmd/testwrapper/args_test.go

@@ -0,0 +1,97 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package main
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestSplitArgs(t *testing.T) {
+	tests := []struct {
+		name            string
+		in              []string
+		pre, pkgs, post []string
+	}{
+		{
+			name: "empty",
+		},
+		{
+			name: "all",
+			in:   []string{"-v", "pkg1", "pkg2", "-run", "TestFoo", "-timeout=20s"},
+			pre:  []string{"-v"},
+			pkgs: []string{"pkg1", "pkg2"},
+			post: []string{"-run", "TestFoo", "-timeout=20s"},
+		},
+		{
+			name: "only_pkgs",
+			in:   []string{"./..."},
+			pkgs: []string{"./..."},
+		},
+		{
+			name: "pkgs_and_post",
+			in:   []string{"pkg1", "-run", "TestFoo"},
+			pkgs: []string{"pkg1"},
+			post: []string{"-run", "TestFoo"},
+		},
+		{
+			name: "pkgs_and_post",
+			in:   []string{"-v", "pkg2"},
+			pre:  []string{"-v"},
+			pkgs: []string{"pkg2"},
+		},
+		{
+			name: "only_args",
+			in:   []string{"-v", "-run=TestFoo"},
+			pre:  []string{"-run", "TestFoo", "-v"}, // sorted
+		},
+		{
+			name: "space_in_pre_arg",
+			in:   []string{"-run", "TestFoo", "./cmd/testwrapper"},
+			pre:  []string{"-run", "TestFoo"},
+			pkgs: []string{"./cmd/testwrapper"},
+		},
+		{
+			name: "space_in_arg",
+			in:   []string{"-exec", "sudo -E", "./cmd/testwrapper"},
+			pre:  []string{"-exec", "sudo -E"},
+			pkgs: []string{"./cmd/testwrapper"},
+		},
+		{
+			name: "test-arg",
+			in:   []string{"-exec", "sudo -E", "./cmd/testwrapper", "--", "--some-flag"},
+			pre:  []string{"-exec", "sudo -E"},
+			pkgs: []string{"./cmd/testwrapper"},
+			post: []string{"--", "--some-flag"},
+		},
+		{
+			name: "dupe-args",
+			in:   []string{"-v", "-v", "-race", "-race", "./cmd/testwrapper", "--", "--some-flag"},
+			pre:  []string{"-race", "-v"},
+			pkgs: []string{"./cmd/testwrapper"},
+			post: []string{"--", "--some-flag"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pre, pkgs, post, err := splitArgs(tt.in)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !slices.Equal(pre, tt.pre) {
+				t.Errorf("pre = %q; want %q", pre, tt.pre)
+			}
+			if !slices.Equal(pkgs, tt.pkgs) {
+				t.Errorf("pattern = %q; want %q", pkgs, tt.pkgs)
+			}
+			if !slices.Equal(post, tt.post) {
+				t.Errorf("post = %q; want %q", post, tt.post)
+			}
+			if t.Failed() {
+				t.Logf("SplitArgs(%q) = %q %q %q", tt.in, pre, pkgs, post)
+			}
+		})
+	}
+}

cmd/testwrapper/testwrapper.go

@@ -13,7 +13,6 @@
 	"context"
 	"encoding/json"
 	"errors"
-	"flag"
 	"fmt"
 	"io"
 	"log"
@@ -71,18 +70,17 @@ type goTestOutput struct {
 // set to true. Package build errors will not emit a testAttempt (as no valid
 // JSON is produced) but the [os/exec.ExitError] will be returned.
 // It calls close(ch) when it's done.
-func runTests(ctx context.Context, attempt int, pt *packageTests, otherArgs []string, ch chan<- *testAttempt) error {
+func runTests(ctx context.Context, attempt int, pt *packageTests, goTestArgs, testArgs []string, ch chan<- *testAttempt) error {
 	defer close(ch)
-	args := []string{"test", "--json"}
+	args := []string{"test"}
-	if *flagSudo {
+	args = append(args, goTestArgs...)
-		args = append(args, "--exec", "sudo -E")
-	}
 	args = append(args, pt.Pattern)
-	args = append(args, otherArgs...)
 	if len(pt.Tests) > 0 {
 		runArg := strings.Join(pt.Tests, "|")
 		args = append(args, "--run", runArg)
 	}
+	args = append(args, testArgs...)
+	args = append(args, "-json")
 	if debug {
 		fmt.Println("running", strings.Join(args, " "))
 	}
@@ -181,59 +179,29 @@ func runTests(ctx context.Context, attempt int, pt *packageTests, otherArgs []st
 	return nil
 }
 
-var (
-	flagVerbose = flag.Bool("v", false, "verbose")
-	flagSudo    = flag.Bool("sudo", false, "run tests with -exec=sudo")
-)
-
 func main() {
+	goTestArgs, packages, testArgs, err := splitArgs(os.Args[1:])
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+	if len(packages) == 0 {
+		fmt.Println("testwrapper: no packages specified")
+		return
+	}
+
 	ctx := context.Background()
-
-	// We only need to parse the -v flag to figure out whether to print the logs
-	// for a test. We don't need to parse any other flags, so we just use the
-	// flag package to parse the -v flag and then pass the rest of the args
-	// through to 'go test'.
-	// We run `go test -json` which returns the same information as `go test -v`,
-	// but in a machine-readable format. So this flag is only for testwrapper's
-	// output.
-
-	flag.Usage = func() {
-		fmt.Println("usage: testwrapper [testwrapper-flags] [pattern] [build/test flags & test binary flags]")
-		fmt.Println()
-		fmt.Println("testwrapper-flags:")
-		flag.CommandLine.PrintDefaults()
-		fmt.Println()
-		fmt.Println("examples:")
-		fmt.Println("\ttestwrapper -v ./... -count=1")
-		fmt.Println("\ttestwrapper ./pkg/foo -run TestBar -count=1")
-		fmt.Println()
-		fmt.Println("Unlike 'go test', testwrapper requires a package pattern as the first positional argument and only supports a single pattern.")
-	}
-	flag.Parse()
-
-	args := flag.Args()
-	if len(args) < 1 || strings.HasPrefix(args[0], "-") {
-		fmt.Println("no pattern specified")
-		flag.Usage()
-		os.Exit(1)
-	} else if len(args) > 1 && !strings.HasPrefix(args[1], "-") {
-		fmt.Println("expected single pattern")
-		flag.Usage()
-		os.Exit(1)
-	}
-	pattern, otherArgs := args[0], args[1:]
-
 	type nextRun struct {
 		tests   []*packageTests
 		attempt int // starting at 1
 	}
-
+	firstRun := &nextRun{
-	toRun := []*nextRun{
+		attempt: 1,
-		{
-			tests:   []*packageTests{{Pattern: pattern}},
-			attempt: 1,
-		},
 	}
+	for _, pkg := range packages {
+		firstRun.tests = append(firstRun.tests, &packageTests{Pattern: pkg})
+	}
+	toRun := []*nextRun{firstRun}
 	printPkgOutcome := func(pkg, outcome string, attempt int) {
 		if outcome == "skip" {
 			fmt.Printf("?\t%s [skipped/no tests] \n", pkg)
@@ -271,7 +239,7 @@ type nextRun struct {
 			runErr := make(chan error, 1)
 			go func() {
 				defer close(runErr)
-				runErr <- runTests(ctx, thisRun.attempt, pt, otherArgs, ch)
+				runErr <- runTests(ctx, thisRun.attempt, pt, goTestArgs, testArgs, ch)
 			}()
 
 			var failed bool
@@ -281,7 +249,7 @@ type nextRun struct {
 				// convenient for us to to specify files in tests, so fix tr.pkg
 				// so that subsequent testwrapper attempts run correctly.
 				if tr.pkg == "command-line-arguments" {
-					tr.pkg = pattern
+					tr.pkg = packages[0]
 				}
 				if tr.pkgFinished {
 					if tr.outcome == "fail" && len(toRetry[tr.pkg]) == 0 {
@@ -293,7 +261,7 @@ type nextRun struct {
 					printPkgOutcome(tr.pkg, tr.outcome, thisRun.attempt)
 					continue
 				}
-				if *flagVerbose || tr.outcome == "fail" {
+				if testingVerbose || tr.outcome == "fail" {
 					io.Copy(os.Stdout, &tr.logs)
 				}
 				if tr.outcome != "fail" {