From f9b505ae56b7d6bf0a18d6f5b49400554dbaf8be Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Sun, 17 Jan 2016 14:11:30 -0800
Subject: [PATCH] integration: use WaitResponse for watch tests

---
 integration/v3_grpc_test.go | 34 ++++++----------------------------
 1 file changed, 6 insertions(+), 28 deletions(-)

diff --git a/integration/v3_grpc_test.go b/integration/v3_grpc_test.go
index 238cf058e..7cf6fcb88 100644
--- a/integration/v3_grpc_test.go
+++ b/integration/v3_grpc_test.go
@@ -412,20 +412,9 @@ func TestV3WatchFromCurrentRevision(t *testing.T) {
 			}
 		}
 
-		rCh := make(chan *pb.WatchResponse)
-		go func() {
-			resp, _ := wStream.Recv()
-			rCh <- resp
-		}()
-		select {
-		case nr := <-rCh:
-			t.Errorf("#%d: unexpected response is received %+v", i, nr)
-		case <-time.After(2 * time.Second):
-		}
-		wStream.CloseSend()
-		rv, ok := <-rCh
-		if rv != nil || !ok {
-			t.Errorf("#%d: rv, ok got = %v %v, want = nil true", i, rv, ok)
+		rok, nr := WaitResponse(wStream, 1*time.Second)
+		if !rok {
+			t.Errorf("unexpected pb.WatchResponse is received %+v", nr)
 		}
 
 		// can't defer because tcp ports will be in use
@@ -473,20 +462,9 @@ func TestV3WatchCancel(t *testing.T) {
 	}
 
 	// watch got canceled, so this should block
-	rCh := make(chan *pb.WatchResponse)
-	go func() {
-		resp, _ := wStream.Recv()
-		rCh <- resp
-	}()
-	select {
-	case nr := <-rCh:
-		t.Errorf("unexpected response is received %+v", nr)
-	case <-time.After(2 * time.Second):
-	}
-	wStream.CloseSend()
-	rv, ok := <-rCh
-	if rv != nil || !ok {
-		t.Errorf("rv, ok got = %v %v, want = nil true", rv, ok)
+	rok, nr := WaitResponse(wStream, 1*time.Second)
+	if !rok {
+		t.Errorf("unexpected pb.WatchResponse is received %+v", nr)
 	}
 
 	clus.Terminate(t)