From e6ab7d3c14d2c7d248cf24061ecd72900bc6d2f3 Mon Sep 17 00:00:00 2001 From: Maisem Ali Date: Sat, 14 Oct 2023 16:56:08 -0700 Subject: [PATCH] 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 --- .github/workflows/test.yml | 2 +- cmd/testwrapper/args.go | 130 +++++++++++++++++++++++++++++++++ cmd/testwrapper/args_test.go | 97 ++++++++++++++++++++++++ cmd/testwrapper/testwrapper.go | 80 ++++++-------------- 4 files changed, 252 insertions(+), 57 deletions(-) create mode 100644 cmd/testwrapper/args.go create mode 100644 cmd/testwrapper/args_test.go diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 965b34ad1..6a0e4a5ed 100644 --- a/.github/workflows/test.yml +++ b/.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: diff --git a/cmd/testwrapper/args.go b/cmd/testwrapper/args.go new file mode 100644 index 000000000..20c64e0e3 --- /dev/null +++ b/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 +}() diff --git a/cmd/testwrapper/args_test.go b/cmd/testwrapper/args_test.go new file mode 100644 index 000000000..10063d7bc --- /dev/null +++ b/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) + } + }) + } +} diff --git a/cmd/testwrapper/testwrapper.go b/cmd/testwrapper/testwrapper.go index c85275a51..bab62f48d 100644 --- a/cmd/testwrapper/testwrapper.go +++ b/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"} - if *flagSudo { - args = append(args, "--exec", "sudo -E") - } + args := []string{"test"} + args = append(args, goTestArgs...) 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 } - - toRun := []*nextRun{ - { - tests: []*packageTests{{Pattern: pattern}}, - attempt: 1, - }, + firstRun := &nextRun{ + 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" {