Add support for custom test commands in RunTestsCmd

Author: lukaszcl

tools/flakeguard/cmd/run.go

@@ -21,6 +21,7 @@ var RunTestsCmd = &cobra.Command{
 		projectPath, _ := cmd.Flags().GetString("project-path")
 		testPackagesJson, _ := cmd.Flags().GetString("test-packages-json")
 		testPackagesArg, _ := cmd.Flags().GetStringSlice("test-packages")
+		testCmdStrings, _ := cmd.Flags().GetStringArray("test-cmd")
 		runCount, _ := cmd.Flags().GetInt("run-count")
 		timeout, _ := cmd.Flags().GetDuration("timeout")
 		tags, _ := cmd.Flags().GetStringArray("tags")
@@ -40,14 +41,17 @@ var RunTestsCmd = &cobra.Command{
 
 		// Determine test packages
 		var testPackages []string
-		if testPackagesJson != "" {
-			if err := json.Unmarshal([]byte(testPackagesJson), &testPackages); err != nil {
-				log.Fatal().Err(err).Msg("Error decoding test packages JSON")
+		if len(testCmdStrings) == 0 {
+			// No custom command -> parse packages
+			if testPackagesJson != "" {
+				if err := json.Unmarshal([]byte(testPackagesJson), &testPackages); err != nil {
+					log.Fatal().Err(err).Msg("Error decoding test packages JSON")
+				}
+			} else if len(testPackagesArg) > 0 {
+				testPackages = testPackagesArg
+			} else {
+				log.Fatal().Msg("Error: must specify either --test-packages-json or --test-packages (or use --test-cmd).")
 			}
-		} else if len(testPackagesArg) > 0 {
-			testPackages = testPackagesArg
-		} else {
-			log.Fatal().Msg("Error: must specify either --test-packages-json or --test-packages")
 		}
 
 		// Initialize the runner
@@ -67,9 +71,22 @@ var RunTestsCmd = &cobra.Command{
 		}
 
 		// Run the tests
-		testReport, err := testRunner.RunTests()
-		if err != nil {
-			log.Fatal().Err(err).Msg("Error running tests")
+		var (
+			testReport *reports.TestReport
+			err        error
+		)
+
+		if len(testCmdStrings) > 0 {
+			testReport, err = testRunner.RunTestsByCmd(testCmdStrings)
+			if err != nil {
+				log.Fatal().Err(err).Msg("Error running custom test command")
+			}
+		} else {
+			// Otherwise, use the normal go test approach
+			testReport, err = testRunner.RunTests()
+			if err != nil {
+				log.Fatal().Err(err).Msg("Error running tests")
+			}
 		}
 
 		// Save the test results in JSON format
@@ -108,6 +125,9 @@ func init() {
 	RunTestsCmd.Flags().StringP("project-path", "r", ".", "The path to the Go project. Default is the current directory. Useful for subprojects")
 	RunTestsCmd.Flags().String("test-packages-json", "", "JSON-encoded string of test packages")
 	RunTestsCmd.Flags().StringSlice("test-packages", nil, "Comma-separated list of test packages to run")
+	RunTestsCmd.Flags().StringArray("test-cmd", nil,
+		"Optional custom test command (e.g. 'go run e2e_test.go -someflag'), which must produce go test -json output.",
+	)
 	RunTestsCmd.Flags().Bool("run-all-packages", false, "Run all test packages in the project. This flag overrides --test-packages and --test-packages-json")
 	RunTestsCmd.Flags().IntP("run-count", "c", 1, "Number of times to run the tests")
 	RunTestsCmd.Flags().Duration("timeout", 0, "Passed on to the 'go test' command as the -timeout flag")

tools/flakeguard/runner/runner.go

@@ -160,6 +160,93 @@ func (r *Runner) runTests(packageName string) (string, bool, error) {
 	return tmpFile.Name(), true, nil // Test succeeded
 }
 
+// RunTestsByCmd runs an arbitrary command testCmd (like ["go", "run", "my_test.go", ...])
+// that produces the same JSON lines that 'go test -json' would produce on stdout.
+// It captures those lines in a temp file, then parses them for pass/fail/panic/race data.
+func (r *Runner) RunTestsByCmd(testCmd []string) (*reports.TestReport, error) {
+	var jsonFilePaths []string
+
+	// Run the command r.RunCount times
+	for i := 0; i < r.RunCount; i++ {
+		jsonFilePath, passed, err := r.runCmd(testCmd, i)
+		if err != nil {
+			return nil, fmt.Errorf("failed to run test command: %w", err)
+		}
+		jsonFilePaths = append(jsonFilePaths, jsonFilePath)
+		if !passed && r.FailFast {
+			break
+		}
+	}
+
+	results, err := r.parseTestResults(jsonFilePaths)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse test results: %w", err)
+	}
+
+	// Build a TestReport, same shape as RunTests()
+	return &reports.TestReport{
+		GoProject:     r.prettyProjectPath,
+		TestRunCount:  r.RunCount,
+		RaceDetection: r.UseRace,
+		Results:       results,
+	}, nil
+}
+
+// runCmd is a helper that runs the user-supplied command once, captures its JSON output,
+// and returns (tempFilePath, passed, error).
+func (r *Runner) runCmd(testCmd []string, runIndex int) (string, bool, error) {
+	// Create temp file for JSON output
+	tmpFile, err := os.CreateTemp("", fmt.Sprintf("test-output-cmd-run%d-*.json", runIndex+1))
+	if err != nil {
+		return "", false, fmt.Errorf("failed to create temp file: %w", err)
+	}
+	defer tmpFile.Close()
+
+	if r.Verbose {
+		log.Info().Msgf("Running custom test command (%d/%d): %s", runIndex+1, r.RunCount, strings.Join(testCmd, " "))
+	}
+
+	cmd := exec.Command(testCmd[0], testCmd[1:]...)
+	cmd.Dir = r.ProjectPath
+
+	// If collecting raw output, write to both file & buffer
+	if r.CollectRawOutput {
+		if r.rawOutputs == nil {
+			r.rawOutputs = make(map[string]*bytes.Buffer)
+		}
+		key := fmt.Sprintf("customCmd-run%d", runIndex+1)
+		if _, exists := r.rawOutputs[key]; !exists {
+			r.rawOutputs[key] = &bytes.Buffer{}
+		}
+		cmd.Stdout = io.MultiWriter(tmpFile, r.rawOutputs[key])
+	} else {
+		cmd.Stdout = tmpFile
+	}
+	// Handle stderr however you like; here we just dump it to console
+	cmd.Stderr = os.Stderr
+
+	// Run it
+	err = cmd.Run()
+
+	// Determine pass/fail from exit code
+	type exitCoder interface {
+		ExitCode() int
+	}
+	var ec exitCoder
+	if errors.As(err, &ec) {
+		// Non-zero exit code => test failure
+		if ec.ExitCode() != 0 {
+			return tmpFile.Name(), false, nil
+		}
+	} else if err != nil {
+		// Some other error that doesn't implement ExitCode() => real error
+		return "", false, fmt.Errorf("error running test command: %w", err)
+	}
+
+	// Otherwise, assume success
+	return tmpFile.Name(), true, nil
+}
+
 type entry struct {
 	Action  string  `json:"Action"`
 	Test    string  `json:"Test"`