Fix flake in Quic EndToEndTest.KeyUpdateInitiatedByServer & KeyUpdateInitiatedByBoth

The server may not have received an ACK for the current key phase by the time the test tried to call InitiateKeyUpdate().

PiperOrigin-RevId: 337116333
Change-Id: Id009f4d080232781ee15b5e2a3dced7a8d04da2e

quic/core/http/end_to_end_test.cc

diff --git a/quic/core/http/end_to_end_test.cc b/quic/core/http/end_to_end_test.cc
index 43d95a3..dddeb1f 100644
--- a/quic/core/http/end_to_end_test.cc
+++ b/quic/core/http/end_to_end_test.cc
@@ -4878,6 +4878,11 @@
       [this]() {
         QuicConnection* server_connection = GetServerConnection();
         if (server_connection != nullptr) {
+          if (!server_connection->IsKeyUpdateAllowed()) {
+            // Server may not have received ack from client yet for the current
+            // key phase, wait a bit and try again.
+            return false;
+          }
           EXPECT_TRUE(server_connection->InitiateKeyUpdate());
         } else {
           ADD_FAILURE() << "Missing server connection";
@@ -4896,6 +4901,9 @@
       [this]() {
         QuicConnection* server_connection = GetServerConnection();
         if (server_connection != nullptr) {
+          if (!server_connection->IsKeyUpdateAllowed()) {
+            return false;
+          }
           EXPECT_TRUE(server_connection->InitiateKeyUpdate());
         } else {
           ADD_FAILURE() << "Missing server connection";
@@ -4943,6 +4951,11 @@
       [this]() {
         QuicConnection* server_connection = GetServerConnection();
         if (server_connection != nullptr) {
+          if (!server_connection->IsKeyUpdateAllowed()) {
+            // Server may not have received ack from client yet for the current
+            // key phase, wait a bit and try again.
+            return false;
+          }
           EXPECT_TRUE(server_connection->InitiateKeyUpdate());
         } else {
           ADD_FAILURE() << "Missing server connection";
@@ -4964,6 +4977,9 @@
       [this]() {
         QuicConnection* server_connection = GetServerConnection();
         if (server_connection != nullptr) {
+          if (!server_connection->IsKeyUpdateAllowed()) {
+            return false;
+          }
           EXPECT_TRUE(server_connection->InitiateKeyUpdate());
         } else {
           ADD_FAILURE() << "Missing server connection";