From d1ad68790600d62407e1b13783c202e293ff3fcc Mon Sep 17 00:00:00 2001 From: dafnamatsry <92669167+dafnamatsry@users.noreply.github.com> Date: Mon, 2 Dec 2024 09:06:31 +0200 Subject: [PATCH] refactor(starknet_batcher): refactor batcher test mock dependencies (#2307) --- crates/starknet_batcher/src/batcher_test.rs | 202 +++++++++----------- 1 file changed, 93 insertions(+), 109 deletions(-) diff --git a/crates/starknet_batcher/src/batcher_test.rs b/crates/starknet_batcher/src/batcher_test.rs index c8a3b14379..79cab97a12 100644 --- a/crates/starknet_batcher/src/batcher_test.rs +++ b/crates/starknet_batcher/src/batcher_test.rs @@ -9,7 +9,7 @@ use futures::future::BoxFuture; use futures::FutureExt; use mockall::automock; use mockall::predicate::{always, eq}; -use rstest::{fixture, rstest}; +use rstest::rstest; use starknet_api::block::BlockNumber; use starknet_api::core::{ContractAddress, Nonce, StateDiffCommitment}; use starknet_api::executable_transaction::Transaction; @@ -72,64 +72,39 @@ fn deadline() -> chrono::DateTime { chrono::Utc::now() + BLOCK_GENERATION_TIMEOUT } -#[fixture] -fn storage_reader() -> MockBatcherStorageReaderTrait { - let mut storage = MockBatcherStorageReaderTrait::new(); - storage.expect_height().returning(|| Ok(INITIAL_HEIGHT)); - storage -} - -#[fixture] -fn storage_writer() -> MockBatcherStorageWriterTrait { - MockBatcherStorageWriterTrait::new() -} - -#[fixture] -fn batcher_config() -> BatcherConfig { - BatcherConfig { outstream_content_buffer_size: STREAMING_CHUNK_SIZE, ..Default::default() } +struct MockDependencies { + storage_reader: MockBatcherStorageReaderTrait, + storage_writer: MockBatcherStorageWriterTrait, + mempool_client: MockMempoolClient, + proposal_manager: MockProposalManagerTraitWrapper, + block_builder_factory: MockBlockBuilderFactoryTrait, } -#[fixture] -fn mempool_client() -> MockMempoolClient { - MockMempoolClient::new() +impl Default for MockDependencies { + fn default() -> Self { + let mut storage_reader = MockBatcherStorageReaderTrait::new(); + storage_reader.expect_height().returning(|| Ok(INITIAL_HEIGHT)); + Self { + storage_reader, + storage_writer: MockBatcherStorageWriterTrait::new(), + mempool_client: MockMempoolClient::new(), + proposal_manager: MockProposalManagerTraitWrapper::new(), + block_builder_factory: MockBlockBuilderFactoryTrait::new(), + } + } } -fn batcher(proposal_manager: MockProposalManagerTraitWrapper) -> Batcher { +fn create_batcher(mock_dependencies: MockDependencies) -> Batcher { Batcher::new( - batcher_config(), - Arc::new(storage_reader()), - Box::new(storage_writer()), - Arc::new(mempool_client()), - Box::new(MockBlockBuilderFactoryTrait::new()), - Box::new(proposal_manager), + BatcherConfig { outstream_content_buffer_size: STREAMING_CHUNK_SIZE, ..Default::default() }, + Arc::new(mock_dependencies.storage_reader), + Box::new(mock_dependencies.storage_writer), + Arc::new(mock_dependencies.mempool_client), + Box::new(mock_dependencies.block_builder_factory), + Box::new(mock_dependencies.proposal_manager), ) } -fn create_batcher( - proposal_manager: MockProposalManagerTraitWrapper, - block_builder_factory: MockBlockBuilderFactoryTrait, -) -> Batcher { - Batcher::new( - batcher_config(), - Arc::new(storage_reader()), - Box::new(storage_writer()), - Arc::new(mempool_client()), - Box::new(block_builder_factory), - Box::new(proposal_manager), - ) -} - -fn mock_proposal_manager_common_expectations( - proposal_manager: &mut MockProposalManagerTraitWrapper, -) { - proposal_manager.expect_wrap_reset().times(1).return_once(|| async {}.boxed()); - proposal_manager - .expect_wrap_await_proposal_commitment() - .times(1) - .with(eq(PROPOSAL_ID)) - .return_once(move |_| { async move { Ok(proposal_commitment()) } }.boxed()); -} - fn abort_signal_sender() -> AbortSignalSender { tokio::sync::oneshot::channel().0 } @@ -168,6 +143,17 @@ fn mock_create_builder_for_propose_block( block_builder_factory } +fn mock_proposal_manager_common_expectations( + proposal_manager: &mut MockProposalManagerTraitWrapper, +) { + proposal_manager.expect_wrap_reset().times(1).return_once(|| async {}.boxed()); + proposal_manager + .expect_wrap_await_proposal_commitment() + .times(1) + .with(eq(PROPOSAL_ID)) + .return_once(move |_| { async move { Ok(proposal_commitment()) } }.boxed()); +} + fn mock_proposal_manager_validate_flow() -> MockProposalManagerTraitWrapper { let mut proposal_manager = MockProposalManagerTraitWrapper::new(); mock_proposal_manager_common_expectations(&mut proposal_manager); @@ -190,7 +176,7 @@ async fn start_height_success() { let mut proposal_manager = MockProposalManagerTraitWrapper::new(); proposal_manager.expect_wrap_reset().times(1).return_once(|| async {}.boxed()); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); assert_eq!(batcher.start_height(StartHeightInput { height: INITIAL_HEIGHT }).await, Ok(())); } @@ -214,7 +200,7 @@ async fn start_height_fail(#[case] height: BlockNumber, #[case] expected_error: let mut proposal_manager = MockProposalManagerTraitWrapper::new(); proposal_manager.expect_wrap_reset().never(); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); assert_eq!(batcher.start_height(StartHeightInput { height }).await, Err(expected_error)); } @@ -224,7 +210,7 @@ async fn duplicate_start_height() { let mut proposal_manager = MockProposalManagerTraitWrapper::new(); proposal_manager.expect_wrap_reset().times(1).return_once(|| async {}.boxed()); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); let initial_height = StartHeightInput { height: INITIAL_HEIGHT }; assert_eq!(batcher.start_height(initial_height.clone()).await, Ok(())); @@ -235,7 +221,7 @@ async fn duplicate_start_height() { #[tokio::test] async fn no_active_height() { let proposal_manager = MockProposalManagerTraitWrapper::new(); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); // Calling `propose_block` and `validate_block` without starting a height should fail. @@ -263,7 +249,11 @@ async fn no_active_height() { async fn validate_block_full_flow() { let block_builder_factory = mock_create_builder_for_validate_block(); let proposal_manager = mock_proposal_manager_validate_flow(); - let mut batcher = create_batcher(proposal_manager, block_builder_factory); + let mut batcher = create_batcher(MockDependencies { + proposal_manager, + block_builder_factory, + ..Default::default() + }); batcher.start_height(StartHeightInput { height: INITIAL_HEIGHT }).await.unwrap(); @@ -304,7 +294,7 @@ async fn send_content_after_proposal_already_finished() { .times(1) .returning(|_| async move { InternalProposalStatus::Finished }.boxed()); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); // Send transactions after the proposal has finished. let send_proposal_input_txs = SendProposalContentInput { @@ -325,7 +315,7 @@ async fn send_content_to_unknown_proposal() { .with(eq(PROPOSAL_ID)) .return_once(move |_| async move { InternalProposalStatus::NotFound }.boxed()); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); // Send transactions to an unknown proposal. let send_proposal_input_txs = SendProposalContentInput { @@ -352,7 +342,7 @@ async fn send_txs_to_an_invalid_proposal() { .with(eq(PROPOSAL_ID)) .return_once(move |_| async move { InternalProposalStatus::Failed }.boxed()); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); let send_proposal_input_txs = SendProposalContentInput { proposal_id: PROPOSAL_ID, @@ -383,7 +373,11 @@ async fn send_finish_to_an_invalid_proposal() { .with(eq(PROPOSAL_ID)) .return_once(move |_| { async move { Err(proposal_error) } }.boxed()); - let mut batcher = create_batcher(proposal_manager, block_builder_factory); + let mut batcher = create_batcher(MockDependencies { + proposal_manager, + block_builder_factory, + ..Default::default() + }); batcher.start_height(StartHeightInput { height: INITIAL_HEIGHT }).await.unwrap(); let validate_block_input = ValidateBlockInput { @@ -414,7 +408,11 @@ async fn propose_block_full_flow() { .times(1) .return_once(|_, _, _| { async move { Ok(()) } }.boxed()); - let mut batcher = create_batcher(proposal_manager, block_builder_factory); + let mut batcher = create_batcher(MockDependencies { + proposal_manager, + block_builder_factory, + ..Default::default() + }); batcher.start_height(StartHeightInput { height: INITIAL_HEIGHT }).await.unwrap(); batcher @@ -454,6 +452,7 @@ async fn propose_block_full_flow() { assert_matches!(exhausted, Err(BatcherError::ProposalNotFound { .. })); } +#[rstest] #[tokio::test] async fn propose_block_without_retrospective_block_hash() { let mut proposal_manager = MockProposalManagerTraitWrapper::new(); @@ -464,14 +463,8 @@ async fn propose_block_without_retrospective_block_hash() { .expect_height() .returning(|| Ok(BlockNumber(constants::STORED_BLOCK_HASH_BUFFER))); - let mut batcher = Batcher::new( - batcher_config(), - Arc::new(storage_reader), - Box::new(storage_writer()), - Arc::new(mempool_client()), - Box::new(MockBlockBuilderFactoryTrait::new()), - Box::new(proposal_manager), - ); + let mut batcher = + create_batcher(MockDependencies { proposal_manager, storage_reader, ..Default::default() }); batcher .start_height(StartHeightInput { height: BlockNumber(constants::STORED_BLOCK_HASH_BUFFER) }) @@ -494,7 +487,7 @@ async fn get_content_from_unknown_proposal() { let mut proposal_manager = MockProposalManagerTraitWrapper::new(); proposal_manager.expect_wrap_await_proposal_commitment().times(0); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); let get_proposal_content_input = GetProposalContentInput { proposal_id: PROPOSAL_ID }; let result = batcher.get_proposal_content(get_proposal_content_input).await; @@ -503,52 +496,43 @@ async fn get_content_from_unknown_proposal() { #[rstest] #[tokio::test] -async fn decision_reached( - batcher_config: BatcherConfig, - storage_reader: MockBatcherStorageReaderTrait, - mut storage_writer: MockBatcherStorageWriterTrait, - mut mempool_client: MockMempoolClient, -) { - let expected_state_diff = ThinStateDiff::default(); - let state_diff_clone = expected_state_diff.clone(); - let expected_proposal_commitment = ProposalCommitment::default(); - let tx_hashes = test_tx_hashes(0..5); - let tx_hashes_clone = tx_hashes.clone(); - let address_to_nonce = test_contract_nonces(0..3); - let nonces_clone = address_to_nonce.clone(); +async fn decision_reached() { + let mut mock_dependencies = MockDependencies::default(); - let mut proposal_manager = MockProposalManagerTraitWrapper::new(); - proposal_manager.expect_wrap_take_proposal_result().times(1).with(eq(PROPOSAL_ID)).return_once( - move |_| { + mock_dependencies + .proposal_manager + .expect_wrap_take_proposal_result() + .times(1) + .with(eq(PROPOSAL_ID)) + .return_once(move |_| { async move { Ok(ProposalOutput { - state_diff: state_diff_clone, - commitment: expected_proposal_commitment, - tx_hashes: tx_hashes_clone, - nonces: nonces_clone, + state_diff: ThinStateDiff::default(), + commitment: ProposalCommitment::default(), + tx_hashes: test_tx_hashes(), + nonces: test_contract_nonces(), }) } .boxed() - }, - ); - mempool_client + }); + + mock_dependencies + .mempool_client .expect_commit_block() - .with(eq(CommitBlockArgs { address_to_nonce, tx_hashes })) + .with(eq(CommitBlockArgs { + address_to_nonce: test_contract_nonces(), + tx_hashes: test_tx_hashes(), + })) .returning(|_| Ok(())); - storage_writer + mock_dependencies + .storage_writer .expect_commit_proposal() - .with(eq(INITIAL_HEIGHT), eq(expected_state_diff)) + .with(eq(INITIAL_HEIGHT), eq(ThinStateDiff::default())) .returning(|_, _| Ok(())); - let mut batcher = Batcher::new( - batcher_config, - Arc::new(storage_reader), - Box::new(storage_writer), - Arc::new(mempool_client), - Box::new(MockBlockBuilderFactoryTrait::new()), - Box::new(proposal_manager), - ); + let mut batcher = create_batcher(mock_dependencies); + batcher.decision_reached(DecisionReachedInput { proposal_id: PROPOSAL_ID }).await.unwrap(); } @@ -564,7 +548,7 @@ async fn decision_reached_no_executed_proposal() { }, ); - let mut batcher = batcher(proposal_manager); + let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); let decision_reached_result = batcher.decision_reached(DecisionReachedInput { proposal_id: PROPOSAL_ID }).await; assert_eq!(decision_reached_result, Err(expected_error)); @@ -638,10 +622,10 @@ impl ProposalManagerTrait for T { } } -fn test_tx_hashes(range: std::ops::Range) -> HashSet { - range.map(|i| tx_hash!(i)).collect() +fn test_tx_hashes() -> HashSet { + (0..5u8).map(|i| tx_hash!(i + 12)).collect() } -fn test_contract_nonces(range: std::ops::Range) -> HashMap { - HashMap::from_iter(range.map(|i| (contract_address!(i), nonce!(i)))) +fn test_contract_nonces() -> HashMap { + HashMap::from_iter((0..3u8).map(|i| (contract_address!(i + 33), nonce!(i + 9)))) }