talent-plan-tinykv/scheduler/server/testutil.go
Connor 5e089a2cd1 init course framework
Signed-off-by: Connor <zbk602423539@gmail.com>
Co-authored-by: Nick Cameron <nrc@ncameron.org>
Co-authored-by: linning <linningde25@gmail.com>
Co-authored-by: YangKeao <keao.yang@yahoo.com>
Co-authored-by: andylokandy <andylokandy@hotmail.com>
Co-authored-by: Iosmanthus Teng <myosmanthustree@gmail.com>
2020-04-30 15:25:07 +08:00

118 lines
3.2 KiB
Go

// Copyright 2017 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 server
import (
"context"
"fmt"
"io/ioutil"
"strings"
"sync"
"time"
"github.com/pingcap-incubator/tinykv/scheduler/pkg/tempurl"
"github.com/pingcap-incubator/tinykv/scheduler/pkg/testutil"
"github.com/pingcap-incubator/tinykv/scheduler/pkg/typeutil"
"github.com/pingcap-incubator/tinykv/scheduler/server/config"
"github.com/pingcap/check"
"github.com/pingcap/log"
"go.etcd.io/etcd/embed"
// Register schedulers
_ "github.com/pingcap-incubator/tinykv/scheduler/server/schedulers"
)
// CleanupFunc closes test pd server(s) and deletes any files left behind.
type CleanupFunc func()
// NewTestServer creates a pd server for testing.
func NewTestServer(c *check.C) (*Server, CleanupFunc, error) {
ctx, cancel := context.WithCancel(context.Background())
cfg := NewTestSingleConfig(c)
s, err := CreateServer(cfg)
if err != nil {
cancel()
return nil, nil, err
}
if err = s.Run(ctx); err != nil {
cancel()
return nil, nil, err
}
cleanup := func() {
cancel()
s.Close()
testutil.CleanServer(cfg)
}
return s, cleanup, nil
}
var zapLogOnce sync.Once
// NewTestSingleConfig is only for test to create one pd.
// Because PD client also needs this, so export here.
func NewTestSingleConfig(c *check.C) *config.Config {
cfg := &config.Config{
Name: "pd",
ClientUrls: tempurl.Alloc(),
PeerUrls: tempurl.Alloc(),
InitialClusterState: embed.ClusterStateFlagNew,
LeaderLease: 1,
TsoSaveInterval: typeutil.NewDuration(200 * time.Millisecond),
}
cfg.AdvertiseClientUrls = cfg.ClientUrls
cfg.AdvertisePeerUrls = cfg.PeerUrls
cfg.DataDir, _ = ioutil.TempDir("/tmp", "test_pd")
cfg.InitialCluster = fmt.Sprintf("pd=%s", cfg.PeerUrls)
cfg.DisableStrictReconfigCheck = true
cfg.TickInterval = typeutil.NewDuration(100 * time.Millisecond)
cfg.ElectionInterval = typeutil.NewDuration(3 * time.Second)
cfg.LeaderPriorityCheckInterval = typeutil.NewDuration(100 * time.Millisecond)
err := cfg.SetupLogger()
c.Assert(err, check.IsNil)
zapLogOnce.Do(func() {
log.ReplaceGlobals(cfg.GetZapLogger(), cfg.GetZapLogProperties())
})
c.Assert(cfg.Adjust(nil), check.IsNil)
return cfg
}
// NewTestMultiConfig is only for test to create multiple pd configurations.
// Because PD client also needs this, so export here.
func NewTestMultiConfig(c *check.C, count int) []*config.Config {
cfgs := make([]*config.Config, count)
var clusters []string
for i := 1; i <= count; i++ {
cfg := NewTestSingleConfig(c)
cfg.Name = fmt.Sprintf("pd%d", i)
clusters = append(clusters, fmt.Sprintf("%s=%s", cfg.Name, cfg.PeerUrls))
cfgs[i-1] = cfg
}
initialCluster := strings.Join(clusters, ",")
for _, cfg := range cfgs {
cfg.InitialCluster = initialCluster
}
return cfgs
}