From 4786460027d21d6dbd23e82d7ffe6ba13896798a 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

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

diff --git a/crates/tests-integration/Cargo.toml b/crates/tests-integration/Cargo.toml
index b6e36b7d..0060d088 100644
--- a/crates/tests-integration/Cargo.toml
+++ b/crates/tests-integration/Cargo.toml
@@ -23,6 +23,7 @@ starknet_api.workspace = true
 starknet_client.workspace = true
 starknet_gateway = { path = "../gateway", version = "0.0", features = ["testing"] }
 starknet_mempool = { path = "../mempool", version = "0.0" }
+starknet_mempool_infra = { path = "../mempool_infra" }
 starknet_mempool_types = { path = "../mempool_types", version = "0.0" }
 starknet_task_executor = { path = "../task_executor", version = "0.0" }
 strum.workspace = true
diff --git a/crates/tests-integration/src/integration_test_setup.rs b/crates/tests-integration/src/integration_test_setup.rs
index b51db583..3dddabe8 100644
--- a/crates/tests-integration/src/integration_test_setup.rs
+++ b/crates/tests-integration/src/integration_test_setup.rs
@@ -7,9 +7,7 @@ use starknet_gateway::config::GatewayNetworkConfig;
 use starknet_gateway::errors::GatewayError;
 use starknet_mempool::communication::create_mempool_server;
 use starknet_mempool::mempool::Mempool;
-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;
@@ -18,12 +16,12 @@ 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<()>,
@@ -55,8 +53,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);
@@ -64,13 +61,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 {
@@ -82,10 +73,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 mock_batcher = self.batcher.clone();
+        self.task_executor.spawn(async move { mock_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..7a3f4a7a
--- /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 {
+        MockBatcher { 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()
+    }
+}