Skip to content

Commit

Permalink
fix ctx
Browse files Browse the repository at this point in the history 
Signed-off-by: yisaer <disxiaofei@163.com>
  • Loading branch information
@Yisaer
Yisaer committed May 8, 2021
1 parent 04de5a9 commit 3f94c65
Show file tree
Hide file tree
Showing 20 changed files with 305 additions and 143 deletions.
22 changes: 17 additions & 5 deletions pkg/autoscaling/calculation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package autoscaling

import (
"context"
"encoding/json"
"fmt"
"math"
Expand All @@ -32,11 +33,22 @@ func Test(t *testing.T) {

var _ = Suite(&calculationTestSuite{})

type calculationTestSuite struct{}
type calculationTestSuite struct {
ctx context.Context
cancel context.CancelFunc
}

func (s *calculationTestSuite) SetUpSuite(c *C) {
s.ctx, s.cancel = context.WithCancel(context.Background())
}

func (s *calculationTestSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *calculationTestSuite) TestGetScaledTiKVGroups(c *C) {
// case1 indicates the tikv cluster with not any group existed
case1 := mockcluster.NewCluster(config.NewTestOptions())
case1 := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
case1.AddLabelsStore(1, 1, map[string]string{})
case1.AddLabelsStore(2, 1, map[string]string{
"foo": "bar",
Expand All @@ -46,7 +58,7 @@ func (s *calculationTestSuite) TestGetScaledTiKVGroups(c *C) {
})

// case2 indicates the tikv cluster with 1 auto-scaling group existed
case2 := mockcluster.NewCluster(config.NewTestOptions())
case2 := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
case2.AddLabelsStore(1, 1, map[string]string{})
case2.AddLabelsStore(2, 1, map[string]string{
groupLabelKey: fmt.Sprintf("%s-%s-0", autoScalingGroupLabelKeyPrefix, TiKV.String()),
Expand All @@ -58,7 +70,7 @@ func (s *calculationTestSuite) TestGetScaledTiKVGroups(c *C) {
})

// case3 indicates the tikv cluster with other group existed
case3 := mockcluster.NewCluster(config.NewTestOptions())
case3 := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
case3.AddLabelsStore(1, 1, map[string]string{})
case3.AddLabelsStore(2, 1, map[string]string{
groupLabelKey: "foo",
Expand Down Expand Up @@ -323,7 +335,7 @@ func (s *calculationTestSuite) TestStrategyChangeCount(c *C) {
}

// tikv cluster with 1 auto-scaling group existed
cluster := mockcluster.NewCluster(config.NewTestOptions())
cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
cluster.AddLabelsStore(1, 1, map[string]string{})
cluster.AddLabelsStore(2, 1, map[string]string{
groupLabelKey: fmt.Sprintf("%s-%s-0", autoScalingGroupLabelKeyPrefix, TiKV.String()),
Expand Down
3 changes: 1 addition & 2 deletions pkg/mock/mockcluster/mockcluster.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,7 @@ type Cluster struct {
}

// NewCluster creates a new Cluster
func NewCluster(opts *config.PersistOptions) *Cluster {
ctx := context.Background()
func NewCluster(ctx context.Context, opts *config.PersistOptions) *Cluster {
clus := &Cluster{
BasicCluster: core.NewBasicCluster(),
IDAllocator: mockid.NewIDAllocator(),
Expand Down
23 changes: 17 additions & 6 deletions server/replication/replication_mode_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,23 @@ func TestReplicationMode(t *testing.T) {

var _ = Suite(&testReplicationMode{})

type testReplicationMode struct{}
type testReplicationMode struct {
ctx context.Context
cancel context.CancelFunc
}

func (s *testReplicationMode) SetUpSuite(c *C) {
s.ctx, s.cancel = context.WithCancel(context.Background())
}

func (s *testReplicationMode) TearDownTest(c *C) {
s.cancel()
}

func (s *testReplicationMode) TestInitial(c *C) {
store := core.NewStorage(kv.NewMemoryKV())
conf := config.ReplicationModeConfig{ReplicationMode: modeMajority}
cluster := mockcluster.NewCluster(config.NewTestOptions())
cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
rep, err := NewReplicationModeManager(conf, store, cluster, nil)
c.Assert(err, IsNil)
c.Assert(rep.GetReplicationStatus(), DeepEquals, &pb.ReplicationStatus{Mode: pb.ReplicationMode_MAJORITY})
Expand Down Expand Up @@ -72,7 +83,7 @@ func (s *testReplicationMode) TestStatus(c *C) {
LabelKey: "dr-label",
WaitSyncTimeout: typeutil.Duration{Duration: time.Minute},
}}
cluster := mockcluster.NewCluster(config.NewTestOptions())
cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
rep, err := NewReplicationModeManager(conf, store, cluster, nil)
c.Assert(err, IsNil)
c.Assert(rep.GetReplicationStatus(), DeepEquals, &pb.ReplicationStatus{
Expand Down Expand Up @@ -147,7 +158,7 @@ func (s *testReplicationMode) TestStateSwitch(c *C) {
WaitStoreTimeout: typeutil.Duration{Duration: time.Minute},
WaitSyncTimeout: typeutil.Duration{Duration: time.Minute},
}}
cluster := mockcluster.NewCluster(config.NewTestOptions())
cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
var replicator mockFileReplicator
rep, err := NewReplicationModeManager(conf, store, cluster, &replicator)
c.Assert(err, IsNil)
Expand Down Expand Up @@ -251,7 +262,7 @@ func (s *testReplicationMode) TestAsynctimeout(c *C) {
WaitSyncTimeout: typeutil.Duration{Duration: time.Minute},
WaitAsyncTimeout: typeutil.Duration{Duration: 2 * time.Minute},
}}
cluster := mockcluster.NewCluster(config.NewTestOptions())
cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
var replicator mockFileReplicator
rep, err := NewReplicationModeManager(conf, store, cluster, &replicator)
c.Assert(err, IsNil)
Expand Down Expand Up @@ -302,7 +313,7 @@ func (s *testReplicationMode) TestRecoverProgress(c *C) {
WaitStoreTimeout: typeutil.Duration{Duration: time.Minute},
WaitSyncTimeout: typeutil.Duration{Duration: time.Minute},
}}
cluster := mockcluster.NewCluster(config.NewTestOptions())
cluster := mockcluster.NewCluster(s.ctx, config.NewTestOptions())
cluster.AddLabelsStore(1, 1, map[string]string{})
rep, err := NewReplicationModeManager(conf, store, cluster, nil)
c.Assert(err, IsNil)
Expand Down
14 changes: 13 additions & 1 deletion server/schedule/checker/joint_state_checker_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
package checker

import (
"context"

. "github.com/pingcap/check"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/tikv/pd/pkg/mock/mockcluster"
Expand All @@ -27,10 +29,20 @@ var _ = Suite(&testJointStateCheckerSuite{})
type testJointStateCheckerSuite struct {
cluster *mockcluster.Cluster
jsc *JointStateChecker
ctx context.Context
cancel context.CancelFunc
}

func (s *testJointStateCheckerSuite) SetUpSuite(c *C) {
s.ctx, s.cancel = context.WithCancel(context.Background())
}

func (s *testJointStateCheckerSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *testJointStateCheckerSuite) SetUpTest(c *C) {
s.cluster = mockcluster.NewCluster(config.NewTestOptions())
s.cluster = mockcluster.NewCluster(s.ctx, config.NewTestOptions())
s.jsc = NewJointStateChecker(s.cluster)
for id := uint64(1); id <= 10; id++ {
s.cluster.PutStoreWithLabels(id)
Expand Down
13 changes: 11 additions & 2 deletions server/schedule/checker/learner_checker_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,12 @@
package checker

import (
"context"
"github.com/tikv/pd/server/config"

. "github.com/pingcap/check"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/tikv/pd/pkg/mock/mockcluster"
"github.com/tikv/pd/server/config"
"github.com/tikv/pd/server/core"
"github.com/tikv/pd/server/schedule/operator"
"github.com/tikv/pd/server/versioninfo"
Expand All @@ -28,17 +30,24 @@ var _ = Suite(&testLearnerCheckerSuite{})
type testLearnerCheckerSuite struct {
cluster *mockcluster.Cluster
lc *LearnerChecker
ctx context.Context
cancel context.CancelFunc
}

func (s *testLearnerCheckerSuite) SetUpTest(c *C) {
s.cluster = mockcluster.NewCluster(config.NewTestOptions())
s.ctx, s.cancel = context.WithCancel(context.Background())
s.cluster = mockcluster.NewCluster(s.ctx, config.NewTestOptions())
s.cluster.DisableFeature(versioninfo.JointConsensus)
s.lc = NewLearnerChecker(s.cluster)
for id := uint64(1); id <= 10; id++ {
s.cluster.PutStoreWithLabels(id)
}
}

func (s *testLearnerCheckerSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *testLearnerCheckerSuite) TestPromoteLearner(c *C) {
lc := s.lc

Expand Down
17 changes: 10 additions & 7 deletions server/schedule/checker/merge_checker_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,17 @@ type testMergeCheckerSuite struct {
regions []*core.RegionInfo
}

func (s *testMergeCheckerSuite) SetUpSuite(c *C) {
s.ctx, s.cancel = context.WithCancel(context.Background())
}

func (s *testMergeCheckerSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *testMergeCheckerSuite) SetUpTest(c *C) {
cfg := config.NewTestOptions()
s.cluster = mockcluster.NewCluster(cfg)
s.cluster = mockcluster.NewCluster(s.ctx, cfg)
s.cluster.SetMaxMergeRegionSize(2)
s.cluster.SetMaxMergeRegionKeys(2)
s.cluster.SetLabelPropertyConfig(config.LabelPropertyConfig{
Expand Down Expand Up @@ -130,14 +138,9 @@ func (s *testMergeCheckerSuite) SetUpTest(c *C) {
for _, region := range s.regions {
s.cluster.PutRegion(region)
}
s.ctx, s.cancel = context.WithCancel(context.Background())
s.mc = NewMergeChecker(s.ctx, s.cluster)
}

func (s *testMergeCheckerSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *testMergeCheckerSuite) TestBasic(c *C) {
s.cluster.SetSplitMergeInterval(0)

Expand Down Expand Up @@ -434,7 +437,7 @@ type testSplitMergeSuite struct{}

func (s *testMergeCheckerSuite) TestCache(c *C) {
cfg := config.NewTestOptions()
s.cluster = mockcluster.NewCluster(cfg)
s.cluster = mockcluster.NewCluster(s.ctx, cfg)
s.cluster.SetMaxMergeRegionSize(2)
s.cluster.SetMaxMergeRegionKeys(2)
s.cluster.SetSplitMergeInterval(time.Hour)
Expand Down
27 changes: 19 additions & 8 deletions server/schedule/checker/replica_checker_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package checker

import (
"context"
"time"

. "github.com/pingcap/check"
Expand All @@ -39,11 +40,21 @@ var _ = Suite(&testReplicaCheckerSuite{})
type testReplicaCheckerSuite struct {
cluster *mockcluster.Cluster
rc *ReplicaChecker
ctx context.Context
cancel context.CancelFunc
}

func (s *testReplicaCheckerSuite) SetUpSuite(c *C) {
s.ctx, s.cancel = context.WithCancel(context.Background())
}

func (s *testReplicaCheckerSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *testReplicaCheckerSuite) SetUpTest(c *C) {
cfg := config.NewTestOptions()
s.cluster = mockcluster.NewCluster(cfg)
s.cluster = mockcluster.NewCluster(s.ctx, cfg)
s.cluster.DisableFeature(versioninfo.JointConsensus)
s.rc = NewReplicaChecker(s.cluster, cache.NewDefaultCache(10))
stats := &pdpb.StoreStats{
Expand Down Expand Up @@ -196,7 +207,7 @@ func (s *testReplicaCheckerSuite) downPeerAndCheck(c *C, aliveRole metapb.PeerRo

func (s *testReplicaCheckerSuite) TestBasic(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)
tc.SetMaxSnapshotCount(2)
tc.DisableFeature(versioninfo.JointConsensus)
rc := NewReplicaChecker(tc, cache.NewDefaultCache(10))
Expand Down Expand Up @@ -269,7 +280,7 @@ func (s *testReplicaCheckerSuite) TestBasic(c *C) {

func (s *testReplicaCheckerSuite) TestLostStore(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)

tc.DisableFeature(versioninfo.JointConsensus)
tc.AddRegionStore(1, 1)
Expand All @@ -288,7 +299,7 @@ func (s *testReplicaCheckerSuite) TestLostStore(c *C) {

func (s *testReplicaCheckerSuite) TestOffline(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)
tc.DisableFeature(versioninfo.JointConsensus)
tc.SetMaxReplicas(3)
tc.SetLocationLabels([]string{"zone", "rack", "host"})
Expand Down Expand Up @@ -340,7 +351,7 @@ func (s *testReplicaCheckerSuite) TestOffline(c *C) {

func (s *testReplicaCheckerSuite) TestDistinctScore(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)
tc.DisableFeature(versioninfo.JointConsensus)
tc.SetMaxReplicas(3)
tc.SetLocationLabels([]string{"zone", "rack", "host"})
Expand Down Expand Up @@ -419,7 +430,7 @@ func (s *testReplicaCheckerSuite) TestDistinctScore(c *C) {

func (s *testReplicaCheckerSuite) TestDistinctScore2(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)
tc.DisableFeature(versioninfo.JointConsensus)
tc.SetMaxReplicas(5)
tc.SetLocationLabels([]string{"zone", "host"})
Expand Down Expand Up @@ -449,7 +460,7 @@ func (s *testReplicaCheckerSuite) TestDistinctScore2(c *C) {

func (s *testReplicaCheckerSuite) TestStorageThreshold(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)
tc.SetLocationLabels([]string{"zone"})
tc.DisableFeature(versioninfo.JointConsensus)
rc := NewReplicaChecker(tc, cache.NewDefaultCache(10))
Expand Down Expand Up @@ -485,7 +496,7 @@ func (s *testReplicaCheckerSuite) TestStorageThreshold(c *C) {

func (s *testReplicaCheckerSuite) TestOpts(c *C) {
opt := config.NewTestOptions()
tc := mockcluster.NewCluster(opt)
tc := mockcluster.NewCluster(s.ctx, opt)
tc.DisableFeature(versioninfo.JointConsensus)
rc := NewReplicaChecker(tc, cache.NewDefaultCache(10))

Expand Down
13 changes: 12 additions & 1 deletion server/schedule/checker/rule_checker_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package checker

import (
"context"
"encoding/hex"

. "github.com/pingcap/check"
Expand All @@ -34,11 +35,21 @@ type testRuleCheckerSuite struct {
cluster *mockcluster.Cluster
ruleManager *placement.RuleManager
rc *RuleChecker
ctx context.Context
cancel context.CancelFunc
}

func (s *testRuleCheckerSuite) SetUpSuite(c *C) {
s.ctx, s.cancel = context.WithCancel(context.Background())
}

func (s *testRuleCheckerSuite) TearDownTest(c *C) {
s.cancel()
}

func (s *testRuleCheckerSuite) SetUpTest(c *C) {
cfg := config.NewTestOptions()
s.cluster = mockcluster.NewCluster(cfg)
s.cluster = mockcluster.NewCluster(s.ctx, cfg)
s.cluster.DisableFeature(versioninfo.JointConsensus)
s.cluster.SetEnablePlacementRules(true)
s.ruleManager = s.cluster.RuleManager
Expand Down
Loading

0 comments on commit 3f94c65

Please sign in to comment.