From eb6a47f87ee760f5004e03efa0765ca3739f755b Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Fri, 10 Mar 2017 12:18:24 -0800
Subject: [PATCH] testutil: add CheckAfterTest for calling AfterTest without a
 testing.T

---
 pkg/testutil/leak.go | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/pkg/testutil/leak.go b/pkg/testutil/leak.go
index 80bc0eebc..874dfcc21 100644
--- a/pkg/testutil/leak.go
+++ b/pkg/testutil/leak.go
@@ -62,10 +62,11 @@ func CheckLeakedGoroutine() bool {
 	return true
 }
 
-func AfterTest(t *testing.T) {
+// CheckAfterTest returns an error if AfterTest would fail with an error.
+func CheckAfterTest(d time.Duration) error {
 	http.DefaultTransport.(*http.Transport).CloseIdleConnections()
 	if testing.Short() {
-		return
+		return nil
 	}
 	var bad string
 	badSubstring := map[string]string{
@@ -78,7 +79,8 @@ func AfterTest(t *testing.T) {
 	}
 
 	var stacks string
-	for i := 0; i < 6; i++ {
+	begin := time.Now()
+	for time.Since(begin) < d {
 		bad = ""
 		stacks = strings.Join(interestingGoroutines(), "\n\n")
 		for substr, what := range badSubstring {
@@ -87,13 +89,22 @@ func AfterTest(t *testing.T) {
 			}
 		}
 		if bad == "" {
-			return
+			return nil
 		}
 		// Bad stuff found, but goroutines might just still be
 		// shutting down, so give it some time.
 		time.Sleep(50 * time.Millisecond)
 	}
-	t.Errorf("Test appears to have leaked %s:\n%s", bad, stacks)
+	return fmt.Errorf("appears to have leaked %s:\n%s", bad, stacks)
+}
+
+// AfterTest is meant to run in a defer that executes after a test completes.
+// It will detect common goroutine leaks, retrying in case there are goroutines
+// not synchronously torn down, and fail the test if any goroutines are stuck.
+func AfterTest(t *testing.T) {
+	if err := CheckAfterTest(300 * time.Millisecond); err != nil {
+		t.Errorf("Test %v", err)
+	}
 }
 
 func interestingGoroutines() (gs []string) {