From 1a571b746ace04b72607c3f9ad2399442b58614d Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Mon, 18 Jul 2016 16:00:17 -0700
Subject: [PATCH] Add TestFailFastRPCWithNoBalancerErrorOnBadCertificates
 TestNonFailFastRPCWithNoBalancerErrorOnBadCertificates

---
 test/end2end_test.go | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/test/end2end_test.go b/test/end2end_test.go
index b258893a..32dd4e69 100644
--- a/test/end2end_test.go
+++ b/test/end2end_test.go
@@ -369,6 +369,7 @@ type test struct {
 	serverCompression bool
 	unaryInt          grpc.UnaryServerInterceptor
 	streamInt         grpc.StreamServerInterceptor
+	balancer          grpc.Balancer
 
 	// srv and srvAddr are set once startServer is called.
 	srv     *grpc.Server
@@ -404,6 +405,8 @@ func newTest(t *testing.T, e env) *test {
 		maxStream: math.MaxUint32,
 	}
 	te.ctx, te.cancel = context.WithCancel(context.Background())
+	// Install roundrobin balancer by default.
+	te.balancer = grpc.RoundRobin(nil)
 	return te
 }
 
@@ -502,6 +505,9 @@ func (te *test) clientConn() *grpc.ClientConn {
 	default:
 		opts = append(opts, grpc.WithInsecure())
 	}
+	if te.balancer != nil {
+		opts = append(opts, grpc.WithBalancer(te.balancer))
+	}
 	var err error
 	te.cc, err = grpc.Dial(te.srvAddr, opts...)
 	if err != nil {
@@ -2317,6 +2323,34 @@ func TestFailFastRPCErrorOnBadCertificates(t *testing.T) {
 	}
 }
 
+func TestFailFastRPCWithNoBalancerErrorOnBadCertificates(t *testing.T) {
+	te := newTest(t, env{name: "bad-cred", network: "tcp", security: "clientAlwaysFailCred"})
+	// Uninstall balancer.
+	te.balancer = nil
+	te.startServer(&testServer{security: "clientAlwaysFailCred"})
+	defer te.tearDown()
+
+	cc := te.clientConn()
+	tc := testpb.NewTestServiceClient(cc)
+	if _, err := tc.EmptyCall(context.Background(), &testpb.Empty{}); !strings.Contains(err.Error(), clientAlwaysFailCredErrorMsg) {
+		te.t.Fatalf("Dial(%q) = %v, want err.Error() contains %q", te.srvAddr, err, clientAlwaysFailCredErrorMsg)
+	}
+}
+
+func TestNonFailFastRPCWithNoBalancerErrorOnBadCertificates(t *testing.T) {
+	te := newTest(t, env{name: "bad-cred", network: "tcp", security: "clientAlwaysFailCred"})
+	// Uninstall balancer.
+	te.balancer = nil
+	te.startServer(&testServer{security: "clientAlwaysFailCred"})
+	defer te.tearDown()
+
+	cc := te.clientConn()
+	tc := testpb.NewTestServiceClient(cc)
+	if _, err := tc.EmptyCall(context.Background(), &testpb.Empty{}, grpc.FailFast(false)); !strings.Contains(err.Error(), clientAlwaysFailCredErrorMsg) {
+		te.t.Fatalf("Dial(%q) = %v, want err.Error() contains %q", te.srvAddr, err, clientAlwaysFailCredErrorMsg)
+	}
+}
+
 // interestingGoroutines returns all goroutines we care about for the purpose
 // of leak checking. It excludes testing or runtime ones.
 func interestingGoroutines() (gs []string) {