util/profile: migrate test-infra to testify (pingcap#26205)

testkit/testkit.go

@@ -0,0 +1,117 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// +build !codes
+
+package testkit
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/session"
+	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/sqlexec"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
+)
+
+var idGenerator atomic.Uint64
+
+// TestKit is a utility to run sql test.
+type TestKit struct {
+	require *require.Assertions
+	assert  *assert.Assertions
+	store   kv.Storage
+	session session.Session
+}
+
+// NewTestKit returns a new *TestKit.
+func NewTestKit(t *testing.T, store kv.Storage) *TestKit {
+	return &TestKit{
+		require: require.New(t),
+		assert:  assert.New(t),
+		store:   store,
+		session: newSession(t, store),
+	}
+}
+
+// MustExec executes a sql statement and asserts nil error.
+func (tk *TestKit) MustExec(sql string, args ...interface{}) {
+	res, err := tk.Exec(sql, args...)
+	comment := fmt.Sprintf("sql:%s, %v, error stack %v", sql, args, errors.ErrorStack(err))
+	tk.require.Nil(err, comment)
+
+	if res != nil {
+		tk.require.Nil(res.Close())
+	}
+}
+
+// Exec executes a sql statement using the prepared stmt API
+func (tk *TestKit) Exec(sql string, args ...interface{}) (sqlexec.RecordSet, error) {
+	ctx := context.Background()
+	if len(args) == 0 {
+		sc := tk.session.GetSessionVars().StmtCtx
+		prevWarns := sc.GetWarnings()
+		stmts, err := tk.session.Parse(ctx, sql)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		warns := sc.GetWarnings()
+		parserWarns := warns[len(prevWarns):]
+		var rs0 sqlexec.RecordSet
+		for i, stmt := range stmts {
+			rs, err := tk.session.ExecuteStmt(ctx, stmt)
+			if i == 0 {
+				rs0 = rs
+			}
+			if err != nil {
+				tk.session.GetSessionVars().StmtCtx.AppendError(err)
+				return nil, errors.Trace(err)
+			}
+		}
+		if len(parserWarns) > 0 {
+			tk.session.GetSessionVars().StmtCtx.AppendWarnings(parserWarns)
+		}
+		return rs0, nil
+	}
+
+	stmtID, _, _, err := tk.session.PrepareStmt(sql)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	params := make([]types.Datum, len(args))
+	for i := 0; i < len(params); i++ {
+		params[i] = types.NewDatum(args[i])
+	}
+	rs, err := tk.session.ExecutePreparedStmt(ctx, stmtID, params)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	err = tk.session.DropPreparedStmt(stmtID)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	return rs, nil
+}
+
+func newSession(t *testing.T, store kv.Storage) session.Session {
+	se, err := session.CreateSession4Test(store)
+	require.Nil(t, err)
+	se.SetConnectionID(idGenerator.Inc())
+	return se
+}

util/profile/flamegraph_test.go

@@ -14,28 +14,27 @@
 package profile
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
-	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-type profileInternalSuite struct{}
+func TestProfileToDatum(t *testing.T) {
+	t.Parallel()
 
-var _ = Suite(&profileInternalSuite{})
-
-func TestT(t *testing.T) {
-	TestingT(t)
-}
-
-func (s *profileInternalSuite) TestProfileToDatum(c *C) {
 	file, err := os.Open("testdata/test.pprof")
-	c.Assert(err, IsNil)
-	defer file.Close()
+	require.Nil(t, err)
+	defer func() {
+		err := file.Close()
+		require.Nil(t, err)
+	}()
 
 	data, err := (&Collector{}).ProfileReaderToDatums(file)
-	c.Assert(err, IsNil)
+	require.Nil(t, err)
 
 	datums := [][]types.Datum{
 		types.MakeDatums(`root`, "100%", "100%", 0, 0, `root`),
@@ -82,13 +81,13 @@ func (s *profileInternalSuite) TestProfileToDatum(c *C) {
 	}
 
 	for i, row := range data {
-		comment := Commentf("row %2d", i)
+		comment := fmt.Sprintf("row %2d", i)
 		rowStr, err := types.DatumsToString(row, true)
-		c.Assert(err, IsNil, comment)
+		assert.Nil(t, err, comment)
 		expectStr, err := types.DatumsToString(datums[i], true)
-		c.Assert(err, IsNil, comment)
+		assert.Nil(t, err, comment)
 
-		comment = Commentf("row %2d, actual (%s), expected (%s)", i, rowStr, expectStr)
+		comment = fmt.Sprintf("row %2d, actual (%s), expected (%s)", i, rowStr, expectStr)
 		equal := true
 		for j, r := range row {
 			v, err := r.CompareDatum(nil, &datums[i][j])
@@ -97,7 +96,7 @@ func (s *profileInternalSuite) TestProfileToDatum(c *C) {
 				break
 			}
 		}
-		c.Assert(err, IsNil, comment)
-		c.Assert(equal, IsTrue, comment)
+		assert.Nil(t, err, comment)
+		assert.True(t, equal, comment)
 	}
 }

util/profile/main_test.go

@@ -0,0 +1,30 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package profile
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/util/testbridge"
+	"go.uber.org/goleak"
+)
+
+func TestMain(m *testing.M) {
+	testbridge.WorkaroundGoCheckFlags()
+	opts := []goleak.Option{
+		goleak.IgnoreTopFunction("go.etcd.io/etcd/pkg/logutil.(*MergeLogger).outputLoop"),
+		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
+	}
+	goleak.VerifyTestMain(m, opts...)
+}

util/profile/profile_test.go

@@ -14,45 +14,43 @@
 package profile_test
 
 import (
+	"testing"
 	"time"
 
-	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/store/mockstore"
+	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/util/profile"
-	"github.com/pingcap/tidb/util/testkit"
+	"github.com/stretchr/testify/require"
 )
 
-type profileSuite struct {
-	store kv.Storage
-	dom   *domain.Domain
-}
+func TestProfiles(t *testing.T) {
+	t.Parallel()
+	var err error
+	var store kv.Storage
+	var dom *domain.Domain
 
-var _ = Suite(&profileSuite{})
+	store, err = mockstore.NewMockStore()
+	require.Nil(t, err)
+	defer func() {
+		err := store.Close()
+		require.Nil(t, err)
+	}()
 
-func (s *profileSuite) SetUpSuite(c *C) {
-	var err error
-	s.store, err = mockstore.NewMockStore()
-	c.Assert(err, IsNil)
 	session.DisableStats4Test()
-	s.dom, err = session.BootstrapSession(s.store)
-	c.Assert(err, IsNil)
-}
-
-func (s *profileSuite) TearDownSuite(c *C) {
-	s.dom.Close()
-	s.store.Close()
-}
+	dom, err = session.BootstrapSession(store)
+	require.Nil(t, err)
+	defer dom.Close()
 
-func (s *profileSuite) TestProfiles(c *C) {
 	oldValue := profile.CPUProfileInterval
 	profile.CPUProfileInterval = 2 * time.Second
 	defer func() {
 		profile.CPUProfileInterval = oldValue
 	}()
-	tk := testkit.NewTestKit(c, s.store)
+
+	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("select * from performance_schema.tidb_profile_cpu")
 	tk.MustExec("select * from performance_schema.tidb_profile_memory")
 	tk.MustExec("select * from performance_schema.tidb_profile_allocs")

util/profile/trackerrecorder_test.go

@@ -15,16 +15,44 @@ package profile
 
 import (
 	"math/rand"
+	"testing"
 	"time"
 	"unsafe"
 
-	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb/util/kvcache"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Suite(&trackRecorderSuite{})
+func TestHeapProfileRecorder(t *testing.T) {
+	t.Parallel()
 
-type trackRecorderSuite struct {
+	// As runtime.MemProfileRate default values is 512 KB , so the num should be greater than 60000
+	// that the memory usage of the values would be greater than 512 KB.
+	num := 60000
+	lru := kvcache.NewSimpleLRUCache(uint(num), 0, 0)
+
+	keys := make([]*mockCacheKey, num)
+	for i := 0; i < num; i++ {
+		keys[i] = newMockHashKey(int64(i))
+		v := getRandomString(10)
+		lru.Put(keys[i], v)
+	}
+	for _, k := range lru.Keys() {
+		assert.Len(t, k.Hash(), 8)
+	}
+	for _, v := range lru.Values() {
+		assert.Len(t, v.(string), 10)
+	}
+
+	bytes, err := col.getFuncMemUsage(kvcache.ProfileName)
+	require.Nil(t, err)
+
+	valueSize := int(unsafe.Sizeof(getRandomString(10)))
+	// ensure that the consumed bytes is at least larger than num * size of value
+	assert.LessOrEqual(t, int64(valueSize*num), bytes)
+	// we should assert lru size last and value size to reference lru in order to avoid gc
+	assert.Equal(t, num, lru.Size())
 }
 
 type mockCacheKey struct {
@@ -59,32 +87,3 @@ func getRandomString(l int) string {
 	}
 	return string(result)
 }
-
-func (t *trackRecorderSuite) TestHeapProfileRecorder(c *C) {
-	// As runtime.MemProfileRate default values is 512 KB , so the num should be greater than 60000
-	// that the memory usage of the values would be greater than 512 KB.
-	num := 60000
-	lru := kvcache.NewSimpleLRUCache(uint(num), 0, 0)
-
-	keys := make([]*mockCacheKey, num)
-	for i := 0; i < num; i++ {
-		keys[i] = newMockHashKey(int64(i))
-		v := getRandomString(10)
-		lru.Put(keys[i], v)
-	}
-	for _, k := range lru.Keys() {
-		c.Assert(len(k.Hash()), Equals, 8)
-	}
-	for _, v := range lru.Values() {
-		val := v.(string)
-		c.Assert(len(val), Equals, 10)
-	}
-
-	bytes, err := col.getFuncMemUsage(kvcache.ProfileName)
-	c.Assert(err, IsNil)
-	valueSize := int(unsafe.Sizeof(getRandomString(10)))
-	// ensure that the consumed bytes is at least larger than num * size of value
-	c.Assert(int64(valueSize*num), LessEqual, bytes)
-	// we should assert lru size last and value size to reference lru in order to avoid gc
-	c.Assert(lru.Size(), Equals, num)
-}