deflake TestWebSocketClient_HeartbeatSucceeds

Change-Id: I1f3331255ed81183b25ec24084c00d0cf7dfb510

staging/src/k8s.io/client-go/tools/remotecommand/websocket_test.go

@@ -1057,8 +1057,7 @@ func TestWebSocketClient_HeartbeatSucceeds(t *testing.T) {
 		for {
 			_, _, err := conn.ReadMessage()
 			if err != nil {
-				t.Logf("server err reading message: %v", err)
-				return
+				break
 			}
 		}
 	}))
@@ -1085,7 +1084,11 @@ func TestWebSocketClient_HeartbeatSucceeds(t *testing.T) {
 		return pongHandler(msg)
 	})
 	go heartbeat.start()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		for {
 			_, _, err := client.ReadMessage()
 			if err != nil {
@@ -1107,6 +1110,7 @@ func TestWebSocketClient_HeartbeatSucceeds(t *testing.T) {
 		close(heartbeat.closer)
 		t.Errorf("unexpected heartbeat timeout")
 	}
+	wg.Wait()
 }
 
 func TestWebSocketClient_StreamsAndExpectedErrors(t *testing.T) {