From d1a9ccb2b9f754976a7a1f6173f48bf040987a95 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Tue, 21 Mar 2017 16:09:45 -0700
Subject: [PATCH] clientv3: wait for Get goroutine in TestDialCancel

---
 clientv3/client_test.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/clientv3/client_test.go b/clientv3/client_test.go
index 8aba02538..8a61aba5f 100644
--- a/clientv3/client_test.go
+++ b/clientv3/client_test.go
@@ -49,7 +49,13 @@ func TestDialCancel(t *testing.T) {
 	c.SetEndpoints("http://254.0.0.1:12345")
 
 	// issue Get to force redial attempts
-	go c.Get(context.TODO(), "abc")
+	getc := make(chan struct{})
+	go func() {
+		defer close(getc)
+		// Get may hang forever on grpc's Stream.Header() if its
+		// context is never canceled.
+		c.Get(c.Ctx(), "abc")
+	}()
 
 	// wait a little bit so client close is after dial starts
 	time.Sleep(100 * time.Millisecond)
@@ -65,6 +71,11 @@ func TestDialCancel(t *testing.T) {
 		t.Fatalf("failed to close")
 	case <-donec:
 	}
+	select {
+	case <-time.After(5 * time.Second):
+		t.Fatalf("get failed to exit")
+	case <-getc:
+	}
 }
 
 func TestDialTimeout(t *testing.T) {