Make websocket heartbeat test timing less flaky

Kubernetes-commit: 26484df2108eff8ad6e06dfc960eae3bdfbf4663

tools/remotecommand/websocket_test.go

@@ -817,15 +817,16 @@ func TestWebSocketClient_BadHandshake(t *testing.T) {
 // TestWebSocketClient_HeartbeatTimeout tests the heartbeat by forcing a
 // timeout by setting the ping period greater than the deadline.
 func TestWebSocketClient_HeartbeatTimeout(t *testing.T) {
+ blockRequestCtx, unblockRequest := context.WithCancel(context.Background())
+ defer unblockRequest()
  // Create fake WebSocket server which blocks.
  websocketServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
  conns, err := webSocketServerStreams(req, w, streamOptionsFromRequest(req))
  if err != nil {
  t.Fatalf("error on webSocketServerStreams: %v", err)
  }
  defer conns.conn.Close()
- // Block server; heartbeat timeout (or test timeout) will fire before this returns.
- time.Sleep(1 * time.Second)
+ <-blockRequestCtx.Done()
  }))
  defer websocketServer.Close()
  // Create websocket client connecting to fake server.
@@ -840,8 +841,8 @@ func TestWebSocketClient_HeartbeatTimeout(t *testing.T) {
  }
  streamExec := exec.(*wsStreamExecutor)
  // Ping period is greater than the ping deadline, forcing the timeout to fire.
- pingPeriod := 20 * time.Millisecond
- pingDeadline := 5 * time.Millisecond
+ pingPeriod := wait.ForeverTestTimeout // this lets the heartbeat deadline expire without renewing it
+ pingDeadline := time.Second // this gives setup 1 second to establish streams
  streamExec.heartbeatPeriod = pingPeriod
  streamExec.heartbeatDeadline = pingDeadline
  // Send some random data to the websocket server through STDIN.
@@ -859,8 +860,7 @@ func TestWebSocketClient_HeartbeatTimeout(t *testing.T) {
  }()
 
  select {
- case <-time.After(pingPeriod * 5):
- // Give up after about five ping attempts
+ case <-time.After(wait.ForeverTestTimeout):
  t.Fatalf("expected heartbeat timeout, got none.")
  case err := <-errorChan:
  // Expecting heartbeat timeout error.