test(sync): in p2p sync client tests, simulate time has passed

crates/papyrus_p2p_sync/Cargo.toml

@@ -37,6 +37,7 @@ papyrus_storage = { workspace = true, features = ["testing"] }
 papyrus_test_utils.workspace = true
 rand_chacha.workspace = true
 static_assertions.workspace = true
+tokio = { workspace = true, features = ["test-util"] }
 
 # The `metrics` crate is used by `latency_histogram` proc macro, which is used in this crate.
 [package.metadata.cargo-machete]

crates/papyrus_p2p_sync/src/client/header_test.rs

@@ -19,6 +19,7 @@ use super::test_utils::{
     HEADER_QUERY_LENGTH,
     SLEEP_DURATION_TO_LET_SYNC_ADVANCE,
     TIMEOUT_FOR_NEW_QUERY_AFTER_PARTIAL_RESPONSE,
+    WAIT_PERIOD_FOR_NEW_DATA,
 };
 
 #[tokio::test]
@@ -146,6 +147,13 @@ async fn sync_sends_new_header_query_if_it_got_partial_responses() {
         }
         headers_sender.send(Ok(DataOrFin(None))).await.unwrap();
 
+        // Wait for the sync to enter sleep due to partial responses. Then, simulate time has
+        // passed.
+        tokio::time::sleep(SLEEP_DURATION_TO_LET_SYNC_ADVANCE).await;
+        tokio::time::pause();
+        tokio::time::advance(WAIT_PERIOD_FOR_NEW_DATA).await;
+        tokio::time::resume();
+
         // First unwrap is for the timeout. Second unwrap is for the Option returned from Stream.
         let SqmrClientPayload {
             query,

crates/papyrus_p2p_sync/src/client/state_diff_test.rs

@@ -32,6 +32,7 @@ use super::test_utils::{
     HEADER_QUERY_LENGTH,
     SLEEP_DURATION_TO_LET_SYNC_ADVANCE,
     STATE_DIFF_QUERY_LENGTH,
+    WAIT_PERIOD_FOR_NEW_DATA,
 };
 use super::{P2PSyncClientError, StateDiffQuery};
 
@@ -94,6 +95,16 @@ async fn state_diff_basic_flow() {
                 .await
                 .unwrap();
         }
+
+        // We wait for the header sync to write the new headers.
+        tokio::time::sleep(SLEEP_DURATION_TO_LET_SYNC_ADVANCE).await;
+
+        // Simulate time has passed so that state diff sync will resend query after it waited for
+        // new header
+        tokio::time::pause();
+        tokio::time::advance(WAIT_PERIOD_FOR_NEW_DATA).await;
+        tokio::time::resume();
+
         for (start_block_number, num_blocks) in [
             (0u64, STATE_DIFF_QUERY_LENGTH),
             (STATE_DIFF_QUERY_LENGTH, HEADER_QUERY_LENGTH - STATE_DIFF_QUERY_LENGTH),
@@ -111,7 +122,9 @@ async fn state_diff_basic_flow() {
                     direction: Direction::Forward,
                     limit: num_blocks,
                     step: 1,
-                })
+                }),
+                "If the limit of the query is too low, try to increase \
+                 SLEEP_DURATION_TO_LET_SYNC_ADVANCE",
             );
 
             for block_number in start_block_number..(start_block_number + num_blocks) {
@@ -344,6 +357,15 @@ async fn validate_state_diff_fails(
             .await
             .unwrap();
 
+        // We wait for the header sync to write the new headers.
+        tokio::time::sleep(SLEEP_DURATION_TO_LET_SYNC_ADVANCE).await;
+
+        // Simulate time has passed so that state diff sync will resend query after it waited for
+        // new header
+        tokio::time::pause();
+        tokio::time::advance(WAIT_PERIOD_FOR_NEW_DATA).await;
+        tokio::time::resume();
+
         // Get a state diff query and validate it
         let SqmrClientPayload {
             query,

crates/papyrus_p2p_sync/src/client/test_utils.rs

@@ -26,10 +26,9 @@ pub const HEADER_QUERY_LENGTH: u64 = 5;
 pub const STATE_DIFF_QUERY_LENGTH: u64 = 3;
 pub const TRANSACTION_QUERY_LENGTH: u64 = 3;
 pub const SLEEP_DURATION_TO_LET_SYNC_ADVANCE: Duration = Duration::from_millis(10);
-// This should be substantially bigger than SLEEP_DURATION_TO_LET_SYNC_ADVANCE.
-pub const WAIT_PERIOD_FOR_NEW_DATA: Duration = Duration::from_millis(50);
+pub const WAIT_PERIOD_FOR_NEW_DATA: Duration = Duration::from_secs(1);
 pub const TIMEOUT_FOR_NEW_QUERY_AFTER_PARTIAL_RESPONSE: Duration =
-    WAIT_PERIOD_FOR_NEW_DATA.saturating_add(SLEEP_DURATION_TO_LET_SYNC_ADVANCE.saturating_mul(10));
+    WAIT_PERIOD_FOR_NEW_DATA.saturating_add(Duration::from_secs(1));
 
 lazy_static! {
     static ref TEST_CONFIG: P2PSyncClientConfig = P2PSyncClientConfig {