From b00c2a9535a15b621130a65572d3a2b8b8399999 Mon Sep 17 00:00:00 2001
From: 3pointer <luancheng@pingcap.com>
Date: Thu, 6 Jan 2022 12:20:36 +0800
Subject: [PATCH] br: fix unstable unit test (#31317)

close pingcap/tidb#31356
---
 br/pkg/restore/client_test.go | 26 ++++++--------------------
 br/pkg/restore/main_test.go   | 25 +++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/br/pkg/restore/client_test.go b/br/pkg/restore/client_test.go
index 1480aa4eba7b4..4baa028f8ea29 100644
--- a/br/pkg/restore/client_test.go
+++ b/br/pkg/restore/client_test.go
@@ -23,17 +23,15 @@ import (
 	"google.golang.org/grpc/keepalive"
 )
 
+var mc *mock.Cluster
+
 var defaultKeepaliveCfg = keepalive.ClientParameters{
 	Time:    3 * time.Second,
 	Timeout: 10 * time.Second,
 }
 
 func TestCreateTables(t *testing.T) {
-	m, err := mock.NewCluster()
-	require.NoError(t, err)
-	require.NoError(t, m.Start())
-	defer m.Stop()
-
+	m := mc
 	client, err := restore.NewRestoreClient(gluetidb.New(), m.PDClient, m.Storage, nil, defaultKeepaliveCfg)
 	require.NoError(t, err)
 
@@ -89,11 +87,7 @@ func TestCreateTables(t *testing.T) {
 }
 
 func TestIsOnline(t *testing.T) {
-	m, err := mock.NewCluster()
-	require.NoError(t, err)
-	require.NoError(t, m.Start())
-	defer m.Stop()
-
+	m := mc
 	client, err := restore.NewRestoreClient(gluetidb.New(), m.PDClient, m.Storage, nil, defaultKeepaliveCfg)
 	require.NoError(t, err)
 
@@ -103,11 +97,7 @@ func TestIsOnline(t *testing.T) {
 }
 
 func TestPreCheckTableClusterIndex(t *testing.T) {
-	m, err := mock.NewCluster()
-	require.NoError(t, err)
-	require.NoError(t, m.Start())
-	defer m.Stop()
-
+	m := mc
 	client, err := restore.NewRestoreClient(gluetidb.New(), m.PDClient, m.Storage, nil, defaultKeepaliveCfg)
 	require.NoError(t, err)
 
@@ -178,11 +168,7 @@ func (fpdc fakePDClient) GetAllStores(context.Context, ...pd.GetStoreOption) ([]
 }
 
 func TestPreCheckTableTiFlashReplicas(t *testing.T) {
-	m, err := mock.NewCluster()
-	require.NoError(t, err)
-	require.NoError(t, m.Start())
-	defer m.Stop()
-
+	m := mc
 	mockStores := []*metapb.Store{
 		{
 			Id: 1,
diff --git a/br/pkg/restore/main_test.go b/br/pkg/restore/main_test.go
index a29fbdba1868a..0b2cea2e47030 100644
--- a/br/pkg/restore/main_test.go
+++ b/br/pkg/restore/main_test.go
@@ -12,11 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package restore
+package restore_test
 
 import (
+	"fmt"
+	"os"
 	"testing"
 
+	"github.com/pingcap/tidb/br/pkg/mock"
 	"github.com/pingcap/tidb/util/testbridge"
 	"go.uber.org/goleak"
 )
@@ -28,5 +31,23 @@ func TestMain(m *testing.M) {
 		goleak.IgnoreTopFunction("go.etcd.io/etcd/pkg/logutil.(*MergeLogger).outputLoop"),
 		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
 	}
-	goleak.VerifyTestMain(m, opts...)
+
+	var err error
+	mc, err = mock.NewCluster()
+	if err != nil {
+		panic(err)
+	}
+	err = mc.Start()
+	if err != nil {
+		panic(err)
+	}
+	exitCode := m.Run()
+	mc.Stop()
+	if exitCode == 0 {
+		if err := goleak.Find(opts...); err != nil {
+			fmt.Fprintf(os.Stderr, "goleak: Errors on successful test run: %v\n", err)
+			exitCode = 1
+		}
+	}
+	os.Exit(exitCode)
 }