mockhbstream, movingaverage, netutil, progress: testify the tests (#5086

) ref #4813 Testify the pkg/mockhbstream, pkg/movingaverage, pkg/netutil, pkg/progress tests. Signed-off-by: JmPotato <[email protected]>
tikv · Jun 1, 2022 · 422acfb · 422acfb
1 parent ce6b35d
commit 422acfb
Show file tree

Hide file tree

Showing 7 changed files with 128 additions and 160 deletions.
diff --git a/pkg/mock/mockhbstream/mockhbstream_test.go b/pkg/mock/mockhbstream/mockhbstream_test.go
@@ -19,40 +19,22 @@ import (
 	"testing"
 
 	"github.com/gogo/protobuf/proto"
-	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/eraftpb"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
+	"github.com/stretchr/testify/require"
 	"github.com/tikv/pd/pkg/mock/mockcluster"
 	"github.com/tikv/pd/pkg/testutil"
 	"github.com/tikv/pd/server/config"
 	"github.com/tikv/pd/server/schedule/hbstream"
 )
 
-func TestHeaertbeatStreams(t *testing.T) {
-	TestingT(t)
-}
-
-var _ = Suite(&testHeartbeatStreamSuite{})
-
-type testHeartbeatStreamSuite struct {
-	ctx    context.Context
-	cancel context.CancelFunc
-}
-
-func (s *testHeartbeatStreamSuite) SetUpSuite(c *C) {
-	s.ctx, s.cancel = context.WithCancel(context.Background())
-}
-
-func (s *testHeartbeatStreamSuite) TearDownTest(c *C) {
-	s.cancel()
-}
-
-func (s *testHeartbeatStreamSuite) TestActivity(c *C) {
+func TestActivity(t *testing.T) {
+	re := require.New(t)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
+	cluster := mockcluster.NewCluster(ctx, config.NewTestOptions())
 	cluster.AddRegionStore(1, 1)
 	cluster.AddRegionStore(2, 0)
 	cluster.AddLeaderRegion(1, 1)
@@ -66,24 +48,24 @@ func (s *testHeartbeatStreamSuite) TestActivity(c *C) {
 
 	// Active stream is stream1.
 	hbs.BindStream(1, stream1)
-	testutil.WaitUntil(c, func() bool {
+	testutil.WaitUntilWithTestingT(t, func() bool {
 		hbs.SendMsg(region, proto.Clone(msg).(*pdpb.RegionHeartbeatResponse))
 		return stream1.Recv() != nil && stream2.Recv() == nil
 	})
 	// Rebind to stream2.
 	hbs.BindStream(1, stream2)
-	testutil.WaitUntil(c, func() bool {
+	testutil.WaitUntilWithTestingT(t, func() bool {
 		hbs.SendMsg(region, proto.Clone(msg).(*pdpb.RegionHeartbeatResponse))
 		return stream1.Recv() == nil && stream2.Recv() != nil
 	})
 	// SendErr to stream2.
 	hbs.SendErr(pdpb.ErrorType_UNKNOWN, "test error", &metapb.Peer{Id: 1, StoreId: 1})
 	res := stream2.Recv()
-	c.Assert(res, NotNil)
-	c.Assert(res.GetHeader().GetError(), NotNil)
+	re.NotNil(res)
+	re.NotNil(res.GetHeader().GetError())
 	// Switch back to 1 again.
 	hbs.BindStream(1, stream1)
-	testutil.WaitUntil(c, func() bool {
+	testutil.WaitUntilWithTestingT(t, func() bool {
 		hbs.SendMsg(region, proto.Clone(msg).(*pdpb.RegionHeartbeatResponse))
 		return stream1.Recv() != nil && stream2.Recv() == nil
 	})

diff --git a/pkg/movingaverage/avg_over_time_test.go b/pkg/movingaverage/avg_over_time_test.go
@@ -16,16 +16,14 @@ package movingaverage
 
 import (
 	"math/rand"
+	"testing"
 	"time"
 
-	. "github.com/pingcap/check"
+	"github.com/stretchr/testify/require"
 )
 
-var _ = Suite(&testAvgOverTimeSuite{})
-
-type testAvgOverTimeSuite struct{}
-
-func (t *testAvgOverTimeSuite) TestPulse(c *C) {
+func TestPulse(t *testing.T) {
+	re := require.New(t)
 	aot := NewAvgOverTime(5 * time.Second)
 	// warm up
 	for i := 0; i < 5; i++ {
@@ -38,27 +36,28 @@ func (t *testAvgOverTimeSuite) TestPulse(c *C) {
 		} else {
 			aot.Add(0, time.Second)
 		}
-		c.Assert(aot.Get(), LessEqual, 600.)
-		c.Assert(aot.Get(), GreaterEqual, 400.)
+		re.LessOrEqual(aot.Get(), 600.)
+		re.GreaterOrEqual(aot.Get(), 400.)
 	}
 }
 
-func (t *testAvgOverTimeSuite) TestChange(c *C) {
+func TestChange(t *testing.T) {
+	re := require.New(t)
 	aot := NewAvgOverTime(5 * time.Second)
 
 	// phase 1: 1000
 	for i := 0; i < 20; i++ {
 		aot.Add(1000, time.Second)
 	}
-	c.Assert(aot.Get(), LessEqual, 1010.)
-	c.Assert(aot.Get(), GreaterEqual, 990.)
+	re.LessOrEqual(aot.Get(), 1010.)
+	re.GreaterOrEqual(aot.Get(), 990.)
 
 	// phase 2: 500
 	for i := 0; i < 5; i++ {
 		aot.Add(500, time.Second)
 	}
-	c.Assert(aot.Get(), LessEqual, 900.)
-	c.Assert(aot.Get(), GreaterEqual, 495.)
+	re.LessOrEqual(aot.Get(), 900.)
+	re.GreaterOrEqual(aot.Get(), 495.)
 	for i := 0; i < 15; i++ {
 		aot.Add(500, time.Second)
 	}
@@ -67,32 +66,34 @@ func (t *testAvgOverTimeSuite) TestChange(c *C) {
 	for i := 0; i < 5; i++ {
 		aot.Add(100, time.Second)
 	}
-	c.Assert(aot.Get(), LessEqual, 678.)
-	c.Assert(aot.Get(), GreaterEqual, 99.)
+	re.LessOrEqual(aot.Get(), 678.)
+	re.GreaterOrEqual(aot.Get(), 99.)
 
 	// clear
 	aot.Set(10)
-	c.Assert(aot.Get(), Equals, 10.)
+	re.Equal(10., aot.Get())
 }
 
-func (t *testAvgOverTimeSuite) TestMinFilled(c *C) {
+func TestMinFilled(t *testing.T) {
+	re := require.New(t)
 	interval := 10 * time.Second
 	rate := 1.0
 	for aotSize := 2; aotSize < 10; aotSize++ {
 		for mfSize := 2; mfSize < 10; mfSize++ {
 			tm := NewTimeMedian(aotSize, mfSize, interval)
 			for i := 0; i < tm.GetFilledPeriod(); i++ {
-				c.Assert(tm.Get(), Equals, 0.0)
+				re.Equal(0.0, tm.Get())
 				tm.Add(rate*interval.Seconds(), interval)
 			}
-			c.Assert(tm.Get(), Equals, rate)
+			re.Equal(rate, tm.Get())
 		}
 	}
 }
 
-func (t *testAvgOverTimeSuite) TestUnstableInterval(c *C) {
+func TestUnstableInterval(t *testing.T) {
+	re := require.New(t)
 	aot := NewAvgOverTime(5 * time.Second)
-	c.Assert(aot.Get(), Equals, 0.)
+	re.Equal(0., aot.Get())
 	// warm up
 	for i := 0; i < 5; i++ {
 		aot.Add(1000, time.Second)
@@ -101,8 +102,8 @@ func (t *testAvgOverTimeSuite) TestUnstableInterval(c *C) {
 	for i := 0; i < 1000; i++ {
 		r := float64(rand.Intn(5))
 		aot.Add(1000*r, time.Second*time.Duration(r))
-		c.Assert(aot.Get(), LessEqual, 1010.)
-		c.Assert(aot.Get(), GreaterEqual, 990.)
+		re.LessOrEqual(aot.Get(), 1010.)
+		re.GreaterOrEqual(aot.Get(), 990.)
 	}
 	// warm up
 	for i := 0; i < 5; i++ {
@@ -112,7 +113,7 @@ func (t *testAvgOverTimeSuite) TestUnstableInterval(c *C) {
 	for i := 0; i < 1000; i++ {
 		rate := float64(i%5*100) + 500
 		aot.Add(rate*3, time.Second*3)
-		c.Assert(aot.Get(), LessEqual, 910.)
-		c.Assert(aot.Get(), GreaterEqual, 490.)
+		re.LessOrEqual(aot.Get(), 910.)
+		re.GreaterOrEqual(aot.Get(), 490.)
 	}
 }
diff --git a/pkg/movingaverage/max_filter_test.go b/pkg/movingaverage/max_filter_test.go
@@ -15,22 +15,21 @@
 package movingaverage
 
 import (
-	. "github.com/pingcap/check"
-)
-
-var _ = Suite(&testMaxFilter{})
+	"testing"
 
-type testMaxFilter struct{}
+	"github.com/stretchr/testify/require"
+)
 
-func (t *testMaxFilter) TestMaxFilter(c *C) {
+func TestMaxFilter(t *testing.T) {
+	re := require.New(t)
 	var empty float64 = 0
 	data := []float64{2, 1, 3, 4, 1, 1, 3, 3, 2, 0, 5}
 	expected := []float64{2, 2, 3, 4, 4, 4, 4, 4, 3, 3, 5}
 
 	mf := NewMaxFilter(5)
-	c.Assert(mf.Get(), Equals, empty)
+	re.Equal(empty, mf.Get())
 
-	checkReset(c, mf, empty)
-	checkAdd(c, mf, data, expected)
-	checkSet(c, mf, data, expected)
+	checkReset(re, mf, empty)
+	checkAdd(re, mf, data, expected)
+	checkSet(re, mf, data, expected)
 }
diff --git a/pkg/movingaverage/moving_average_test.go b/pkg/movingaverage/moving_average_test.go
@@ -20,17 +20,9 @@ import (
 	"testing"
 	"time"
 
-	. "github.com/pingcap/check"
+	"github.com/stretchr/testify/require"
 )
 
-func Test(t *testing.T) {
-	TestingT(t)
-}
-
-var _ = Suite(&testMovingAvg{})
-
-type testMovingAvg struct{}
-
 func addRandData(ma MovingAvg, n int, mx float64) {
 	rand.Seed(time.Now().UnixNano())
 	for i := 0; i < n; i++ {
@@ -40,63 +32,65 @@ func addRandData(ma MovingAvg, n int, mx float64) {
 
 // checkReset checks the Reset works properly.
 // emptyValue is the moving average of empty data set.
-func checkReset(c *C, ma MovingAvg, emptyValue float64) {
+func checkReset(re *require.Assertions, ma MovingAvg, emptyValue float64) {
 	addRandData(ma, 100, 1000)
 	ma.Reset()
-	c.Assert(ma.Get(), Equals, emptyValue)
+	re.Equal(emptyValue, ma.Get())
 }
 
 // checkAddGet checks Add works properly.
-func checkAdd(c *C, ma MovingAvg, data []float64, expected []float64) {
-	c.Assert(len(data), Equals, len(expected))
+func checkAdd(re *require.Assertions, ma MovingAvg, data []float64, expected []float64) {
+	re.Equal(len(expected), len(data))
 	for i, x := range data {
 		ma.Add(x)
-		c.Assert(math.Abs(ma.Get()-expected[i]), LessEqual, 1e-7)
+		re.LessOrEqual(math.Abs(ma.Get()-expected[i]), 1e-7)
 	}
 }
 
 // checkSet checks Set = Reset + Add
-func checkSet(c *C, ma MovingAvg, data []float64, expected []float64) {
-	c.Assert(len(data), Equals, len(expected))
+func checkSet(re *require.Assertions, ma MovingAvg, data []float64, expected []float64) {
+	re.Equal(len(expected), len(data))
 
 	// Reset + Add
 	addRandData(ma, 100, 1000)
 	ma.Reset()
-	checkAdd(c, ma, data, expected)
+	checkAdd(re, ma, data, expected)
 
 	// Set
 	addRandData(ma, 100, 1000)
 	ma.Set(data[0])
-	c.Assert(ma.Get(), Equals, expected[0])
-	checkAdd(c, ma, data[1:], expected[1:])
+	re.Equal(expected[0], ma.Get())
+	checkAdd(re, ma, data[1:], expected[1:])
 }
 
 // checkInstantaneous checks GetInstantaneous
-func checkInstantaneous(c *C, ma MovingAvg) {
+func checkInstantaneous(re *require.Assertions, ma MovingAvg) {
 	value := 100.000000
 	ma.Add(value)
-	c.Assert(ma.GetInstantaneous(), Equals, value)
+	re.Equal(value, ma.GetInstantaneous())
 }
 
-func (t *testMovingAvg) TestMedianFilter(c *C) {
+func TestMedianFilter(t *testing.T) {
+	re := require.New(t)
 	var empty float64 = 0
 	data := []float64{2, 4, 2, 800, 600, 6, 3}
 	expected := []float64{2, 3, 2, 3, 4, 6, 6}
 
 	mf := NewMedianFilter(5)
-	c.Assert(mf.Get(), Equals, empty)
+	re.Equal(empty, mf.Get())
 
-	checkReset(c, mf, empty)
-	checkAdd(c, mf, data, expected)
-	checkSet(c, mf, data, expected)
+	checkReset(re, mf, empty)
+	checkAdd(re, mf, data, expected)
+	checkSet(re, mf, data, expected)
 }
 
 type testCase struct {
 	ma       MovingAvg
 	expected []float64
 }
 
-func (t *testMovingAvg) TestMovingAvg(c *C) {
+func TestMovingAvg(t *testing.T) {
+	re := require.New(t)
 	var empty float64 = 0
 	data := []float64{1, 1, 1, 1, 5, 1, 1, 1}
 	testCases := []testCase{{
@@ -116,11 +110,11 @@ func (t *testMovingAvg) TestMovingAvg(c *C) {
 		expected: []float64{1.000000, 1.000000, 1.000000, 1.000000, 5.000000, 5.000000, 5.000000, 5.000000},
 	},
 	}
-	for _, test := range testCases {
-		c.Assert(test.ma.Get(), Equals, empty)
-		checkReset(c, test.ma, empty)
-		checkAdd(c, test.ma, data, test.expected)
-		checkSet(c, test.ma, data, test.expected)
-		checkInstantaneous(c, test.ma)
+	for _, testCase := range testCases {
+		re.Equal(empty, testCase.ma.Get())
+		checkReset(re, testCase.ma, empty)
+		checkAdd(re, testCase.ma, data, testCase.expected)
+		checkSet(re, testCase.ma, data, testCase.expected)
+		checkInstantaneous(re, testCase.ma)
 	}
 }
diff --git a/pkg/movingaverage/queue_test.go b/pkg/movingaverage/queue_test.go
@@ -15,26 +15,30 @@
 package movingaverage
 
 import (
-	. "github.com/pingcap/check"
+	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
-func (t *testMovingAvg) TestQueue(c *C) {
+func TestQueue(t *testing.T) {
+	re := require.New(t)
 	sq := NewSafeQueue()
 	sq.PushBack(1)
 	sq.PushBack(2)
 	v1 := sq.PopFront()
 	v2 := sq.PopFront()
-	c.Assert(1, Equals, v1.(int))
-	c.Assert(2, Equals, v2.(int))
+	re.Equal(1, v1.(int))
+	re.Equal(2, v2.(int))
 }
 
-func (t *testMovingAvg) TestClone(c *C) {
+func TestClone(t *testing.T) {
+	re := require.New(t)
 	s1 := NewSafeQueue()
 	s1.PushBack(1)
 	s1.PushBack(2)
 	s2 := s1.Clone()
 	s2.PopFront()
 	s2.PopFront()
-	c.Assert(s1.que.Len(), Equals, 2)
-	c.Assert(s2.que.Len(), Equals, 0)
+	re.Equal(2, s1.que.Len())
+	re.Equal(0, s2.que.Len())
 }