From d01510416cddf4d3e509265e55262525640afa0e Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Mon, 9 Oct 2017 11:23:30 -0700
Subject: [PATCH] clientv3/integration: add balancer network partition tests

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

diff --git a/clientv3/integration/network_partition_test.go b/clientv3/integration/network_partition_test.go
new file mode 100644
index 000000000000..f13fae197769
--- /dev/null
+++ b/clientv3/integration/network_partition_test.go
@@ -0,0 +1,100 @@
+// 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.
+
+// +build !cluster_proxy
+
+package integration
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/coreos/etcd/clientv3"
+	"github.com/coreos/etcd/integration"
+	"github.com/coreos/etcd/pkg/testutil"
+)
+
+// TestNetworkPartitionBalancerWithPut tests when one member becomes isolated,
+// first Put requests fail and following retry succeeds with client balancer
+// switching to others.
+func TestNetworkPartitionBalancerWithPut(t *testing.T) {
+	testNetworkPartitionBalancer(t, func(cli *clientv3.Client) {
+		var err error
+		for i := 0; i < 2; i++ {
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			_, err = cli.Put(ctx, "a", "b")
+			cancel()
+			if err == nil {
+				break
+			}
+			if err != context.DeadlineExceeded {
+				t.Fatalf("#%d: expected %v, got %v", i, context.DeadlineExceeded, err)
+			}
+			if i == 0 { // give enough time for endpoint switch
+				time.Sleep(5 * time.Second)
+			}
+		}
+		if err != nil {
+			t.Fatalf("balancer did not switch in time (%v)", err)
+		}
+	})
+}
+
+// TestNetworkPartitionBalancerWithGet tests when one member becomes isolated,
+// quorum Get request succeeds via internal retry function, while triggering
+// endpoint switch out of network partition.
+func TestNetworkPartitionBalancerWithGet(t *testing.T) {
+	testNetworkPartitionBalancer(t, func(cli *clientv3.Client) {
+		// give enough timeout for balancer endpoint switch
+		ctx, cancel := context.WithTimeout(context.Background(), 7*time.Second)
+		_, err := cli.Get(ctx, "a")
+		cancel()
+		if err != nil {
+			t.Errorf("want err == nil, got %v", err)
+		}
+	})
+}
+
+func testNetworkPartitionBalancer(t *testing.T, op func(*clientv3.Client)) {
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{
+		Size:                 3,
+		GRPCKeepAliveMinTime: time.Millisecond, // avoid too_many_pings
+	})
+	defer clus.Terminate(t)
+
+	// expect pin ep[0]
+	ccfg := clientv3.Config{
+		Endpoints:            []string{clus.Members[0].GRPCAddr()},
+		DialTimeout:          3 * time.Second,
+		DialKeepAliveTime:    2 * time.Second,
+		DialKeepAliveTimeout: 2 * time.Second,
+	}
+	cli, err := clientv3.New(ccfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cli.Close()
+
+	// add other endpoints for later endpoint switch
+	cli.SetEndpoints(clus.Members[0].GRPCAddr(), clus.Members[1].GRPCAddr(), clus.Members[1].GRPCAddr())
+
+	time.Sleep(3 * time.Second)
+	clus.Members[0].InjectPartition(t, clus.Members[1:])
+	defer clus.Members[0].RecoverPartition(t, clus.Members[1:])
+
+	op(cli)
+}