From 158e523d173f4973de53b698888789adcd90c52b Mon Sep 17 00:00:00 2001
From: Ayelet Zilber <ayelet@starkware.co>
Date: Thu, 4 Jul 2024 16:47:10 +0300
Subject: [PATCH] refactor: mock batcher channel to unbounded channel

---
 crates/tests-integration/src/integration_test_setup.rs | 8 ++++----
 crates/tests-integration/src/mock_batcher.rs           | 6 +++---
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/crates/tests-integration/src/integration_test_setup.rs b/crates/tests-integration/src/integration_test_setup.rs
index 8ca5843e..e83dfb3e 100644
--- a/crates/tests-integration/src/integration_test_setup.rs
+++ b/crates/tests-integration/src/integration_test_setup.rs
@@ -12,7 +12,7 @@ use starknet_mempool_types::communication::{MempoolClientImpl, MempoolRequestAnd
 use starknet_mempool_types::mempool_types::ThinTransaction;
 use starknet_task_executor::tokio_executor::TokioExecutor;
 use tokio::runtime::Handle;
-use tokio::sync::mpsc::{channel, Sender};
+use tokio::sync::mpsc::{channel, unbounded_channel, UnboundedSender};
 use tokio::sync::oneshot;
 use tokio::task::JoinHandle;
 
@@ -24,7 +24,7 @@ pub struct IntegrationTestSetup {
     pub gateway_client: GatewayClient,
 
     pub batcher_handle: JoinHandle<()>,
-    tx_batcher: Sender<BatcherCommand>,
+    tx_batcher: UnboundedSender<BatcherCommand>,
 
     pub gateway_handle: JoinHandle<()>,
     pub mempool_handle: JoinHandle<()>,
@@ -56,7 +56,7 @@ impl IntegrationTestSetup {
         tokio::time::sleep(std::time::Duration::from_millis(100)).await;
 
         // Build Batcher.
-        let (tx_batcher, rx_batcher) = channel::<BatcherCommand>(MESSAGE_QUEUE_SIZE);
+        let (tx_batcher, rx_batcher) = unbounded_channel::<BatcherCommand>();
         let mut batcher = MockBatcher::new(rx_batcher, tx_mempool);
         let batcher_handle = task_executor.spawn_with_handle(async move {
             batcher.run().await;
@@ -88,7 +88,7 @@ impl IntegrationTestSetup {
 
     pub async fn get_txs(&self, n_txs: usize) -> Vec<ThinTransaction> {
         let (response_tx, response_rx) = oneshot::channel::<Vec<ThinTransaction>>();
-        self.tx_batcher.send(BatcherCommand::GetTxs(n_txs, response_tx)).await.unwrap();
+        self.tx_batcher.send(BatcherCommand::GetTxs(n_txs, response_tx)).unwrap();
         response_rx.await.unwrap()
     }
 }
diff --git a/crates/tests-integration/src/mock_batcher.rs b/crates/tests-integration/src/mock_batcher.rs
index c5ce2f5c..316d73b6 100644
--- a/crates/tests-integration/src/mock_batcher.rs
+++ b/crates/tests-integration/src/mock_batcher.rs
@@ -3,18 +3,18 @@ use starknet_mempool_types::communication::{
     MempoolClient, MempoolClientImpl, MempoolRequest, MempoolResponse,
 };
 use starknet_mempool_types::mempool_types::ThinTransaction;
-use tokio::sync::mpsc::{Receiver, Sender};
+use tokio::sync::mpsc::{Sender, UnboundedReceiver};
 
 use crate::integration_test_utils::BatcherCommand;
 
 pub struct MockBatcher {
-    rx_commands: Receiver<BatcherCommand>,
+    rx_commands: UnboundedReceiver<BatcherCommand>,
     mempool_client: MempoolClientImpl,
 }
 
 impl MockBatcher {
     pub fn new(
-        rx_commands: Receiver<BatcherCommand>,
+        rx_commands: UnboundedReceiver<BatcherCommand>,
         mempool_sender: Sender<ComponentRequestAndResponseSender<MempoolRequest, MempoolResponse>>,
     ) -> Self {
         Self { rx_commands, mempool_client: MempoolClientImpl::new(mempool_sender) }