From 9d95cfb105f635fc3c936c9ea87814f65864bc2c Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Thu, 26 Oct 2017 21:39:12 -0700
Subject: [PATCH] clientv3/integration: add
 TestBalancerUnderServerShutdownImmutable

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

diff --git a/clientv3/integration/server_shutdown_test.go b/clientv3/integration/server_shutdown_test.go
index de128eb8d50..4cfe218887b 100644
--- a/clientv3/integration/server_shutdown_test.go
+++ b/clientv3/integration/server_shutdown_test.go
@@ -179,3 +179,58 @@ func testBalancerUnderServerShutdownMutable(t *testing.T, op func(*clientv3.Clie
 		t.Fatal(err)
 	}
 }
+
+func TestBalancerUnderServerShutdownGetLinearizable(t *testing.T) {
+	testBalancerUnderServerShutdownImmutable(t, func(cli *clientv3.Client, ctx context.Context) error {
+		_, err := cli.Get(ctx, "foo")
+		return err
+	}, 7*time.Second) // give enough time for leader election, balancer switch
+}
+
+func TestBalancerUnderServerShutdownGetSerializable(t *testing.T) {
+	testBalancerUnderServerShutdownImmutable(t, func(cli *clientv3.Client, ctx context.Context) error {
+		_, err := cli.Get(ctx, "foo", clientv3.WithSerializable())
+		return err
+	}, 2*time.Second)
+}
+
+// testBalancerUnderServerShutdownImmutable expects that when the member of
+// the pinned endpoint is shut down, the balancer switches its endpoints
+// and all subsequent range requests succeed with new endpoints.
+func testBalancerUnderServerShutdownImmutable(t *testing.T, op func(*clientv3.Client, context.Context) error, timeout time.Duration) {
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{
+		Size:               3,
+		SkipCreatingClient: true,
+	})
+	defer clus.Terminate(t)
+
+	eps := []string{clus.Members[0].GRPCAddr(), clus.Members[1].GRPCAddr(), clus.Members[2].GRPCAddr()}
+
+	// pin eps[0]
+	cli, err := clientv3.New(clientv3.Config{Endpoints: []string{eps[0]}})
+	if err != nil {
+		t.Errorf("failed to create client: %v", err)
+	}
+	defer cli.Close()
+
+	// wait for eps[0] to be pinned
+	waitPinReady(t, cli)
+
+	// add all eps to list, so that when the original pined one fails
+	// the client can switch to other available eps
+	cli.SetEndpoints(eps...)
+
+	// shut down eps[0]
+	clus.Members[0].Terminate(t)
+
+	// switched to others when eps[0] was explicitly shut down
+	// and following request should succeed
+	cctx, ccancel := context.WithTimeout(context.Background(), timeout)
+	err = op(cli, cctx)
+	ccancel()
+	if err != nil {
+		t.Errorf("failed to finish range request in time %v (timeout %v)", err, timeout)
+	}
+}