From eaa5d9772ff11b970ed70d49f4939e074ad0eec0 Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Fri, 19 Aug 2016 13:02:22 -0700
Subject: [PATCH] integration: improve TestTransferLeader

so that it can check leader transition
---
 integration/cluster.go      | 14 +++++-----
 integration/cluster_test.go | 53 ++++++++++++++++++++++---------------
 2 files changed, 39 insertions(+), 28 deletions(-)

diff --git a/integration/cluster.go b/integration/cluster.go
index 326a12d6d..ab4832d0a 100644
--- a/integration/cluster.go
+++ b/integration/cluster.go
@@ -702,13 +702,13 @@ func (m *member) Stop(t *testing.T) {
 	plog.Printf("stopped %s (%s)", m.Name, m.grpcAddr)
 }
 
-// StopWithAutoLeaderTransfer stops the member with auto leader transfer.
-func (m *member) StopWithAutoLeaderTransfer(t *testing.T) {
-	plog.Printf("stopping %s (%s)", m.Name, m.grpcAddr)
-	m.s.TransferLeadership()
-	m.Close()
-	m.hss = nil
-	plog.Printf("stopped %s (%s)", m.Name, m.grpcAddr)
+// checkLeaderTransition waits for leader transition, returning the new leader ID.
+func checkLeaderTransition(t *testing.T, m *member, oldLead uint64) uint64 {
+	interval := time.Duration(m.s.Cfg.TickMs) * time.Millisecond
+	for m.s.Lead() == 0 || (m.s.Lead() == oldLead) {
+		time.Sleep(interval)
+	}
+	return m.s.Lead()
 }
 
 // StopNotify unblocks when a member stop completes
diff --git a/integration/cluster_test.go b/integration/cluster_test.go
index f91bc428c..c3658c4ba 100644
--- a/integration/cluster_test.go
+++ b/integration/cluster_test.go
@@ -26,7 +26,6 @@ import (
 
 	"github.com/coreos/etcd/client"
 	"github.com/coreos/etcd/etcdserver"
-	"github.com/coreos/etcd/etcdserver/etcdserverpb"
 	"github.com/coreos/etcd/pkg/testutil"
 
 	"golang.org/x/net/context"
@@ -470,29 +469,41 @@ func TestTransferLeader(t *testing.T) {
 	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
 	defer clus.Terminate(t)
 
-	leaderIdx := clus.WaitLeader(t)
-
-	err := clus.Members[leaderIdx].s.TransferLeadership()
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
-func TestTransferLeaderStopTrigger(t *testing.T) {
-	defer testutil.AfterTest(t)
-
-	clus := NewClusterV3(t, &ClusterConfig{Size: 3})
-	defer clus.Terminate(t)
-
 	oldLeadIdx := clus.WaitLeader(t)
-	clus.Members[oldLeadIdx].StopWithAutoLeaderTransfer(t)
+	oldLeadID := uint64(clus.Members[oldLeadIdx].s.ID())
 
-	// issue put to one of the other member
-	kvc := toGRPC(clus.Client((oldLeadIdx + 1) % 3)).KV
-	sctx, scancel := context.WithTimeout(context.TODO(), clus.Members[oldLeadIdx].electionTimeout())
-	_, err := kvc.Range(sctx, &etcdserverpb.RangeRequest{Key: []byte("foo")})
-	scancel()
+	// ensure followers go through leader transition while learship transfer
+	idc := make(chan uint64)
+	for i := range clus.Members {
+		if oldLeadIdx != i {
+			go func(m *member) {
+				idc <- checkLeaderTransition(t, m, oldLeadID)
+			}(clus.Members[i])
+		}
+	}
+
+	err := clus.Members[oldLeadIdx].s.TransferLeadership()
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	// wait until leader transitions have happened
+	var newLeadIDs [2]uint64
+	for i := range newLeadIDs {
+		select {
+		case newLeadIDs[i] = <-idc:
+		case <-time.After(time.Second):
+			t.Fatal("timed out waiting for leader transition")
+		}
+	}
+
+	// remaining members must agree on the same leader
+	if newLeadIDs[0] != newLeadIDs[1] {
+		t.Fatalf("expected same new leader %d == %d", newLeadIDs[0], newLeadIDs[1])
+	}
+
+	// new leader must be different than the old leader
+	if oldLeadID == newLeadIDs[0] {
+		t.Fatalf("expected old leader %d != new leader %d", oldLeadID, newLeadIDs[0])
+	}
 }