Add Retry to Some Tests (microsoft#2511)

G-arj · Mar 11, 2022 · e4ffd5e · e4ffd5e
1 parent 02c36e4
commit e4ffd5e
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 19 deletions.
diff --git a/src/test/lib/HandshakeTest.cpp b/src/test/lib/HandshakeTest.cpp
@@ -2000,10 +2000,15 @@ QuicTestKeyUpdate(
                     // Send some data to perform the key update.
                     // TODO: Update this to send stream data, like QuicConnectAndPing does.
                     //
-                    TEST_QUIC_SUCCEEDED(Client.SetPeerBidiStreamCount((uint16_t)(101+i)));
-                    TEST_EQUAL((uint16_t)(101+i), Client.GetPeerBidiStreamCount());
-                    CxPlatSleep(100);
-                    TEST_EQUAL((uint16_t)(101+i), Server->GetLocalBidiStreamCount());
+                    uint16_t PeerCount, Expected = 101+i, Tries = 0;
+                    TEST_QUIC_SUCCEEDED(Client.SetPeerBidiStreamCount(Expected));
+                    TEST_EQUAL(Expected, Client.GetPeerBidiStreamCount());
+
+                    do {
+                        CxPlatSleep(100);
+                        PeerCount =  Server->GetLocalBidiStreamCount();
+                    } while (PeerCount != Expected && Tries++ < 10);
+                    TEST_EQUAL(Expected, PeerCount);
 
                     //
                     // Force a client key update to occur again to check for double update
@@ -2013,10 +2018,16 @@ QuicTestKeyUpdate(
                         TEST_QUIC_SUCCEEDED(Client.ForceKeyUpdate());
                     }
 
-                    TEST_QUIC_SUCCEEDED(Server->SetPeerBidiStreamCount((uint16_t)(100+i)));
-                    TEST_EQUAL((uint16_t)(100+i), Server->GetPeerBidiStreamCount());
-                    CxPlatSleep(100);
-                    TEST_EQUAL((uint16_t)(100+i), Client.GetLocalBidiStreamCount());
+                    Expected = 100+i;
+                    TEST_QUIC_SUCCEEDED(Server->SetPeerBidiStreamCount(Expected));
+                    TEST_EQUAL(Expected, Server->GetPeerBidiStreamCount());
+
+                    Tries = 0;
+                    do {
+                        CxPlatSleep(100);
+                        PeerCount =  Client.GetLocalBidiStreamCount();
+                    } while (PeerCount != Expected && Tries++ < 10);
+                    TEST_EQUAL(Expected, PeerCount);
                 }
 
                 CxPlatSleep(100);

diff --git a/src/test/lib/QuicDrill.cpp b/src/test/lib/QuicDrill.cpp
@@ -311,17 +311,15 @@ QuicDrillInitialPacketFailureTest(
             return false;
         }
 
-        //
-        // Generously wait for server to process packet.
-        //
-        CxPlatSleep(100);
-
-        Status = Listener.GetStatistics(Stats);
-        if (QUIC_FAILED(Status)) {
-            TEST_FAILURE("Get Listener statistics after test failed, 0x%x.", Status);
-            return false;
-        }
-        DroppedPacketsAfter = Stats.BindingRecvDroppedPackets;
+        uint32_t Tries = 0;
+        do {
+            CxPlatSleep(100);
+            if (QUIC_FAILED(Status = Listener.GetStatistics(Stats))) {
+                TEST_FAILURE("Get Listener statistics after test failed, 0x%x.", Status);
+                return false;
+            }
+            DroppedPacketsAfter = Stats.BindingRecvDroppedPackets;
+        } while (DroppedPacketsAfter - DroppedPacketsBefore != 1 && Tries++ < 10);
 
         //
         // Validate the server rejected the packet just sent.