From f1490ea85b2d343f92ed4c2712a2d332b6b9cdee Mon Sep 17 00:00:00 2001
From: tison <wander4096@gmail.com>
Date: Sat, 4 Dec 2021 16:48:36 +0800
Subject: [PATCH] expression: migrate test-infra to testify for
 flag_simplify_test.go

This closes #29136.

Signed-off-by: tison <wander4096@gmail.com>
---
 expression/flag_simplify_test.go | 48 ++++++++++----------------------
 expression/main_test.go          | 15 +++++++++-
 2 files changed, 28 insertions(+), 35 deletions(-)

diff --git a/expression/flag_simplify_test.go b/expression/flag_simplify_test.go
index a9f23e19ea099..a3900f4a31149 100644
--- a/expression/flag_simplify_test.go
+++ b/expression/flag_simplify_test.go
@@ -15,41 +15,20 @@
 package expression_test
 
 import (
-	. "github.com/pingcap/check"
-	"github.com/pingcap/tidb/domain"
-	"github.com/pingcap/tidb/kv"
-	"github.com/pingcap/tidb/sessionctx"
-	"github.com/pingcap/tidb/util/mock"
-	"github.com/pingcap/tidb/util/testkit"
-	"github.com/pingcap/tidb/util/testutil"
-)
-
-var _ = Suite(&testFlagSimplifySuite{})
+	"testing"
 
-type testFlagSimplifySuite struct {
-	store    kv.Storage
-	dom      *domain.Domain
-	ctx      sessionctx.Context
-	testData testutil.TestData
-}
+	"github.com/pingcap/tidb/expression"
+	"github.com/pingcap/tidb/testkit"
+	"github.com/pingcap/tidb/testkit/testdata"
+)
 
-func (s *testFlagSimplifySuite) SetUpSuite(c *C) {
-	var err error
-	s.store, s.dom, err = newStoreWithBootstrap()
-	c.Assert(err, IsNil)
-	s.ctx = mock.NewContext()
-	s.testData, err = testutil.LoadTestSuiteData("testdata", "flag_simplify")
-	c.Assert(err, IsNil)
-}
+func TestSimplifyExpressionByFlag(t *testing.T) {
+	t.Parallel()
 
-func (s *testFlagSimplifySuite) TearDownSuite(c *C) {
-	c.Assert(s.testData.GenerateOutputIfNeeded(), IsNil)
-	s.dom.Close()
-	s.store.Close()
-}
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
 
-func (s *testFlagSimplifySuite) TestSimplifyExpressionByFlag(c *C) {
-	tk := testkit.NewTestKit(c, s.store)
+	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
 	tk.MustExec("drop table if exists t")
 	tk.MustExec("create table t(id int primary key, a bigint unsigned not null, b bigint unsigned)")
@@ -59,11 +38,12 @@ func (s *testFlagSimplifySuite) TestSimplifyExpressionByFlag(c *C) {
 		SQL  string
 		Plan []string
 	}
-	s.testData.GetTestCases(c, &input, &output)
+	flagSimplifyData := expression.GetFlagSimplifyData()
+	flagSimplifyData.GetTestCases(t, &input, &output)
 	for i, tt := range input {
-		s.testData.OnRecord(func() {
+		testdata.OnRecord(func() {
 			output[i].SQL = tt
-			output[i].Plan = s.testData.ConvertRowsToStrings(tk.MustQuery(tt).Rows())
+			output[i].Plan = testdata.ConvertRowsToStrings(tk.MustQuery(tt).Rows())
 		})
 		tk.MustQuery(tt).Check(testkit.Rows(output[i].Plan...))
 	}
diff --git a/expression/main_test.go b/expression/main_test.go
index 84a4ad15392ad..ecbd5f70f93e7 100644
--- a/expression/main_test.go
+++ b/expression/main_test.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/testkit/testdata"
 	"github.com/pingcap/tidb/testkit/testmain"
 	"github.com/pingcap/tidb/util/mock"
 	"github.com/pingcap/tidb/util/testbridge"
@@ -28,6 +29,8 @@ import (
 	"go.uber.org/goleak"
 )
 
+var testDataMap = make(testdata.BookKeeper)
+
 func TestMain(m *testing.M) {
 	testbridge.WorkaroundGoCheckFlags()
 	testmain.ShortCircuitForBench(m)
@@ -45,13 +48,19 @@ func TestMain(m *testing.M) {
 	// Note, SetSystemTZ() is a sync.Once operation.
 	timeutil.SetSystemTZ("system")
 
+	testDataMap.LoadTestSuiteData("testdata", "flag_simplify")
+
 	opts := []goleak.Option{
 		goleak.IgnoreTopFunction("go.etcd.io/etcd/pkg/logutil.(*MergeLogger).outputLoop"),
 		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
 		goleak.IgnoreTopFunction("github.com/pingcap/tidb/table/tables.mockRemoteService"),
 	}
 
-	goleak.VerifyTestMain(m, opts...)
+	callback := func(i int) int {
+		testDataMap.GenerateOutputIfNeeded()
+		return i
+	}
+	goleak.VerifyTestMain(testmain.WrapTestingM(m, callback), opts...)
 }
 
 func createContext(t *testing.T) *mock.Context {
@@ -63,3 +72,7 @@ func createContext(t *testing.T) *mock.Context {
 	ctx.GetSessionVars().PlanColumnID = 0
 	return ctx
 }
+
+func GetFlagSimplifyData() testdata.TestData {
+	return testDataMap["flag_simplify"]
+}