Skip to content

Commit

Permalink
feat: add mock batcher and implement as a wrapper in integration test
Browse files Browse the repository at this point in the history
  • Loading branch information
ayeletstarkware committed Jul 2, 2024
1 parent a7c19f4 commit 4786460
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 19 deletions.
1 change: 1 addition & 0 deletions crates/tests-integration/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
26 changes: 7 additions & 19 deletions crates/tests-integration/src/integration_test_setup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<()>,
Expand Down Expand Up @@ -55,22 +53,15 @@ 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);
let mempool_handle = task_executor.spawn_with_handle(async move {
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 {
Expand All @@ -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()
}
}
1 change: 1 addition & 0 deletions crates/tests-integration/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
pub mod integration_test_setup;
pub mod integration_test_utils;
pub mod mock_batcher;
pub mod state_reader;
23 changes: 23 additions & 0 deletions crates/tests-integration/src/mock_batcher.rs
Original file line number Diff line number Diff line change
@@ -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()
}
}

0 comments on commit 4786460

Please sign in to comment.