address comment

tikv · Jun 5, 2018 · af05403 · af05403
1 parent 2a1932f
commit af05403
Show file tree

Hide file tree

Showing 8 changed files with 28 additions and 32 deletions.
diff --git a/pd-client/leader_change_test.go b/pd-client/leader_change_test.go
@@ -39,9 +39,8 @@ func (s *testLeaderChangeSuite) prepareClusterN(c *C, n int) (svrs map[string]*s
 	for i := 0; i < n; i++ {
 		cfg := cfgs[i]
 		go func() {
-			svr, err := server.CreateServer(cfg, api.NewHandler)
+			svr, err := server.CreateServer(cfg, api.NewHandler, server.TestMode())
 			c.Assert(err, IsNil)
-			svr.SetUpTestMode()
 			err = svr.Run()
 			c.Assert(err, IsNil)
 			ch <- svr

diff --git a/server/api/server_test.go b/server/api/server_test.go
@@ -78,9 +78,8 @@ func mustNewCluster(c *C, num int) ([]*server.Config, []*server.Server, cleanUpF
 	ch := make(chan *server.Server, num)
 	for _, cfg := range cfgs {
 		go func(cfg *server.Config) {
-			s, err := server.CreateServer(cfg, NewHandler)
+			s, err := server.CreateServer(cfg, NewHandler, server.TestMode())
 			c.Assert(err, IsNil)
-			s.SetUpTestMode()
 			err = s.Run()
 			c.Assert(err, IsNil)
 			ch <- s

diff --git a/server/api/store_ns_test.go b/server/api/store_ns_test.go
@@ -61,9 +61,8 @@ func (s *testStoreNsSuite) SetUpSuite(c *C) {
 
 	cfg := server.NewTestSingleConfig()
 	cfg.NamespaceClassifier = "table"
-	svr, err := server.CreateServer(cfg, NewHandler)
+	svr, err := server.CreateServer(cfg, NewHandler, server.TestMode())
 	c.Assert(err, IsNil)
-	svr.SetUpTestMode()
 	c.Assert(svr.Run(), IsNil)
 	s.svr = svr
 	s.cleanup = func() {

diff --git a/server/join_test.go b/server/join_test.go
@@ -100,12 +100,11 @@ func startPdWith(cfg *Config) (*Server, error) {
 			errCh <- errors.Trace(err)
 			return
 		}
-		svr, err := CreateServer(cfg, nil)
+		svr, err := CreateServer(cfg, nil, TestMode())
 		if err != nil {
 			errCh <- errors.Trace(err)
 			return
 		}
-		svr.SetUpTestMode()
 		err = svr.Run()
 		if err != nil {
 			errCh <- errors.Trace(err)

diff --git a/server/leader_test.go b/server/leader_test.go
@@ -35,9 +35,8 @@ func (s *testGetLeaderSuite) SetUpSuite(c *C) {
 
 	s.wg.Add(1)
 	s.done = make(chan bool)
-	svr, err := CreateServer(cfg, nil)
+	svr, err := CreateServer(cfg, nil, TestMode())
 	c.Assert(err, IsNil)
-	svr.SetUpTestMode()
 	err = svr.Run()
 	// Send requests after server has started.
 	go s.sendRequest(c, cfg.ClientUrls)

diff --git a/server/server.go b/server/server.go
@@ -91,8 +91,22 @@ type Server struct {
 	hbStreams *heartbeatStreams
 }
 
+// CreateServerOption configures server when creating.
+type CreateServerOption func(*Server)
+
+// TestMode for testing.
+func TestMode() CreateServerOption {
+	// The etcd master version has removed embed.Config.SetupLogging.
+	// Now logger is set up automatically based on embed.Config.Logger, embed.Config.LogOutputs, embed.Config.Debug fields.
+	// Use zap logger in the test, otherwise will panic. Reference: https://github.com/coreos/etcd/blob/master/embed/config_logging.go#L45
+	return func(s *Server) {
+		s.etcdCfg.Logger = "zap"
+		s.etcdCfg.LogOutputs = []string{"stdout"}
+	}
+}
+
 // CreateServer creates the UNINITIALIZED pd server with given configuration.
-func CreateServer(cfg *Config, apiRegister func(*Server) http.Handler) (*Server, error) {
+func CreateServer(cfg *Config, apiRegister func(*Server) http.Handler, ops ...CreateServerOption) (*Server, error) {
 	log.Infof("PD config - %v", cfg)
 	rand.Seed(time.Now().UnixNano())
 
@@ -114,7 +128,9 @@ func CreateServer(cfg *Config, apiRegister func(*Server) http.Handler) (*Server,
 	}
 	etcdCfg.ServiceRegister = func(gs *grpc.Server) { pdpb.RegisterPDServer(gs, s) }
 	s.etcdCfg = etcdCfg
-
+	for _, op := range ops {
+		op(s)
+	}
 	return s, nil
 }
 
@@ -285,16 +301,6 @@ func (s *Server) Run() error {
 	return nil
 }
 
-// SetUpTestMode for testing
-func (s *Server) SetUpTestMode() {
-	// The etcd master version has removed embed.Config.SetupLogging.
-	// Now logger is set up automatically based on embed.Config.Logger, embed.Config.LogOutputs, embed.Config.Debug fields.
-	// Use zap logger in the test, otherwise will panic. Reference: https://github.com/coreos/etcd/blob/master/embed/config_logging.go#L45
-	s.etcdCfg.Logger = "zap"
-	s.etcdCfg.LogOutputs = []string{"stdout"}
-	s.etcdCfg.Debug = false
-}
-
 func (s *Server) bootstrapCluster(req *pdpb.BootstrapRequest) (*pdpb.BootstrapResponse, error) {
 	clusterID := s.clusterID
 

diff --git a/server/server_test.go b/server/server_test.go
@@ -34,9 +34,8 @@ type cleanupFunc func()
 func newTestServer(c *C) (*Server, cleanUpFunc) {
 	cfg := NewTestSingleConfig()
 
-	svr, err := CreateServer(cfg, nil)
+	svr, err := CreateServer(cfg, nil, TestMode())
 	c.Assert(err, IsNil)
-	svr.SetUpTestMode()
 
 	cleanup := func() {
 		svr.Close()
@@ -93,9 +92,8 @@ func (s *testLeaderServerSuite) SetUpSuite(c *C) {
 		cfg := cfgs[i]
 
 		go func() {
-			svr, err := CreateServer(cfg, nil)
+			svr, err := CreateServer(cfg, nil, TestMode())
 			c.Assert(err, IsNil)
-			svr.SetUpTestMode()
 			err = svr.Run()
 			c.Assert(err, IsNil)
 			ch <- svr
@@ -149,9 +147,8 @@ func newTestServersWithCfgs(c *C, cfgs []*Config) ([]*Server, cleanupFunc) {
 	ch := make(chan *Server)
 	for _, cfg := range cfgs {
 		go func(cfg *Config) {
-			svr, err := CreateServer(cfg, nil)
+			svr, err := CreateServer(cfg, nil, TestMode())
 			c.Assert(err, IsNil)
-			svr.SetUpTestMode()
 			err = svr.Run()
 			c.Assert(err, IsNil)
 			ch <- svr
@@ -269,9 +266,8 @@ func (s *testServerSuite) TestCheckClusterID(c *C) {
 
 	// Start previous cluster, expect an error.
 	cfgA.InitialCluster = originInitial
-	svr, err := CreateServer(cfgA, nil)
+	svr, err := CreateServer(cfgA, nil, TestMode())
 	c.Assert(err, IsNil)
-	svr.SetUpTestMode()
 	err = svr.Run()
 	c.Assert(err, NotNil)
 }
diff --git a/server/testutil.go b/server/testutil.go
@@ -39,11 +39,10 @@ func cleanServer(cfg *Config) {
 // NewTestServer creates a pd server for testing.
 func NewTestServer() (*Config, *Server, CleanupFunc, error) {
 	cfg := NewTestSingleConfig()
-	s, err := CreateServer(cfg, nil)
+	s, err := CreateServer(cfg, nil, TestMode())
 	if err != nil {
 		return nil, nil, nil, errors.Trace(err)
 	}
-	s.SetUpTestMode()
 	if err = s.Run(); err != nil {
 		return nil, nil, nil, errors.Trace(err)
 	}