From 011db228a82ec0481b24ff652ef1927c44b80f0b Mon Sep 17 00:00:00 2001
From: Antonio Ojea <aojea@google.com>
Date: Mon, 25 Sep 2023 10:31:46 +0000
Subject: [PATCH] deflake TestWebSocketClient_HeartbeatSucceeds

Change-Id: I1f3331255ed81183b25ec24084c00d0cf7dfb510

Kubernetes-commit: 4a74958b2df7ecd9e6a4867f7b10f1402cb2e215
---
 tools/remotecommand/websocket_test.go | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tools/remotecommand/websocket_test.go b/tools/remotecommand/websocket_test.go
index 22c28c89..2895ba54 100644
--- a/tools/remotecommand/websocket_test.go
+++ b/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) {