From 9e57a3ef4304de4df575c8783307ba9bcfc39d21 Mon Sep 17 00:00:00 2001
From: Ayelet Zilber <ayelet@starkware.co>
Date: Mon, 1 Jul 2024 13:07:42 +0300
Subject: [PATCH] feat: add mock batcher and implement as a wrapper in
 integration test

---
 .../src/integration_test_setup.rs             | 27 ++++++-------------
 crates/tests-integration/src/lib.rs           |  1 +
 crates/tests-integration/src/mock_batcher.rs  | 23 ++++++++++++++++
 3 files changed, 32 insertions(+), 19 deletions(-)
 create mode 100644 crates/tests-integration/src/mock_batcher.rs

diff --git a/crates/tests-integration/src/integration_test_setup.rs b/crates/tests-integration/src/integration_test_setup.rs
index 4a589175..444dc500 100644
--- a/crates/tests-integration/src/integration_test_setup.rs
+++ b/crates/tests-integration/src/integration_test_setup.rs
@@ -8,9 +8,7 @@ use starknet_gateway::errors::GatewayError;
 use starknet_mempool::communication::create_mempool_server;
 use starknet_mempool::mempool::Mempool;
 use starknet_mempool_infra::component_server::ComponentServerStarter;
-use starknet_mempool_types::communication::{
-    MempoolClient, MempoolClientImpl, MempoolRequestAndResponseSender,
-};
+use starknet_mempool_types::communication::{MempoolClientImpl, MempoolRequestAndResponseSender};
 use starknet_mempool_types::mempool_types::ThinTransaction;
 use starknet_task_executor::executor::TaskExecutor;
 use starknet_task_executor::tokio_executor::TokioExecutor;
@@ -19,12 +17,13 @@ use tokio::sync::mpsc::channel;
 use tokio::task::JoinHandle;
 
 use crate::integration_test_utils::{create_gateway, GatewayClient};
+use crate::mock_batcher::MockBatcher;
 
 pub struct IntegrationTestSetup {
     pub task_executor: TokioExecutor,
     pub gateway_client: GatewayClient,
-    // TODO(MockBatcher).
-    pub batcher_mempool_client: MempoolClientImpl,
+
+    pub batcher: MockBatcher,
 
     pub gateway_handle: JoinHandle<()>,
     pub mempool_handle: JoinHandle<()>,
@@ -56,8 +55,7 @@ impl IntegrationTestSetup {
         tokio::time::sleep(std::time::Duration::from_millis(100)).await;
 
         // Build Batcher.
-        // TODO(MockBatcher)
-        let batcher_mempool_client = MempoolClientImpl::new(tx_mempool.clone());
+        let batcher = MockBatcher::new(tx_mempool.clone());
 
         // Build and run mempool.
         let mut mempool_server = create_mempool_server(Mempool::empty(), rx_mempool);
@@ -65,13 +63,7 @@ impl IntegrationTestSetup {
             mempool_server.start().await;
         });
 
-        Self {
-            task_executor,
-            gateway_client,
-            batcher_mempool_client,
-            gateway_handle,
-            mempool_handle,
-        }
+        Self { task_executor, gateway_client, batcher, gateway_handle, mempool_handle }
     }
 
     pub async fn assert_add_tx_success(&self, tx: &RPCTransaction) -> TransactionHash {
@@ -83,10 +75,7 @@ impl IntegrationTestSetup {
     }
 
     pub async fn get_txs(&mut self, n_txs: usize) -> Vec<ThinTransaction> {
-        let batcher_mempool_client = self.batcher_mempool_client.clone();
-        self.task_executor
-            .spawn(async move { batcher_mempool_client.get_txs(n_txs).await.unwrap() })
-            .await
-            .unwrap()
+        let batcher = self.batcher.clone();
+        self.task_executor.spawn(async move { batcher.get_txs(n_txs).await }).await.unwrap()
     }
 }
diff --git a/crates/tests-integration/src/lib.rs b/crates/tests-integration/src/lib.rs
index d0f04d51..265c50f2 100644
--- a/crates/tests-integration/src/lib.rs
+++ b/crates/tests-integration/src/lib.rs
@@ -1,3 +1,4 @@
 pub mod integration_test_setup;
 pub mod integration_test_utils;
+pub mod mock_batcher;
 pub mod state_reader;
diff --git a/crates/tests-integration/src/mock_batcher.rs b/crates/tests-integration/src/mock_batcher.rs
new file mode 100644
index 00000000..216501d5
--- /dev/null
+++ b/crates/tests-integration/src/mock_batcher.rs
@@ -0,0 +1,23 @@
+use starknet_mempool_infra::component_definitions::ComponentRequestAndResponseSender;
+use starknet_mempool_types::communication::{
+    MempoolClient, MempoolClientImpl, MempoolRequest, MempoolResponse,
+};
+use starknet_mempool_types::mempool_types::ThinTransaction;
+use tokio::sync::mpsc::Sender;
+
+#[derive(Clone)]
+pub struct MockBatcher {
+    mempool_client: MempoolClientImpl,
+}
+
+impl MockBatcher {
+    pub fn new(
+        mempool_client: Sender<ComponentRequestAndResponseSender<MempoolRequest, MempoolResponse>>,
+    ) -> Self {
+        Self { mempool_client: MempoolClientImpl::new(mempool_client) }
+    }
+
+    pub async fn get_txs(&self, n_txs: usize) -> Vec<ThinTransaction> {
+        self.mempool_client.get_txs(n_txs).await.unwrap()
+    }
+}