From 856d32b4a9590a1d8eb6d94349d75644cc837838 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 28 Sep 2023 16:30:12 -0700
Subject: [PATCH] cmd/testwrapper: include flake URL in JSON metadata

Updates tailscale/corp#14975

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/testwrapper/flakytest/flakytest.go |  2 +-
 cmd/testwrapper/testwrapper.go         | 33 +++++++++++++++++++-------
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/cmd/testwrapper/flakytest/flakytest.go b/cmd/testwrapper/flakytest/flakytest.go
index 8b9dcd980..494ed080b 100644
--- a/cmd/testwrapper/flakytest/flakytest.go
+++ b/cmd/testwrapper/flakytest/flakytest.go
@@ -38,7 +38,7 @@ func Mark(t testing.TB, issue string) {
 		// We're being run under cmd/testwrapper so send our sentinel message
 		// to stderr. (We avoid doing this when the env is absent to avoid
 		// spamming people running tests without the wrapper)
-		fmt.Fprintln(os.Stderr, FlakyTestLogMessage)
+		fmt.Fprintf(os.Stderr, "%s: %s\n", FlakyTestLogMessage, issue)
 	}
 	t.Logf("flakytest: issue tracking this flaky test: %s", issue)
 }
diff --git a/cmd/testwrapper/testwrapper.go b/cmd/testwrapper/testwrapper.go
index 5d4d58f20..7f1342a4f 100644
--- a/cmd/testwrapper/testwrapper.go
+++ b/cmd/testwrapper/testwrapper.go
@@ -24,6 +24,7 @@ import (
 	"time"
 
 	xmaps "golang.org/x/exp/maps"
+	"golang.org/x/exp/slices"
 	"tailscale.com/cmd/testwrapper/flakytest"
 )
 
@@ -34,11 +35,16 @@ type testAttempt struct {
 	testName      string // "TestFoo"
 	outcome       string // "pass", "fail", "skip"
 	logs          bytes.Buffer
-	isMarkedFlaky bool // set if the test is marked as flaky
+	isMarkedFlaky bool   // set if the test is marked as flaky
+	issueURL      string // set if the test is marked as flaky
 
 	pkgFinished bool
 }
 
+// packageTests describes what to run.
+// It's also JSON-marshalled to output for analysys tools to parse
+// so the fields are all exported.
+// TODO(bradfitz): move this type to its own types package?
 type packageTests struct {
 	// Pattern is the package Pattern to run.
 	// Must be a single Pattern, not a list of patterns.
@@ -46,6 +52,8 @@ type packageTests struct {
 	// Tests is a list of Tests to run. If empty, all Tests in the package are
 	// run.
 	Tests []string // ["TestFoo", "TestBar"]
+	// IssueURLs maps from a test name to a URL tracking its flake.
+	IssueURLs map[string]string // "TestFoo" => "https://github.com/foo/bar/issue/123"
 }
 
 type goTestOutput struct {
@@ -152,8 +160,9 @@ func runTests(ctx context.Context, attempt int, pt *packageTests, otherArgs []st
 			pkgTests[testName].outcome = goOutput.Action
 			ch <- pkgTests[testName]
 		case "output":
-			if strings.TrimSpace(goOutput.Output) == flakytest.FlakyTestLogMessage {
+			if suffix, ok := strings.CutPrefix(strings.TrimSpace(goOutput.Output), flakytest.FlakyTestLogMessage); ok {
 				pkgTests[testName].isMarkedFlaky = true
+				pkgTests[testName].issueURL = strings.TrimPrefix(suffix, ": ")
 			} else {
 				pkgTests[testName].logs.WriteString(goOutput.Output)
 			}
@@ -244,12 +253,11 @@ func main() {
 			os.Exit(1)
 		}
 		if thisRun.attempt > 1 {
-			fmt.Printf("\n\nAttempt #%d: Retrying flaky tests:\n\n", thisRun.attempt)
 			j, _ := json.Marshal(thisRun.tests)
-			fmt.Printf("\n\nflakytest failures JSON: %s\n\n", j)
+			fmt.Printf("\n\nAttempt #%d: Retrying flaky tests:\n\nflakytest failures JSON: %s\n\n", thisRun.attempt, j)
 		}
 
-		toRetry := make(map[string][]string) // pkg -> tests to retry
+		toRetry := make(map[string][]*testAttempt) // pkg -> tests to retry
 		for _, pt := range thisRun.tests {
 			ch := make(chan *testAttempt)
 			runErr := make(chan error, 1)
@@ -284,7 +292,7 @@ func main() {
 					continue
 				}
 				if tr.isMarkedFlaky {
-					toRetry[tr.pkg] = append(toRetry[tr.pkg], tr.testName)
+					toRetry[tr.pkg] = append(toRetry[tr.pkg], tr)
 				} else {
 					failed = true
 				}
@@ -317,10 +325,17 @@ func main() {
 		}
 		for _, pkg := range pkgs {
 			tests := toRetry[pkg]
-			sort.Strings(tests)
+			slices.SortFunc(tests, func(a, b *testAttempt) int { return strings.Compare(a.testName, b.testName) })
+			issueURLs := map[string]string{} // test name => URL
+			var testNames []string
+			for _, ta := range tests {
+				issueURLs[ta.testName] = ta.issueURL
+				testNames = append(testNames, ta.testName)
+			}
 			nextRun.tests = append(nextRun.tests, &packageTests{
-				Pattern: pkg,
-				Tests:   tests,
+				Pattern:   pkg,
+				Tests:     testNames,
+				IssueURLs: issueURLs,
 			})
 		}
 		toRun = append(toRun, nextRun)