From ff2e5926c2a971c8b781e56415a4f1dc0b14b02e Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Tue, 24 Oct 2017 15:50:26 -0700
Subject: [PATCH] clientv3/integration: add TestBalancerUnderServerFailureWatch

Current Watch integration tests haven't covered the balancer
switch behavior under server failures.

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
---
 clientv3/integration/server_failure_test.go | 99 +++++++++++++++++++++
 1 file changed, 99 insertions(+)
 create mode 100644 clientv3/integration/server_failure_test.go

diff --git a/clientv3/integration/server_failure_test.go b/clientv3/integration/server_failure_test.go
new file mode 100644
index 000000000000..7055f5793973
--- /dev/null
+++ b/clientv3/integration/server_failure_test.go
@@ -0,0 +1,99 @@
+// Copyright 2017 The etcd 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 integration
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/coreos/etcd/clientv3"
+	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
+	"github.com/coreos/etcd/integration"
+	"github.com/coreos/etcd/pkg/testutil"
+)
+
+// TestBalancerUnderServerFailureWatch expects that watch client
+// switch its endpoints when the member of the pinned endpoint fails.
+func TestBalancerUnderServerFailureWatch(t *testing.T) {
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
+	defer clus.Terminate(t)
+
+	eps := []string{clus.Members[0].GRPCAddr(), clus.Members[1].GRPCAddr(), clus.Members[2].GRPCAddr()}
+
+	lead := clus.WaitLeader(t)
+
+	// close unnecessary clients
+	clus.Client(lead).Close()
+	clus.TakeClient(lead)
+	clus.Client((lead + 2) % 3).Close()
+	clus.TakeClient((lead + 2) % 3)
+
+	// pin eps[lead]
+	cli, err := clientv3.New(clientv3.Config{Endpoints: []string{eps[lead]}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cli.Close()
+	cli.SetEndpoints(eps...)
+
+	wch := cli.Watch(context.Background(), "foo", clientv3.WithCreatedNotify())
+	<-wch
+
+	donec := make(chan struct{})
+	go func() {
+		defer close(donec)
+
+		// switch to others when eps[lead] fails
+		for ev := range wch {
+			received := false
+			for _, e := range ev.Events {
+				if string(e.Kv.Value) != "bar" {
+					t.Fatalf("expected 'bar', got %+v", e.Kv)
+				}
+				received = true
+			}
+			if werr := ev.Err(); werr != nil {
+				t.Fatal(werr)
+			}
+			if received {
+				break
+			}
+		}
+	}()
+
+	// fail eps[lead]
+	clus.Members[lead].Terminate(t)
+
+	// writes to eps[lead+1]
+	for {
+		_, err = clus.Client((lead+1)%3).Put(context.Background(), "foo", "bar")
+		if err == nil {
+			break
+		}
+		if err == rpctypes.ErrTimeout || err == rpctypes.ErrTimeoutDueToLeaderFail {
+			continue
+		}
+		t.Fatal(err)
+	}
+
+	select {
+	case <-donec:
+	case <-time.After(5 * time.Second): // enough time for leader election, balancer switch
+		t.Fatal("took too long to receive events")
+	}
+}