-
Notifications
You must be signed in to change notification settings - Fork 1
/
helper.go
85 lines (71 loc) · 1.86 KB
/
helper.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
package mock
import (
"fmt"
"os"
"strings"
"sync"
"testing"
"time"
"github.com/deepfabric/prophet/option"
"github.com/deepfabric/prophet/util"
"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/clientv3"
"go.etcd.io/etcd/embed"
)
var (
mutex sync.Mutex
ports = 10000
)
func nextTestPorts() int {
mutex.Lock()
defer mutex.Unlock()
ports++
return ports
}
// NewEtcdClient create a etcd client
func NewEtcdClient(t *testing.T, port int) *clientv3.Client {
client, err := clientv3.New(clientv3.Config{
Endpoints: strings.Split(fmt.Sprintf("http://127.0.0.1:%d", port), ","),
DialTimeout: option.DefaultTimeout,
})
if err != nil {
assert.FailNowf(t, "create etcd client failed", "error: %+v", err)
}
return client
}
// StartTestSingleEtcd start a single etcd server
func StartTestSingleEtcd(t *testing.T) (chan interface{}, int) {
port := nextTestPorts()
peerPort := nextTestPorts()
now := time.Now().UnixNano()
cfg := embed.NewConfig()
cfg.Name = "p1"
cfg.Dir = fmt.Sprintf("%s/prophet/test-%d", os.TempDir(), now)
cfg.WalDir = ""
cfg.InitialCluster = fmt.Sprintf("p1=http://127.0.0.1:%d", peerPort)
cfg.ClusterState = embed.ClusterStateFlagNew
cfg.EnablePprof = false
cfg.Debug = false
cfg.LPUrls, _ = util.ParseUrls(fmt.Sprintf("http://127.0.0.1:%d", peerPort))
cfg.APUrls = cfg.LPUrls
cfg.LCUrls, _ = util.ParseUrls(fmt.Sprintf("http://127.0.0.1:%d", port))
cfg.ACUrls = cfg.LCUrls
etcd, err := embed.StartEtcd(cfg)
if err != nil {
assert.FailNowf(t, "start embed etcd failed", "error: %+v", err)
}
select {
case <-etcd.Server.ReadyNotify():
time.Sleep(time.Millisecond * 100)
stopC := make(chan interface{})
go func() {
<-stopC
etcd.Server.Stop()
os.RemoveAll(cfg.Dir)
}()
return stopC, port
case <-time.After(time.Minute * 5):
assert.FailNowf(t, "start embed etcd failed", "error: timeout", err)
}
return nil, 0
}