fix(papyrus_p2p_sync): Tx and state_diff tests wait for headers sync

starkware-libs · Feb 2, 2025 · 585542f · 585542f
1 parent e6d9b3f
commit 585542f
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/crates/papyrus_p2p_sync/src/client/state_diff_test.rs b/crates/papyrus_p2p_sync/src/client/state_diff_test.rs
@@ -103,6 +103,8 @@ async fn state_diff_basic_flow() {
         Action::ReceiveQuery(Box::new(|_query| ()), DataType::Header),
     ];
 
+    // Sleep so state diff sync will wait for new data
+    actions.push(Action::SleepToLetSyncAdvance);
     // Send headers with corresponding state diff length
     for (i, (state_diff, _)) in state_diffs_and_chunks.iter().enumerate() {
         actions.push(Action::SendHeader(DataOrFin(Some(random_header(
@@ -114,6 +116,24 @@ async fn state_diff_basic_flow() {
     }
     actions.push(Action::SendHeader(DataOrFin(None)));
 
+    let len = state_diffs_and_chunks.len();
+    // Wait for header sync to finish before continuing state diff sync.
+    actions.push(Action::CheckStorage(Box::new(move |reader| {
+        async move {
+            let block_number = BlockNumber(len.try_into().unwrap());
+            wait_for_marker(
+                DataType::Header,
+                &reader,
+                block_number,
+                SLEEP_DURATION_TO_LET_SYNC_ADVANCE,
+                TIMEOUT_FOR_TEST,
+            )
+            .await;
+        }
+        .boxed()
+    })));
+    actions.push(Action::SimulateWaitPeriodForNewData);
+
     let len = state_diffs_and_chunks.len();
     actions.push(Action::ReceiveQuery(
         Box::new(move |query| {

diff --git a/crates/papyrus_p2p_sync/src/client/transaction_test.rs b/crates/papyrus_p2p_sync/src/client/transaction_test.rs
@@ -44,6 +44,8 @@ async fn transaction_basic_flow() {
         Action::ReceiveQuery(Box::new(|_query| ()), DataType::Header),
     ];
 
+    // Sleep so transaction sync will wait for new data
+    actions.push(Action::SleepToLetSyncAdvance);
     // Send headers with corresponding transaction length
     for (i, block_body) in block_bodies.iter().enumerate() {
         actions.push(Action::SendHeader(DataOrFin(Some(random_header(
@@ -55,6 +57,24 @@ async fn transaction_basic_flow() {
     }
     actions.push(Action::SendHeader(DataOrFin(None)));
 
+    let len = block_bodies.len();
+    // Wait for header sync to finish before continuing transaction sync.
+    actions.push(Action::CheckStorage(Box::new(move |reader| {
+        async move {
+            let block_number = BlockNumber(len.try_into().unwrap());
+            wait_for_marker(
+                DataType::Header,
+                &reader,
+                block_number,
+                SLEEP_DURATION_TO_LET_SYNC_ADVANCE,
+                TIMEOUT_FOR_TEST,
+            )
+            .await;
+        }
+        .boxed()
+    })));
+    actions.push(Action::SimulateWaitPeriodForNewData);
+
     // Send transactions for each block and then validate they were written
     for (i, BlockBody { transactions, transaction_outputs, transaction_hashes }) in
         block_bodies.into_iter().enumerate()