From cbcaa88548c3ce53ea83013f5b5a6726f3172291 Mon Sep 17 00:00:00 2001 From: Yair Bakalchuk Date: Thu, 26 Dec 2024 15:09:08 +0200 Subject: [PATCH] feat(mempool_test_utils,starknet_integration_tests): prepare funded but undeployed accounts --- .../src/starknet_api_test_utils.rs | 6 +-- .../src/flow_test_setup.rs | 11 ++-- .../src/integration_test_setup.rs | 9 ++-- .../src/state_reader.rs | 54 +++++++++++++++---- .../tests/mempool_p2p_flow_test.rs | 2 +- 5 files changed, 61 insertions(+), 21 deletions(-) diff --git a/crates/mempool_test_utils/src/starknet_api_test_utils.rs b/crates/mempool_test_utils/src/starknet_api_test_utils.rs index 6c1b8303abd..c08b2fef25a 100644 --- a/crates/mempool_test_utils/src/starknet_api_test_utils.rs +++ b/crates/mempool_test_utils/src/starknet_api_test_utils.rs @@ -260,8 +260,8 @@ impl MultiAccountTransactionGenerator { }) } - pub fn accounts(&self) -> Vec { - self.account_tx_generators.iter().map(|tx_gen| &tx_gen.account).copied().collect() + pub fn accounts(&self) -> &[AccountTransactionGenerator] { + self.account_tx_generators.as_slice() } } @@ -272,7 +272,7 @@ impl MultiAccountTransactionGenerator { /// with room for future extensions. /// /// TODO: add more transaction generation methods as needed. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct AccountTransactionGenerator { pub account: Contract, nonce_manager: SharedNonceManager, diff --git a/crates/starknet_integration_tests/src/flow_test_setup.rs b/crates/starknet_integration_tests/src/flow_test_setup.rs index 707fb5bbd72..9389e1b8922 100644 --- a/crates/starknet_integration_tests/src/flow_test_setup.rs +++ b/crates/starknet_integration_tests/src/flow_test_setup.rs @@ -1,7 +1,10 @@ use std::net::SocketAddr; use blockifier::context::ChainInfo; -use mempool_test_utils::starknet_api_test_utils::{Contract, MultiAccountTransactionGenerator}; +use mempool_test_utils::starknet_api_test_utils::{ + AccountTransactionGenerator, + MultiAccountTransactionGenerator, +}; use papyrus_network::network_manager::BroadcastTopicChannels; use papyrus_protobuf::consensus::{ProposalPart, StreamMessage}; use starknet_api::rpc_transaction::RpcTransaction; @@ -66,7 +69,7 @@ impl FlowTestSetup { // Create nodes one after the other in order to make sure the ports are not overlapping. let sequencer_0 = FlowSequencerSetup::new( - accounts.clone(), + accounts.to_vec(), SEQUENCER_0, chain_info.clone(), sequencer_0_consensus_manager_config, @@ -76,7 +79,7 @@ impl FlowTestSetup { .await; let sequencer_1 = FlowSequencerSetup::new( - accounts, + accounts.to_vec(), SEQUENCER_1, chain_info, sequencer_1_consensus_manager_config, @@ -115,7 +118,7 @@ pub struct FlowSequencerSetup { impl FlowSequencerSetup { #[instrument(skip(accounts, chain_info, consensus_manager_config), level = "debug")] pub async fn new( - accounts: Vec, + accounts: Vec, sequencer_index: usize, chain_info: ChainInfo, consensus_manager_config: ConsensusManagerConfig, diff --git a/crates/starknet_integration_tests/src/integration_test_setup.rs b/crates/starknet_integration_tests/src/integration_test_setup.rs index 736ff7674fc..406b58e5f57 100644 --- a/crates/starknet_integration_tests/src/integration_test_setup.rs +++ b/crates/starknet_integration_tests/src/integration_test_setup.rs @@ -2,7 +2,10 @@ use std::net::SocketAddr; use std::path::PathBuf; use blockifier::context::ChainInfo; -use mempool_test_utils::starknet_api_test_utils::{Contract, MultiAccountTransactionGenerator}; +use mempool_test_utils::starknet_api_test_utils::{ + AccountTransactionGenerator, + MultiAccountTransactionGenerator, +}; use papyrus_network::network_manager::BroadcastTopicChannels; use papyrus_protobuf::consensus::{ProposalPart, StreamMessage}; use papyrus_storage::{StorageConfig, StorageReader}; @@ -62,7 +65,7 @@ impl IntegrationTestSetup { let consensus_manager_config = consensus_manager_configs.remove(0); let mempool_p2p_config = mempool_p2p_configs.remove(0); let sequencer = IntegrationSequencerSetup::new( - accounts.clone(), + accounts.to_vec(), sequencer_id, chain_info.clone(), consensus_manager_config, @@ -143,7 +146,7 @@ pub struct IntegrationSequencerSetup { impl IntegrationSequencerSetup { #[instrument(skip(accounts, chain_info, consensus_manager_config), level = "debug")] pub async fn new( - accounts: Vec, + accounts: Vec, sequencer_index: usize, chain_info: ChainInfo, consensus_manager_config: ConsensusManagerConfig, diff --git a/crates/starknet_integration_tests/src/state_reader.rs b/crates/starknet_integration_tests/src/state_reader.rs index e75ea8cd0c2..a4306823434 100644 --- a/crates/starknet_integration_tests/src/state_reader.rs +++ b/crates/starknet_integration_tests/src/state_reader.rs @@ -8,7 +8,7 @@ use blockifier::test_utils::{CairoVersion, RunnableCairo1, BALANCE}; use blockifier::versioned_constants::VersionedConstants; use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass; use indexmap::IndexMap; -use mempool_test_utils::starknet_api_test_utils::Contract; +use mempool_test_utils::starknet_api_test_utils::{AccountTransactionGenerator, Contract}; use papyrus_common::pending_classes::PendingClasses; use papyrus_rpc::{run_server, RpcConfig}; use papyrus_storage::body::BodyStorageWriter; @@ -60,7 +60,10 @@ pub struct StorageTestSetup { } impl StorageTestSetup { - pub fn new(test_defined_accounts: Vec, chain_info: &ChainInfo) -> Self { + pub fn new( + test_defined_accounts: Vec, + chain_info: &ChainInfo, + ) -> Self { let ((rpc_storage_reader, mut rpc_storage_writer), _, rpc_storage_file_handle) = TestStorageBuilder::default().chain_id(chain_info.chain_id.clone()).build(); create_test_state(&mut rpc_storage_writer, chain_info, test_defined_accounts.clone()); @@ -94,7 +97,7 @@ impl StorageTestSetup { fn create_test_state( storage_writer: &mut StorageWriter, chain_info: &ChainInfo, - test_defined_accounts: Vec, + test_defined_accounts: Vec, ) { let into_contract = |contract: FeatureContract| Contract { contract, @@ -123,7 +126,7 @@ fn create_test_state( fn initialize_papyrus_test_state( storage_writer: &mut StorageWriter, chain_info: &ChainInfo, - test_defined_accounts: Vec, + test_defined_accounts: Vec, default_test_contracts: Vec, erc20_contract: Contract, ) { @@ -134,8 +137,11 @@ fn initialize_papyrus_test_state( &erc20_contract, ); - let contract_classes_to_retrieve = - test_defined_accounts.into_iter().chain(default_test_contracts).chain([erc20_contract]); + let contract_classes_to_retrieve = test_defined_accounts + .into_iter() + .map(|acc| acc.account) + .chain(default_test_contracts) + .chain([erc20_contract]); let sierra_vec: Vec<_> = prepare_sierra_classes(contract_classes_to_retrieve.clone()); let (cairo0_contract_classes, cairo1_contract_classes) = prepare_compiled_contract_classes(contract_classes_to_retrieve); @@ -151,7 +157,7 @@ fn initialize_papyrus_test_state( fn prepare_state_diff( chain_info: &ChainInfo, - test_defined_accounts: &[Contract], + test_defined_accounts: &[AccountTransactionGenerator], default_test_contracts: &[Contract], erc20_contract: &Contract, ) -> ThinStateDiff { @@ -169,7 +175,17 @@ fn prepare_state_diff( // state_diff_builder.set_contracts(accounts_defined_in_the_test).declare().fund(); // ``` // or use declare txs and transfers for both. - state_diff_builder.inject_accounts_into_state(test_defined_accounts); + let (deployed_accounts, undeployed_accounts): (Vec<_>, Vec<_>) = + test_defined_accounts.iter().partition(|account| account.is_deployed()); + + let deployed_accounts_contracts: Vec<_> = + deployed_accounts.iter().map(|acc| acc.account).collect(); + let undeployed_accounts_contracts: Vec<_> = + undeployed_accounts.iter().map(|acc| acc.account).collect(); + + state_diff_builder.inject_deployed_accounts_into_state(deployed_accounts_contracts.as_slice()); + state_diff_builder + .inject_undeployed_accounts_into_state(undeployed_accounts_contracts.as_slice()); state_diff_builder.build() } @@ -394,8 +410,26 @@ impl<'a> ThinStateDiffBuilder<'a> { } // TODO(deploy_account_support): delete method once we have batcher with execution. - fn inject_accounts_into_state(&mut self, accounts_defined_in_the_test: &'a [Contract]) { - self.set_contracts(accounts_defined_in_the_test).declare().deploy().fund(); + fn inject_deployed_accounts_into_state( + &mut self, + deployed_accounts_defined_in_the_test: &'a [Contract], + ) { + self.set_contracts(deployed_accounts_defined_in_the_test).declare().deploy().fund(); + + // Set nonces as 1 in the state so that subsequent invokes can pass validation. + self.nonces = self + .deployed_contracts + .iter() + .map(|(&address, _)| (address, Nonce(Felt::ONE))) + .collect(); + } + + // TODO(deploy_account_support): delete method once we have batcher with execution. + fn inject_undeployed_accounts_into_state( + &mut self, + undeployed_accounts_defined_in_the_test: &'a [Contract], + ) { + self.set_contracts(undeployed_accounts_defined_in_the_test).declare().fund(); // Set nonces as 1 in the state so that subsequent invokes can pass validation. self.nonces = self diff --git a/crates/starknet_integration_tests/tests/mempool_p2p_flow_test.rs b/crates/starknet_integration_tests/tests/mempool_p2p_flow_test.rs index c69b39b0561..1f275552704 100644 --- a/crates/starknet_integration_tests/tests/mempool_p2p_flow_test.rs +++ b/crates/starknet_integration_tests/tests/mempool_p2p_flow_test.rs @@ -57,7 +57,7 @@ async fn setup( ) -> (SequencerNodeConfig, BroadcastTopicChannels) { let accounts = tx_generator.accounts(); let chain_info = create_chain_info(); - let storage_for_test = StorageTestSetup::new(accounts, &chain_info); + let storage_for_test = StorageTestSetup::new(accounts.to_vec(), &chain_info); let mut available_ports = AvailablePorts::new(test_identifier.into(), 0); // Spawn a papyrus rpc server for a papyrus storage reader.