From f4f49d552588d69063d0c8948d98e1adab48c2a4 Mon Sep 17 00:00:00 2001
From: Josh Lind <josh.lind@hotmail.com>
Date: Thu, 30 Nov 2023 17:54:25 -0500
Subject: [PATCH] [Forge] Bump catchup time for
 run-forge-realistic-env-workload-sweep.

---
 .../src/tests/storage_synchronizer.rs         | 20 +++++++++++++++----
 testsuite/forge-cli/src/main.rs               |  6 ++++--
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/state-sync/state-sync-driver/src/tests/storage_synchronizer.rs b/state-sync/state-sync-driver/src/tests/storage_synchronizer.rs
index 0b8a57e082f23..bb14888c5d5cd 100644
--- a/state-sync/state-sync-driver/src/tests/storage_synchronizer.rs
+++ b/state-sync/state-sync-driver/src/tests/storage_synchronizer.rs
@@ -38,7 +38,10 @@ use claims::assert_matches;
 use futures::StreamExt;
 use mockall::predicate::always;
 use std::{sync::Arc, time::Duration};
-use tokio::task::JoinHandle;
+use tokio::{task::JoinHandle, time::timeout};
+
+// Useful test constants
+const TEST_TIMEOUT_SECS: u64 = 30;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_apply_transaction_outputs() {
@@ -832,8 +835,12 @@ async fn verify_snapshot_commit_notification(
     commit_listener: &mut CommitNotificationListener,
     expected_committed_transactions: CommittedTransactions,
 ) {
-    let CommitNotification::CommittedStateSnapshot(committed_snapshot) =
-        commit_listener.select_next_some().await;
+    let CommitNotification::CommittedStateSnapshot(committed_snapshot) = timeout(
+        Duration::from_secs(TEST_TIMEOUT_SECS),
+        commit_listener.select_next_some(),
+    )
+    .await
+    .unwrap();
     assert_eq!(
         committed_snapshot.committed_transaction,
         expected_committed_transactions
@@ -845,7 +852,12 @@ async fn verify_error_notification(
     error_listener: &mut ErrorNotificationListener,
     expected_notification_id: NotificationId,
 ) {
-    let error_notification = error_listener.select_next_some().await;
+    let error_notification = timeout(
+        Duration::from_secs(TEST_TIMEOUT_SECS),
+        error_listener.select_next_some(),
+    )
+    .await
+    .unwrap();
     assert_eq!(error_notification.notification_id, expected_notification_id);
     assert_matches!(error_notification.error, Error::UnexpectedError(_));
 }
diff --git a/testsuite/forge-cli/src/main.rs b/testsuite/forge-cli/src/main.rs
index 790713d38d425..463e5d77bd026 100644
--- a/testsuite/forge-cli/src/main.rs
+++ b/testsuite/forge-cli/src/main.rs
@@ -1908,7 +1908,8 @@ fn realistic_network_tuned_for_throughput_test() -> ForgeConfig {
             .with_success_criteria(
                 SuccessCriteria::new(25000)
                     .add_no_restarts()
-                    .add_wait_for_catchup_s(60),
+                    /* This test runs at high load, so we need more catchup time */
+                    .add_wait_for_catchup_s(120),
                 /* Doesn't work without event indices
                 .add_chain_progress(StateProgressThreshold {
                     max_no_progress_secs: 10.0,
@@ -1920,7 +1921,8 @@ fn realistic_network_tuned_for_throughput_test() -> ForgeConfig {
         forge_config = forge_config.with_success_criteria(
             SuccessCriteria::new(12000)
                 .add_no_restarts()
-                .add_wait_for_catchup_s(60)
+                /* This test runs at high load, so we need more catchup time */
+                .add_wait_for_catchup_s(120)
                 .add_system_metrics_threshold(SystemMetricsThreshold::new(
                     // Tuned for throughput uses more cores than regular tests,
                     // as it achieves higher throughput.