Fix flaky test PivotBlockRetrieverTest.shouldRecoverFromUnresponsivePeer (#9408)

* Fix flaky test PivotBlockRetrieverTest.shouldRecoverFromUnresponsivePeer

Signed-off-by: Matilda Clerke <[email protected]>

* Fix flaky test PivotBlockRetrieverTest.shouldRecoverFromUnresponsivePeer again

Signed-off-by: Matilda Clerke <[email protected]>

* Move waitUntilComplete to bottom of file

Signed-off-by: Matilda Clerke <[email protected]>

---------

Signed-off-by: Matilda Clerke <[email protected]>

Author: Matilda-Clerke

ethereum/eth/src/test/java/org/hyperledger/besu/ethereum/eth/sync/fastsync/PivotBlockRetrieverTest.java

@@ -26,20 +26,20 @@
 import org.hyperledger.besu.ethereum.eth.manager.EthProtocolManager;
 import org.hyperledger.besu.ethereum.eth.manager.EthProtocolManagerTestBuilder;
 import org.hyperledger.besu.ethereum.eth.manager.EthProtocolManagerTestUtil;
+import org.hyperledger.besu.ethereum.eth.manager.EthScheduler;
 import org.hyperledger.besu.ethereum.eth.manager.peertask.PeerTaskExecutor;
 import org.hyperledger.besu.ethereum.eth.manager.peertask.PeerTaskExecutorResponseCode;
 import org.hyperledger.besu.ethereum.eth.manager.peertask.PeerTaskExecutorResult;
 import org.hyperledger.besu.ethereum.eth.manager.peertask.task.GetHeadersFromPeerTask;
 import org.hyperledger.besu.ethereum.eth.transactions.TransactionPool;
 import org.hyperledger.besu.ethereum.mainnet.ProtocolSchedule;
+import org.hyperledger.besu.metrics.noop.NoOpMetricsSystem;
 import org.hyperledger.besu.plugin.services.storage.DataStorageFormat;
-import org.hyperledger.besu.testutil.DeterministicEthScheduler;
 import org.hyperledger.besu.util.ExceptionUtils;
 
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Stream;
 
 import org.apache.tuweni.bytes.Bytes;
@@ -55,7 +55,6 @@ public class PivotBlockRetrieverTest {
 
   private static final long PIVOT_BLOCK_NUMBER = 10;
 
-  private final AtomicBoolean timeout = new AtomicBoolean(false);
   private EthProtocolManager ethProtocolManager;
   private MutableBlockchain blockchain;
   private TransactionPool transactionPool;
@@ -82,7 +81,7 @@ public void setUp(final DataStorageFormat storageFormat) {
         EthProtocolManagerTestBuilder.builder()
             .setProtocolSchedule(protocolSchedule)
             .setBlockchain(blockchain)
-            .setEthScheduler(new DeterministicEthScheduler(timeout::get))
+            .setEthScheduler(new EthScheduler(1, 1, 1, new NoOpMetricsSystem()))
             .setWorldStateArchive(blockchainSetupUtil.getWorldArchive())
             .setTransactionPool(transactionPool)
             .setEthereumWireProtocolConfiguration(EthProtocolConfiguration.defaultConfig())
@@ -143,6 +142,8 @@ public void shouldSucceedWhenAllPeersAgree(final DataStorageFormat storageFormat
 
     final CompletableFuture<FastSyncState> future = pivotBlockRetriever.downloadPivotBlockHeader();
 
+    waitUntilComplete(future);
+
     Mockito.verify(peerTaskExecutor)
         .executeAgainstPeer(Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA));
     Mockito.verify(peerTaskExecutor)
@@ -190,6 +191,8 @@ public void shouldQueryBestPeersFirst(final DataStorageFormat storageFormat) {
 
     final CompletableFuture<FastSyncState> future = pivotBlockRetriever.downloadPivotBlockHeader();
 
+    waitUntilComplete(future);
+
     Mockito.verify(peerTaskExecutor)
         .executeAgainstPeer(Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA));
     Mockito.verify(peerTaskExecutor, Mockito.never())
@@ -218,22 +221,24 @@ public void shouldRecoverFromUnresponsivePeer(final DataStorageFormat storageFor
     Mockito.when(
             peerTaskExecutor.executeAgainstPeer(
                 Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA)))
-        .thenReturn(
-            new PeerTaskExecutorResult<>(
-                Optional.of(List.of(blockchain.getBlockHeader(PIVOT_BLOCK_NUMBER).get())),
-                PeerTaskExecutorResponseCode.SUCCESS,
-                List.of(peerA)))
-        .thenReturn(
-            new PeerTaskExecutorResult<>(
-                Optional.of(List.of(blockchain.getBlockHeader(PIVOT_BLOCK_NUMBER).get())),
-                PeerTaskExecutorResponseCode.SUCCESS,
-                List.of(peerA)));
+        .thenAnswer(
+            (invocationOnMock) -> {
+              Thread.sleep(200);
+              return new PeerTaskExecutorResult<>(
+                  Optional.of(List.of(blockchain.getBlockHeader(PIVOT_BLOCK_NUMBER).get())),
+                  PeerTaskExecutorResponseCode.SUCCESS,
+                  List.of(peerA));
+            });
     Mockito.when(
             peerTaskExecutor.executeAgainstPeer(
                 Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerB)))
-        .thenReturn(
-            new PeerTaskExecutorResult<>(
-                Optional.empty(), PeerTaskExecutorResponseCode.TIMEOUT, List.of(peerB)));
+        .thenAnswer(
+            (invocationOnMock) -> {
+              Thread.sleep(200);
+              return new PeerTaskExecutorResult<>(
+                  Optional.empty(), PeerTaskExecutorResponseCode.TIMEOUT, List.of(peerB));
+            });
+
     Mockito.when(
             peerTaskExecutor.executeAgainstPeer(
                 Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerC)))
@@ -245,6 +250,8 @@ public void shouldRecoverFromUnresponsivePeer(final DataStorageFormat storageFor
 
     final CompletableFuture<FastSyncState> future = pivotBlockRetriever.downloadPivotBlockHeader();
 
+    waitUntilComplete(future);
+
     Mockito.verify(peerTaskExecutor)
         .executeAgainstPeer(Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA));
     Mockito.verify(peerTaskExecutor)
@@ -304,6 +311,8 @@ public void shouldRetryWhenPeersDisagreeOnPivot_successfulRetry(
     // Execute task and wait for response
     final CompletableFuture<FastSyncState> future = pivotBlockRetriever.downloadPivotBlockHeader();
 
+    waitUntilComplete(future);
+
     Mockito.verify(peerTaskExecutor, Mockito.times(2))
         .executeAgainstPeer(Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA));
     Mockito.verify(peerTaskExecutor, Mockito.times(2))
@@ -367,6 +376,8 @@ public void shouldRetryWhenPeersDisagreeOnPivot_exceedMaxRetries(
     // Execute task and wait for response
     final CompletableFuture<FastSyncState> future = pivotBlockRetriever.downloadPivotBlockHeader();
 
+    waitUntilComplete(future);
+
     Mockito.verify(peerTaskExecutor, Mockito.times(2))
         .executeAgainstPeer(Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA));
     Mockito.verify(peerTaskExecutor, Mockito.times(2))
@@ -417,6 +428,8 @@ public void shouldRetryWhenPeersDisagreeOnPivot_pivotInvalidOnRetry(
     // Execute task and wait for response
     final CompletableFuture<FastSyncState> future = pivotBlockRetriever.downloadPivotBlockHeader();
 
+    waitUntilComplete(future);
+
     Mockito.verify(peerTaskExecutor)
         .executeAgainstPeer(Mockito.any(GetHeadersFromPeerTask.class), Mockito.eq(peerA));
     Mockito.verify(peerTaskExecutor)
@@ -429,6 +442,14 @@ public void shouldRetryWhenPeersDisagreeOnPivot_pivotInvalidOnRetry(
         .isEqualTo(SyncError.PIVOT_BLOCK_HEADER_MISMATCH);
   }
 
+  private void waitUntilComplete(final CompletableFuture<?> future) {
+    try {
+      future.join();
+    } catch (RuntimeException e) {
+      // do nothing
+    }
+  }
+
   @Test
   void dryRunDetector() {
     assertThat(true)