k8s.io/apiserver@v0.31.1/pkg/storage/etcd3/testserver/test_server.go (about)

     1  /*
     2  Copyright 2021 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package testserver
    18  
    19  import (
    20  	"fmt"
    21  	"net"
    22  	"net/url"
    23  	"os"
    24  	"strconv"
    25  	"testing"
    26  	"time"
    27  
    28  	clientv3 "go.etcd.io/etcd/client/v3"
    29  	"go.etcd.io/etcd/server/v3/embed"
    30  	"go.uber.org/zap/zapcore"
    31  	"go.uber.org/zap/zaptest"
    32  	"google.golang.org/grpc"
    33  )
    34  
    35  // getAvailablePort returns a TCP port that is available for binding.
    36  func getAvailablePorts(count int) ([]int, error) {
    37  	ports := []int{}
    38  	for i := 0; i < count; i++ {
    39  		l, err := net.Listen("tcp", ":0")
    40  		if err != nil {
    41  			return nil, fmt.Errorf("could not bind to a port: %v", err)
    42  		}
    43  		// It is possible but unlikely that someone else will bind this port before we get a chance to use it.
    44  		defer l.Close()
    45  		ports = append(ports, l.Addr().(*net.TCPAddr).Port)
    46  	}
    47  	return ports, nil
    48  }
    49  
    50  // NewTestConfig returns a configuration for an embedded etcd server.
    51  // The configuration is based on embed.NewConfig(), with the following adjustments:
    52  //   - sets UnsafeNoFsync = true to improve test performance (only reasonable in a test-only
    53  //     single-member server we never intend to restart or keep data from)
    54  //   - uses free ports for client and peer listeners
    55  //   - cleans up the data directory on test termination
    56  //   - silences server logs other than errors
    57  func NewTestConfig(t testing.TB) *embed.Config {
    58  	cfg := embed.NewConfig()
    59  
    60  	cfg.UnsafeNoFsync = true
    61  
    62  	ports, err := getAvailablePorts(2)
    63  	if err != nil {
    64  		t.Fatal(err)
    65  	}
    66  	clientURL := url.URL{Scheme: "http", Host: net.JoinHostPort("localhost", strconv.Itoa(ports[0]))}
    67  	peerURL := url.URL{Scheme: "http", Host: net.JoinHostPort("localhost", strconv.Itoa(ports[1]))}
    68  
    69  	cfg.ListenPeerUrls = []url.URL{peerURL}
    70  	cfg.AdvertisePeerUrls = []url.URL{peerURL}
    71  	cfg.ListenClientUrls = []url.URL{clientURL}
    72  	cfg.AdvertiseClientUrls = []url.URL{clientURL}
    73  	cfg.InitialCluster = cfg.InitialClusterFromName(cfg.Name)
    74  
    75  	cfg.ZapLoggerBuilder = embed.NewZapLoggerBuilder(zaptest.NewLogger(t, zaptest.Level(zapcore.ErrorLevel)).Named("etcd-server"))
    76  	cfg.Dir = t.TempDir()
    77  	os.Chmod(cfg.Dir, 0700)
    78  	return cfg
    79  }
    80  
    81  // RunEtcd starts an embedded etcd server with the provided config
    82  // (or NewTestConfig(t) if nil), and returns a client connected to the server.
    83  // The server is terminated when the test ends.
    84  func RunEtcd(t testing.TB, cfg *embed.Config) *clientv3.Client {
    85  	t.Helper()
    86  
    87  	if cfg == nil {
    88  		cfg = NewTestConfig(t)
    89  	}
    90  
    91  	e, err := embed.StartEtcd(cfg)
    92  	if err != nil {
    93  		t.Fatal(err)
    94  	}
    95  	t.Cleanup(e.Close)
    96  
    97  	select {
    98  	case <-e.Server.ReadyNotify():
    99  	case <-time.After(60 * time.Second):
   100  		e.Server.Stop() // trigger a shutdown
   101  		t.Fatal("server took too long to start")
   102  	}
   103  	go func() {
   104  		err := <-e.Err()
   105  		if err != nil {
   106  			t.Error(err)
   107  		}
   108  	}()
   109  
   110  	tlsConfig, err := cfg.ClientTLSInfo.ClientConfig()
   111  	if err != nil {
   112  		t.Fatal(err)
   113  	}
   114  
   115  	client, err := clientv3.New(clientv3.Config{
   116  		TLS:         tlsConfig,
   117  		Endpoints:   e.Server.Cluster().ClientURLs(),
   118  		DialTimeout: 10 * time.Second,
   119  		DialOptions: []grpc.DialOption{grpc.WithBlock()},
   120  		Logger:      zaptest.NewLogger(t, zaptest.Level(zapcore.ErrorLevel)).Named("etcd-client"),
   121  	})
   122  	if err != nil {
   123  		t.Fatal(err)
   124  	}
   125  	return client
   126  }