From c12ec8ffba5e200e348106218c6362a9cd634ace Mon Sep 17 00:00:00 2001
From: Sung Yoon Whang <sungyoon@uber.com>
Date: Mon, 30 Sep 2024 17:45:34 -0700
Subject: [PATCH] Add RunOnFailure option

This adds an Option to run goleak when test runs fail.

Fixes #128.
---
 leaks.go         |  3 +++
 options.go       | 18 ++++++++++++++----
 testmain.go      | 22 +++++++++++++++++++---
 testmain_test.go |  4 ++++
 4 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/leaks.go b/leaks.go
index cc206f1..e0384f8 100644
--- a/leaks.go
+++ b/leaks.go
@@ -59,6 +59,9 @@ func Find(options ...Option) error {
 	if opts.cleanup != nil {
 		return errors.New("Cleanup can only be passed to VerifyNone or VerifyTestMain")
 	}
+	if opts.runOnFailure {
+		return errors.New("RunOnFailure can only be passed to VerifyTestMain")
+	}
 	var stacks []stack.Stack
 	retry := true
 	for i := 0; retry; i++ {
diff --git a/options.go b/options.go
index d18abd9..2f76c05 100644
--- a/options.go
+++ b/options.go
@@ -38,10 +38,11 @@ type Option interface {
 const _defaultRetries = 20
 
 type opts struct {
-	filters    []func(stack.Stack) bool
-	maxRetries int
-	maxSleep   time.Duration
-	cleanup    func(int)
+	filters      []func(stack.Stack) bool
+	maxRetries   int
+	maxSleep     time.Duration
+	cleanup      func(int)
+	runOnFailure bool
 }
 
 // implement apply so that opts struct itself can be used as
@@ -51,6 +52,7 @@ func (o *opts) apply(opts *opts) {
 	opts.maxRetries = o.maxRetries
 	opts.maxSleep = o.maxSleep
 	opts.cleanup = o.cleanup
+	opts.runOnFailure = o.runOnFailure
 }
 
 // optionFunc lets us easily write options without a custom type.
@@ -107,6 +109,14 @@ func IgnoreCurrent() Option {
 	})
 }
 
+// RunOnFailure makes goleak look for leaking goroutines upon test failures.
+// By default goleak only looks for leaking goroutines when tests succeed.
+func RunOnFailure() Option {
+	return optionFunc(func(opts *opts) {
+		opts.runOnFailure = true
+	})
+}
+
 func maxSleep(d time.Duration) Option {
 	return optionFunc(func(opts *opts) {
 		opts.maxSleep = d
diff --git a/testmain.go b/testmain.go
index d883124..7185ef4 100644
--- a/testmain.go
+++ b/testmain.go
@@ -60,10 +60,26 @@ func VerifyTestMain(m TestingM, options ...Option) {
 	}
 	defer func() { cleanup(exitCode) }()
 
-	if exitCode == 0 {
+	var (
+		run      bool
+		errorMsg string
+	)
+
+	if !opts.runOnFailure && exitCode == 0 {
+		errorMsg = "goleak: Errors on successful test run:%v\n"
+		run = true
+	} else if opts.runOnFailure {
+		errorMsg = "goleak: Errors on unsuccessful test run: %v\n"
+		run = true
+	}
+
+	if run {
 		if err := Find(opts); err != nil {
-			fmt.Fprintf(_osStderr, "goleak: Errors on successful test run: %v\n", err)
-			exitCode = 1
+			fmt.Fprintf(_osStderr, errorMsg, err)
+			// rewrite exitCode if test passed and is set to 0.
+			if exitCode == 0 {
+				exitCode = 1
+			}
 		}
 	}
 }
diff --git a/testmain_test.go b/testmain_test.go
index 5945c5c..7b7dd00 100644
--- a/testmain_test.go
+++ b/testmain_test.go
@@ -67,6 +67,10 @@ func TestVerifyTestMain(t *testing.T) {
 	assert.Equal(t, 7, <-exitCode, "Exit code should not be modified")
 	assert.NotContains(t, <-stderr, "goleak: Errors", "Ignore leaks on unsuccessful runs")
 
+	VerifyTestMain(dummyTestMain(7), RunOnFailure())
+	assert.Equal(t, 7, <-exitCode, "Exit code should not be modified")
+	assert.Contains(t, <-stderr, "goleak: Errors", "Find leaks on unsuccessful runs with RunOnFailure specified")
+
 	VerifyTestMain(dummyTestMain(0))
 	assert.Equal(t, 1, <-exitCode, "Expect error due to leaks on successful runs")
 	assert.Contains(t, <-stderr, "goleak: Errors", "Find leaks on successful runs")