test(mempool_p2p): fix CR comments

crates/starknet_integration_tests/tests/mempool_p2p_flow_test.rs

@@ -6,7 +6,7 @@ use futures::StreamExt;
 use mempool_test_utils::starknet_api_test_utils::MultiAccountTransactionGenerator;
 use papyrus_network::gossipsub_impl::Topic;
 use papyrus_network::network_manager::test_utils::create_network_config_connected_to_broadcast_channels;
-use papyrus_network::network_manager::BroadcastTopicClientTrait;
+use papyrus_network::network_manager::{BroadcastTopicChannels, BroadcastTopicClientTrait};
 use papyrus_protobuf::mempool::RpcTransactionWrapper;
 use rstest::{fixture, rstest};
 use starknet_api::executable_transaction::AccountTransaction;
@@ -46,15 +46,11 @@ fn tx_generator() -> MultiAccountTransactionGenerator {
     create_integration_test_tx_generator()
 }
 
-// TODO: remove code duplication with FlowTestSetup
-#[rstest]
-#[tokio::test]
-async fn test_mempool_sends_tx_to_other_peer(tx_generator: MultiAccountTransactionGenerator) {
-    let handle = Handle::current();
-    let task_executor = TokioExecutor::new(handle);
-
-    let chain_info = create_chain_info();
+async fn setup(
+    tx_generator: &MultiAccountTransactionGenerator,
+) -> (SequencerNodeConfig, BroadcastTopicChannels<RpcTransactionWrapper>) {
     let accounts = tx_generator.accounts();
+    let chain_info = create_chain_info();
     let storage_for_test = StorageTestSetup::new(accounts, chain_info.chain_id.clone());
 
     // Spawn a papyrus rpc server for a papyrus storage reader.
@@ -79,12 +75,13 @@ async fn test_mempool_sends_tx_to_other_peer(tx_generator: MultiAccountTransacti
         ..Default::default()
     };
 
+    let chain_info = create_chain_info();
     let batcher_config =
         create_batcher_config(storage_for_test.batcher_storage_config, chain_info.clone());
     let gateway_config = create_gateway_config(chain_info).await;
     let http_server_config = create_http_server_config().await;
     let rpc_state_reader_config = test_rpc_state_reader_config(rpc_server_addr);
-    let (network_config, mut broadcast_channels) =
+    let (network_config, broadcast_channels) =
         create_network_config_connected_to_broadcast_channels::<RpcTransactionWrapper>(Topic::new(
             MEMPOOL_TOPIC,
         ));
@@ -98,7 +95,15 @@ async fn test_mempool_sends_tx_to_other_peer(tx_generator: MultiAccountTransacti
         mempool_p2p_config,
         ..SequencerNodeConfig::default()
     };
+    (config, broadcast_channels)
+}
 
+#[rstest]
+#[tokio::test]
+async fn test_mempool_sends_tx_to_other_peer(tx_generator: MultiAccountTransactionGenerator) {
+    let handle = Handle::current();
+    let task_executor = TokioExecutor::new(handle);
+    let (config, mut broadcast_channels) = setup(&tx_generator).await;
     let (_clients, servers) = create_node_modules(&config);
 
     let HttpServerConfig { ip, port } = config.http_server_config;
@@ -135,48 +140,7 @@ async fn test_mempool_sends_tx_to_other_peer(tx_generator: MultiAccountTransacti
 async fn test_mempool_receives_tx_from_other_peer(tx_generator: MultiAccountTransactionGenerator) {
     let handle = Handle::current();
     let task_executor = TokioExecutor::new(handle);
-    let accounts = tx_generator.accounts();
-    let storage_for_test = StorageTestSetup::new(accounts);
-    // Spawn a papyrus rpc server for a papyrus storage reader.
-    let rpc_server_addr =
-        spawn_test_rpc_state_reader(storage_for_test.rpc_storage_reader, storage_for_test.chain_id)
-            .await;
-    // Derive the configuration for the mempool node.
-    let components = ComponentConfig {
-        consensus_manager: ComponentExecutionConfig {
-            execution_mode: ComponentExecutionMode::Disabled,
-            local_server_config: None,
-            ..Default::default()
-        },
-        batcher: ComponentExecutionConfig {
-            execution_mode: ComponentExecutionMode::Disabled,
-            local_server_config: None,
-            ..Default::default()
-        },
-        ..Default::default()
-    };
-    let chain_id = storage_for_test.batcher_storage_config.db_config.chain_id.clone();
-    let mut chain_info = ChainInfo::create_for_testing();
-    chain_info.chain_id = chain_id.clone();
-    let batcher_config =
-        create_batcher_config(storage_for_test.batcher_storage_config, chain_info.clone());
-    let gateway_config = create_gateway_config(chain_info).await;
-    let http_server_config = create_http_server_config().await;
-    let rpc_state_reader_config = test_rpc_state_reader_config(rpc_server_addr);
-    let (network_config, mut broadcast_channels) =
-        create_network_config_connected_to_broadcast_channels::<RpcTransactionWrapper>(Topic::new(
-            MEMPOOL_TOPIC,
-        ));
-    let mempool_p2p_config = MempoolP2pConfig { network_config, ..Default::default() };
-    let config = SequencerNodeConfig {
-        components,
-        batcher_config,
-        gateway_config,
-        http_server_config,
-        rpc_state_reader_config,
-        mempool_p2p_config,
-        ..SequencerNodeConfig::default()
-    };
+    let (config, mut broadcast_channels) = setup(&tx_generator).await;
     let (clients, servers) = create_node_modules(&config);
     let mempool_client = clients.get_mempool_shared_client().unwrap();
     // Build and run the sequencer node.
@@ -202,15 +166,19 @@ async fn test_mempool_receives_tx_from_other_peer(tx_generator: MultiAccountTran
             .unwrap();
     }
 
-    // waiting for the tx to be received (TODO: figure out a better solution)
-    tokio::time::sleep(std::time::Duration::from_millis(2000)).await;
-
-    let received_tx = mempool_client.get_txs(expected_txs.len()).await.unwrap();
-
-    // make sure we received all the transactions
-    assert_eq!(received_tx.len(), expected_txs.len());
+    let mut received_transactions: Vec<AccountTransaction> = vec![];
+    for _ in 0..20 {
+        if received_transactions.len() == expected_txs.len() {
+            break;
+        }
+        received_transactions
+            .append(&mut mempool_client.get_txs(expected_txs.len()).await.unwrap());
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+    }
+    assert_eq!(received_transactions.len(), expected_txs.len());
 
-    for tx in received_tx.iter() {
+    for tx in received_transactions {
+        // TODO: change mempool to store RpcTransaction
         let converted_tx: RpcTransaction = match tx {
             AccountTransaction::Declare(_declare_tx) => {
                 panic!("No implementation for converting DeclareTransaction to an RpcTransaction")