From 51b2b2ec76f8938f95791bad6d1f63ec12c85d5b Mon Sep 17 00:00:00 2001
From: yux0 <yx@uber.com>
Date: Wed, 18 Dec 2019 14:10:15 -0800
Subject: [PATCH] Fix merge conflicts

---
 test/integration_test.go | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/test/integration_test.go b/test/integration_test.go
index bd20cfbc8..92fa757f2 100644
--- a/test/integration_test.go
+++ b/test/integration_test.go
@@ -100,11 +100,29 @@ func (ts *IntegrationTestSuite) SetupSuite() {
 
 func (ts *IntegrationTestSuite) TearDownSuite() {
 	ts.rpcClient.Close()
-	// sleep for a while to allow the pollers to shutdown
-	// then assert that there are no lingering go routines
-	time.Sleep(1 * time.Minute)
-	// https://github.com/uber-go/cadence-client/issues/739
-	goleak.VerifyNoLeaks(ts.T(), goleak.IgnoreTopFunction("go.uber.org/cadence/internal.(*coroutineState).initialYield"))
+
+	// allow the pollers to shut down, and ensure there are no goroutine leaks.
+	// this will wait for up to 1 minute for leaks to subside, but exit relatively quickly if possible.
+	max := time.After(time.Minute)
+	var last error
+	for {
+		select {
+		case <-max:
+			if last != nil {
+				ts.NoError(last)
+				return
+			}
+			ts.FailNow("leaks timed out but no error, should be impossible")
+		case <-time.After(time.Second):
+			// https://github.com/uber-go/cadence-client/issues/739
+			last = goleak.FindLeaks(goleak.IgnoreTopFunction("go.uber.org/cadence/internal.(*coroutineState).initialYield"))
+			if last == nil {
+				// no leak, done waiting
+				return
+			}
+			// else wait for another check or the timeout (which will record the latest error)
+		}
+	}
 }
 
 func (ts *IntegrationTestSuite) SetupTest() {