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