From 18318aa45ff3c54ff10a5fc154bcd8930b34c93c Mon Sep 17 00:00:00 2001
From: Lukasz <120112546+lukaszcl@users.noreply.github.com>
Date: Mon, 25 Nov 2024 15:43:06 +0100
Subject: [PATCH] Add shuffle option to flakeguard (#1379)

---
 tools/flakeguard/cmd/run.go       | 6 ++++++
 tools/flakeguard/runner/runner.go | 9 +++++++++
 2 files changed, 15 insertions(+)

diff --git a/tools/flakeguard/cmd/run.go b/tools/flakeguard/cmd/run.go
index 43d6b0be2..8ff3db6ec 100644
--- a/tools/flakeguard/cmd/run.go
+++ b/tools/flakeguard/cmd/run.go
@@ -26,6 +26,8 @@ var RunTestsCmd = &cobra.Command{
 		maxPassRatio, _ := cmd.Flags().GetFloat64("max-pass-ratio")
 		skipTests, _ := cmd.Flags().GetStringSlice("skip-tests")
 		printFailedTests, _ := cmd.Flags().GetBool("print-failed-tests")
+		useShuffle, _ := cmd.Flags().GetBool("shuffle")
+		shuffleSeed, _ := cmd.Flags().GetString("shuffle-seed")
 
 		// Check if project dependencies are correctly set up
 		if err := checkDependencies(projectPath); err != nil {
@@ -50,6 +52,8 @@ var RunTestsCmd = &cobra.Command{
 			UseRace:              useRace,
 			SkipTests:            skipTests,
 			SelectedTestPackages: testPackages,
+			UseShuffle:           useShuffle,
+			ShuffleSeed:          shuffleSeed,
 		}
 
 		testResults, err := runner.RunTests()
@@ -101,6 +105,8 @@ func init() {
 	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().Bool("race", false, "Enable the race detector")
+	RunTestsCmd.Flags().Bool("shuffle", false, "Enable test shuffling")
+	RunTestsCmd.Flags().String("shuffle-seed", "", "Set seed for test shuffling. Must be used with --shuffle")
 	RunTestsCmd.Flags().Bool("fail-fast", false, "Stop on the first test failure")
 	RunTestsCmd.Flags().String("output-json", "", "Path to output the test results in JSON format")
 	RunTestsCmd.Flags().StringSlice("skip-tests", nil, "Comma-separated list of test names to skip from running")
diff --git a/tools/flakeguard/runner/runner.go b/tools/flakeguard/runner/runner.go
index 9e5e3564e..4dc8980a6 100644
--- a/tools/flakeguard/runner/runner.go
+++ b/tools/flakeguard/runner/runner.go
@@ -29,6 +29,8 @@ type Runner struct {
 	Verbose              bool     // If true, provides detailed logging.
 	RunCount             int      // Number of times to run the tests.
 	UseRace              bool     // Enable race detector.
+	UseShuffle           bool     // Enable test shuffling. -shuffle=on flag.
+	ShuffleSeed          string   // Set seed for test shuffling -shuffle={seed} flag. Must be used with UseShuffle.
 	FailFast             bool     // Stop on first test failure.
 	SkipTests            []string // Test names to exclude.
 	SelectedTestPackages []string // Explicitly selected packages to run.
@@ -93,6 +95,13 @@ func (r *Runner) runTests(packageName string) (string, bool, error) {
 	if r.UseRace {
 		args = append(args, "-race")
 	}
+	if r.UseShuffle {
+		if r.ShuffleSeed != "" {
+			args = append(args, fmt.Sprintf("-shuffle=%s", r.ShuffleSeed))
+		} else {
+			args = append(args, "-shuffle=on")
+		}
+	}
 	if len(r.SkipTests) > 0 {
 		skipPattern := strings.Join(r.SkipTests, "|")
 		args = append(args, fmt.Sprintf("-skip=%s", skipPattern))