diff --git a/leaks.go b/leaks.go index ace2e2b..ee122b7 100644 --- a/leaks.go +++ b/leaks.go @@ -21,6 +21,7 @@ package goleak import ( + "errors" "fmt" "go.uber.org/goleak/internal/stack" @@ -55,6 +56,9 @@ func Find(options ...Option) error { cur := stack.Current().ID() opts := buildOpts(options...) + if opts.cleanup != nil { + return errors.New("Cleanup can only be passed to VerifyNone or VerifyTestMain") + } var stacks []stack.Stack retry := true for i := 0; retry; i++ { @@ -79,12 +83,20 @@ type testHelper interface { // // defer VerifyNone(t) func VerifyNone(t TestingT, options ...Option) { + opts := buildOpts(options...) + var cleanup func(int) + cleanup, opts.cleanup = opts.cleanup, nil + if h, ok := t.(testHelper); ok { // Mark this function as a test helper, if available. h.Helper() } - if err := Find(options...); err != nil { + if err := Find(opts); err != nil { t.Error(err) } + + if cleanup != nil { + cleanup(0) + } } diff --git a/leaks_test.go b/leaks_test.go index f80a4ea..cf24585 100644 --- a/leaks_test.go +++ b/leaks_test.go @@ -40,17 +40,26 @@ func testOptions() Option { } func TestFind(t *testing.T) { - require.NoError(t, Find(), "Should find no leaks by default") + t.Run("Should find no leaks by default", func(t *testing.T) { + require.NoError(t, Find()) + }) - bg := startBlockedG() - err := Find(testOptions()) - require.Error(t, err, "Should find leaks with leaked goroutine") - assert.Contains(t, err.Error(), "blockedG") - assert.Contains(t, err.Error(), "created by go.uber.org/goleak.startBlockedG") - - // Once we unblock the goroutine, we shouldn't have leaks. - bg.unblock() - require.NoError(t, Find(), "Should find no leaks by default") + t.Run("Find leaks with leaked goroutine", func(t *testing.T) { + bg := startBlockedG() + err := Find(testOptions()) + require.Error(t, err, "Should find leaks with leaked goroutine") + assert.Contains(t, err.Error(), "blockedG") + assert.Contains(t, err.Error(), "created by go.uber.org/goleak.startBlockedG") + + // Once we unblock the goroutine, we shouldn't have leaks. + bg.unblock() + require.NoError(t, Find(), "Should find no leaks by default") + }) + + t.Run("Find can't take in Cleanup option", func(t *testing.T) { + err := Find(Cleanup(func(int) { assert.Fail(t, "this should not be called") })) + require.Error(t, err, "Should exit with invalid option") + }) } func TestFindRetry(t *testing.T) { @@ -74,14 +83,26 @@ func (ft *fakeT) Error(args ...interface{}) { } func TestVerifyNone(t *testing.T) { - ft := &fakeT{} - VerifyNone(ft) - require.Empty(t, ft.errors, "Expect no errors from VerifyNone") + t.Run("VerifyNone finds leaks", func(t *testing.T) { + ft := &fakeT{} + VerifyNone(ft) + require.Empty(t, ft.errors, "Expect no errors from VerifyNone") + + bg := startBlockedG() + VerifyNone(ft, testOptions()) + require.NotEmpty(t, ft.errors, "Expect errors from VerifyNone on leaked goroutine") + bg.unblock() + }) - bg := startBlockedG() - VerifyNone(ft, testOptions()) - require.NotEmpty(t, ft.errors, "Expect errors from VerifyNone on leaked goroutine") - bg.unblock() + t.Run("cleanup registered callback should be called", func(t *testing.T) { + ft := &fakeT{} + cleanupCalled := false + VerifyNone(ft, Cleanup(func(c int) { + assert.Equal(t, 0, c) + cleanupCalled = true + })) + require.True(t, cleanupCalled, "expect cleanup registered callback to be called") + }) } func TestIgnoreCurrent(t *testing.T) { diff --git a/options.go b/options.go index 33266f6..d2d473b 100644 --- a/options.go +++ b/options.go @@ -41,6 +41,16 @@ type opts struct { filters []func(stack.Stack) bool maxRetries int maxSleep time.Duration + cleanup func(int) +} + +// implement apply so that opts struct itself can be used as +// an Option. +func (o *opts) apply(opts *opts) { + opts.filters = o.filters + opts.maxRetries = o.maxRetries + opts.maxSleep = o.maxSleep + opts.cleanup = o.cleanup } // optionFunc lets us easily write options without a custom type. @@ -57,6 +67,18 @@ func IgnoreTopFunction(f string) Option { }) } +// Cleanup sets up a cleanup function that will be executed at the +// end of the leak check. +// When passed to [VerifyTestMain], the exit code passed to cleanupFunc +// will be set to the exit code of TestMain. +// When passed to [VerifyNone], the exit code will be set to 0. +// This cannot be passed to [Find]. +func Cleanup(cleanupFunc func(exitCode int)) Option { + return optionFunc(func(opts *opts) { + opts.cleanup = cleanupFunc + }) +} + // IgnoreCurrent records all current goroutines when the option is created, and ignores // them in any future Find/Verify calls. func IgnoreCurrent() Option { @@ -98,8 +120,8 @@ func buildOpts(options ...Option) *opts { return opts } -func (vo *opts) filter(s stack.Stack) bool { - for _, filter := range vo.filters { +func (o *opts) filter(s stack.Stack) bool { + for _, filter := range o.filters { if filter(s) { return true } @@ -107,14 +129,14 @@ func (vo *opts) filter(s stack.Stack) bool { return false } -func (vo *opts) retry(i int) bool { - if i >= vo.maxRetries { +func (o *opts) retry(i int) bool { + if i >= o.maxRetries { return false } d := time.Duration(int(time.Microsecond) << uint(i)) - if d > vo.maxSleep { - d = vo.maxSleep + if d > o.maxSleep { + d = o.maxSleep } time.Sleep(d) return true diff --git a/testmain.go b/testmain.go index e7d190b..7b1a50b 100644 --- a/testmain.go +++ b/testmain.go @@ -51,13 +51,19 @@ type TestingM interface { // for any goroutine leaks and fail the tests if any leaks were found. func VerifyTestMain(m TestingM, options ...Option) { exitCode := m.Run() + opts := buildOpts(options...) + + var cleanup func(int) + cleanup, opts.cleanup = opts.cleanup, nil + if cleanup == nil { + cleanup = _osExit + } + defer func() { cleanup(exitCode) }() if exitCode == 0 { - if err := Find(options...); err != nil { + if err := Find(opts); err != nil { fmt.Fprintf(_osStderr, "goleak: Errors on successful test run: %v\n", err) exitCode = 1 } } - - _osExit(exitCode) } diff --git a/testmain_test.go b/testmain_test.go index af5046b..5945c5c 100644 --- a/testmain_test.go +++ b/testmain_test.go @@ -75,4 +75,13 @@ func TestVerifyTestMain(t *testing.T) { VerifyTestMain(dummyTestMain(0)) assert.Equal(t, 0, <-exitCode, "Expect no errors without leaks") assert.NotContains(t, <-stderr, "goleak: Errors", "No errors on successful run without leaks") + + cleanupCalled := false + cleanupExitcode := 0 + VerifyTestMain(dummyTestMain(3), Cleanup(func(ec int) { + cleanupCalled = true + cleanupExitcode = ec + })) + assert.True(t, cleanupCalled) + assert.Equal(t, 3, cleanupExitcode) }