From d1b0e1af06d9c44c71413410dcf2527f349d4767 Mon Sep 17 00:00:00 2001
From: Paul Scott <408401+icio@users.noreply.github.com>
Date: Mon, 24 Feb 2025 13:26:41 -0800
Subject: [PATCH] cmd/testwrapper/flakytest: add Marked to check if in
 flakytest (#15119)

Updates tailscale/corp#26637

Signed-off-by: Paul Scott <paul@tailscale.com>
---
 cmd/testwrapper/flakytest/flakytest.go      | 29 +++++++++++++
 cmd/testwrapper/flakytest/flakytest_test.go | 46 +++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/cmd/testwrapper/flakytest/flakytest.go b/cmd/testwrapper/flakytest/flakytest.go
index 494ed080b..6302900cb 100644
--- a/cmd/testwrapper/flakytest/flakytest.go
+++ b/cmd/testwrapper/flakytest/flakytest.go
@@ -9,8 +9,12 @@ package flakytest
 import (
 	"fmt"
 	"os"
+	"path"
 	"regexp"
+	"sync"
 	"testing"
+
+	"tailscale.com/util/mak"
 )
 
 // FlakyTestLogMessage is a sentinel value that is printed to stderr when a
@@ -25,6 +29,11 @@ const FlakeAttemptEnv = "TS_TESTWRAPPER_ATTEMPT"
 
 var issueRegexp = regexp.MustCompile(`\Ahttps://github\.com/tailscale/[a-zA-Z0-9_.-]+/issues/\d+\z`)
 
+var (
+	rootFlakesMu sync.Mutex
+	rootFlakes   map[string]bool
+)
+
 // Mark sets the current test as a flaky test, such that if it fails, it will
 // be retried a few times on failure. issue must be a GitHub issue that tracks
 // the status of the flaky test being marked, of the format:
@@ -41,4 +50,24 @@ func Mark(t testing.TB, issue string) {
 		fmt.Fprintf(os.Stderr, "%s: %s\n", FlakyTestLogMessage, issue)
 	}
 	t.Logf("flakytest: issue tracking this flaky test: %s", issue)
+
+	// Record the root test name as flakey.
+	rootFlakesMu.Lock()
+	defer rootFlakesMu.Unlock()
+	mak.Set(&rootFlakes, t.Name(), true)
+}
+
+// Marked reports whether the current test or one of its parents was marked flaky.
+func Marked(t testing.TB) bool {
+	n := t.Name()
+	for {
+		if rootFlakes[n] {
+			return true
+		}
+		n = path.Dir(n)
+		if n == "." || n == "/" {
+			break
+		}
+	}
+	return false
 }
diff --git a/cmd/testwrapper/flakytest/flakytest_test.go b/cmd/testwrapper/flakytest/flakytest_test.go
index 85e77a939..64cbfd9a3 100644
--- a/cmd/testwrapper/flakytest/flakytest_test.go
+++ b/cmd/testwrapper/flakytest/flakytest_test.go
@@ -41,3 +41,49 @@ func TestFlakeRun(t *testing.T) {
 		t.Fatal("First run in testwrapper, failing so that test is retried. This is expected.")
 	}
 }
+
+func TestMarked_Root(t *testing.T) {
+	Mark(t, "https://github.com/tailscale/tailscale/issues/0")
+
+	t.Run("child", func(t *testing.T) {
+		t.Run("grandchild", func(t *testing.T) {
+			if got, want := Marked(t), true; got != want {
+				t.Fatalf("Marked(t) = %t, want %t", got, want)
+			}
+		})
+
+		if got, want := Marked(t), true; got != want {
+			t.Fatalf("Marked(t) = %t, want %t", got, want)
+		}
+	})
+
+	if got, want := Marked(t), true; got != want {
+		t.Fatalf("Marked(t) = %t, want %t", got, want)
+	}
+}
+
+func TestMarked_Subtest(t *testing.T) {
+	t.Run("flaky", func(t *testing.T) {
+		Mark(t, "https://github.com/tailscale/tailscale/issues/0")
+
+		t.Run("child", func(t *testing.T) {
+			t.Run("grandchild", func(t *testing.T) {
+				if got, want := Marked(t), true; got != want {
+					t.Fatalf("Marked(t) = %t, want %t", got, want)
+				}
+			})
+
+			if got, want := Marked(t), true; got != want {
+				t.Fatalf("Marked(t) = %t, want %t", got, want)
+			}
+		})
+
+		if got, want := Marked(t), true; got != want {
+			t.Fatalf("Marked(t) = %t, want %t", got, want)
+		}
+	})
+
+	if got, want := Marked(t), false; got != want {
+		t.Fatalf("Marked(t) = %t, want %t", got, want)
+	}
+}