From f7417cf6937d18d93a4381081657329e3678aa1e Mon Sep 17 00:00:00 2001
From: Yicheng Qin <qycqycqycqycqyc@gmail.com>
Date: Tue, 19 Aug 2014 13:08:49 -0700
Subject: [PATCH] server: add testCluster.Add

and fix possible testing error.
---
 etcd/etcd_start_test.go |  6 ++----
 etcd/etcd_test.go       | 21 +++++++++++++++++----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/etcd/etcd_start_test.go b/etcd/etcd_start_test.go
index 6b5140fad..a8ced924a 100644
--- a/etcd/etcd_start_test.go
+++ b/etcd/etcd_start_test.go
@@ -151,11 +151,9 @@ func TestRunByDiscoveryService(t *testing.T) {
 	c := newTestConfig()
 	c.Name = bootstrapName
 	c.Discovery = ds.URL(0) + "/v2/keys/_etcd/registry/1"
-	ts := testServer{Config: c}
-	ts.Start()
-	defer ts.Destroy()
+	ts := &testServer{Config: c}
+	ds.Add(ts)
 
-	ts.WaitMode(participantMode)
 	// wait for the leader to do a heartbeat
 	// it will update the lead field of the follower
 	time.Sleep(ds.Node(0).e.tickDuration * defaultHeartbeat * 2)
diff --git a/etcd/etcd_test.go b/etcd/etcd_test.go
index 6476bc2bf..527151950 100644
--- a/etcd/etcd_test.go
+++ b/etcd/etcd_test.go
@@ -250,10 +250,8 @@ func TestRestoreSnapshotFromLeader(t *testing.T) {
 	c := newTestConfig()
 	c.Name = "1"
 	c.Peers = []string{cl.URL(0)}
-	ts := testServer{Config: c}
-	ts.Start()
-	defer ts.Destroy()
-	ts.WaitMode(participantMode)
+	ts := &testServer{Config: c}
+	cl.Add(ts)
 
 	// check new proposal could be submitted
 	if _, err := cl.Participant(0).Set("/foo", false, "bar", store.Permanent); err != nil {
@@ -280,6 +278,7 @@ func TestRestoreSnapshotFromLeader(t *testing.T) {
 	}
 	<-wch.EventChan
 
+	// check node map of two machines are the same
 	g := ts.Participant().node.Nodes()
 	w := cl.Participant(0).node.Nodes()
 	if !reflect.DeepEqual(g, w) {
@@ -352,6 +351,20 @@ func (c *testCluster) wait() {
 	}
 }
 
+func (c *testCluster) Add(s *testServer) {
+	lead, _ := c.Leader()
+	// wait for the node to join the cluster
+	// TODO(yichengq): remove this when we get rid of all timeouts
+	wch, err := c.Participant(int(lead)).Watch(v2machineKVPrefix, true, false, 0)
+	if err != nil {
+		panic(err)
+	}
+	s.Start()
+	<-wch.EventChan
+	c.Size++
+	c.nodes = append(c.nodes, s)
+}
+
 func (c *testCluster) Node(i int) *testServer {
 	return c.nodes[i]
 }