From 17d5d3b91b6165693d22f901410ce05005f55a99 Mon Sep 17 00:00:00 2001
From: ShahakShama <70578257+ShahakShama@users.noreply.github.com>
Date: Thu, 5 Dec 2024 13:47:53 +0200
Subject: [PATCH] refactor(papyrus_p2p_sync): add_test receives query size
 instead of constant (#2379)

---
 .../src/client/header_test.rs                 | 37 +++++++++-------
 .../papyrus_p2p_sync/src/client/test_utils.rs | 44 ++++++++++++++-----
 2 files changed, 52 insertions(+), 29 deletions(-)

diff --git a/crates/papyrus_p2p_sync/src/client/header_test.rs b/crates/papyrus_p2p_sync/src/client/header_test.rs
index ad214edbef..c0c706eb1c 100644
--- a/crates/papyrus_p2p_sync/src/client/header_test.rs
+++ b/crates/papyrus_p2p_sync/src/client/header_test.rs
@@ -198,23 +198,26 @@ async fn sync_sends_new_header_query_if_it_got_partial_responses() {
 
 #[tokio::test]
 async fn wrong_block_number() {
-    run_test(vec![
-        // We already validate the query content in other tests.
-        Action::ReceiveQuery(Box::new(|_query| ())),
-        Action::SendHeader(DataOrFin(Some(random_header(
-            &mut get_rng(),
-            BlockNumber(1),
-            None,
-            None,
-        )))),
-        Action::ValidateReportSent,
-        Action::CheckStorage(Box::new(|reader| {
-            async move {
-                assert_eq!(0, reader.begin_ro_txn().unwrap().get_header_marker().unwrap().0);
-            }
-            .boxed()
-        })),
-    ])
+    run_test(
+        1,
+        vec![
+            // We already validate the query content in other tests.
+            Action::ReceiveQuery(Box::new(|_query| ())),
+            Action::SendHeader(DataOrFin(Some(random_header(
+                &mut get_rng(),
+                BlockNumber(1),
+                None,
+                None,
+            )))),
+            Action::ValidateReportSent,
+            Action::CheckStorage(Box::new(|reader| {
+                async move {
+                    assert_eq!(0, reader.begin_ro_txn().unwrap().get_header_marker().unwrap().0);
+                }
+                .boxed()
+            })),
+        ],
+    )
     .await;
 }
 
diff --git a/crates/papyrus_p2p_sync/src/client/test_utils.rs b/crates/papyrus_p2p_sync/src/client/test_utils.rs
index eaf4d87682..d2c2861fd8 100644
--- a/crates/papyrus_p2p_sync/src/client/test_utils.rs
+++ b/crates/papyrus_p2p_sync/src/client/test_utils.rs
@@ -135,18 +135,38 @@ pub enum Action {
 }
 
 // TODO(shahak): add support for state diffs, transactions and classes.
-// TODO(shahak): add max query size as argument instead of constant.
-pub async fn run_test(actions: Vec<Action>) {
-    let TestArgs {
-        p2p_sync,
-        storage_reader,
-        mut mock_header_response_manager,
-        // The test will fail if we drop these
-        mock_state_diff_response_manager: _mock_state_diff_responses_manager,
-        mock_class_response_manager: _mock_class_responses_manager,
-        mock_transaction_response_manager: _mock_transaction_responses_manager,
-        ..
-    } = setup();
+pub async fn run_test(header_max_query_length: u64, actions: Vec<Action>) {
+    let p2p_sync_config = P2PSyncClientConfig {
+        num_headers_per_query: header_max_query_length,
+        num_block_state_diffs_per_query: STATE_DIFF_QUERY_LENGTH,
+        num_block_transactions_per_query: TRANSACTION_QUERY_LENGTH,
+        num_block_classes_per_query: CLASS_DIFF_QUERY_LENGTH,
+        wait_period_for_new_data: WAIT_PERIOD_FOR_NEW_DATA,
+        buffer_size: BUFFER_SIZE,
+        stop_sync_at_block_number: None,
+    };
+    let buffer_size = p2p_sync_config.buffer_size;
+    let ((storage_reader, storage_writer), _temp_dir) = get_test_storage();
+    let (header_sender, mut mock_header_response_manager) =
+        mock_register_sqmr_protocol_client(buffer_size);
+    let (state_diff_sender, _mock_state_diff_response_manager) =
+        mock_register_sqmr_protocol_client(buffer_size);
+    let (transaction_sender, _mock_transaction_response_manager) =
+        mock_register_sqmr_protocol_client(buffer_size);
+    let (class_sender, _mock_class_response_manager) =
+        mock_register_sqmr_protocol_client(buffer_size);
+    let p2p_sync_channels = P2PSyncClientChannels {
+        header_sender,
+        state_diff_sender,
+        transaction_sender,
+        class_sender,
+    };
+    let p2p_sync = P2PSyncClient::new(
+        p2p_sync_config,
+        storage_reader.clone(),
+        storage_writer,
+        p2p_sync_channels,
+    );
 
     let mut headers_current_query_responses_manager = None;