Skip to content

Commit

Permalink
refactor(mempool_test_utils): add non mut account_with_id, rename pre…
Browse files Browse the repository at this point in the history
…vious implementation
  • Loading branch information
yair-starkware committed Dec 17, 2024
1 parent e02ed63 commit da0ded7
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 11 deletions.
24 changes: 18 additions & 6 deletions crates/mempool_test_utils/src/starknet_api_test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ pub fn executable_invoke_tx(cairo_version: CairoVersion) -> AccountTransaction {

let mut tx_generator = MultiAccountTransactionGenerator::new();
tx_generator.register_account(default_account);
tx_generator.account_with_id(0).generate_executable_invoke()
tx_generator.account_with_id_mut(0).generate_executable_invoke()
}

pub fn generate_deploy_account_with_salt(
Expand Down Expand Up @@ -162,9 +162,9 @@ type SharedNonceManager = Rc<RefCell<NonceManager>>;
/// tx_generator.register_account_for_flow_test(some_account_type.clone());
/// tx_generator.register_account_for_flow_test(some_account_type);
///
/// let account_0_tx_with_nonce_0 = tx_generator.account_with_id(0).generate_invoke_with_tip(1);
/// let account_1_tx_with_nonce_0 = tx_generator.account_with_id(1).generate_invoke_with_tip(3);
/// let account_0_tx_with_nonce_1 = tx_generator.account_with_id(0).generate_invoke_with_tip(1);
/// let account_0_tx_with_nonce_0 = tx_generator.account_with_id_mut(0).generate_invoke_with_tip(1);
/// let account_1_tx_with_nonce_0 = tx_generator.account_with_id_mut(1).generate_invoke_with_tip(3);
/// let account_0_tx_with_nonce_1 = tx_generator.account_with_id_mut(0).generate_invoke_with_tip(1);
/// ```
// Note: when moving this to starknet api crate, see if blockifier's
// [blockifier::transaction::test_utils::FaultyAccountTxCreatorArgs] can be made to use this.
Expand Down Expand Up @@ -194,7 +194,10 @@ impl MultiAccountTransactionGenerator {
default_deploy_account_tx
}

pub fn account_with_id(&mut self, account_id: AccountId) -> &mut AccountTransactionGenerator {
pub fn account_with_id_mut(
&mut self,
account_id: AccountId,
) -> &mut AccountTransactionGenerator {
self.account_tx_generators.get_mut(account_id).unwrap_or_else(|| {
panic!(
"{account_id:?} not found! This number should be an index of an account in the \
Expand All @@ -203,6 +206,15 @@ impl MultiAccountTransactionGenerator {
})
}

pub fn account_with_id(&self, account_id: AccountId) -> &AccountTransactionGenerator {
self.account_tx_generators.get(account_id).unwrap_or_else(|| {
panic!(
"{account_id:?} not found! This number should be an index of an account in the \
initialization array. "
)
})
}

// TODO(deploy_account_support): once we support deploy account in tests, remove this method and
// only use new_account_default in tests. In particular, deploy account txs must be then sent to
// the GW via the add tx endpoint just like other txs.
Expand Down Expand Up @@ -280,7 +292,7 @@ impl AccountTransactionGenerator {
rpc_invoke_tx(invoke_tx_args)
}

pub fn sender_address(&mut self) -> ContractAddress {
pub fn sender_address(&self) -> ContractAddress {
self.account.sender_address
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ async fn await_block(
.ok_or(())
}

pub async fn end_to_end_integration(mut tx_generator: MultiAccountTransactionGenerator) {
pub async fn end_to_end_integration(tx_generator: MultiAccountTransactionGenerator) {
const EXPECTED_BLOCK_NUMBER: BlockNumber = BlockNumber(15);

info!("Running integration test setup.");
Expand Down
8 changes: 4 additions & 4 deletions crates/starknet_integration_tests/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -165,11 +165,11 @@ pub fn create_txs_for_integration_test(

// Create RPC transactions.
let account0_invoke_nonce1 =
tx_generator.account_with_id(ACCOUNT_ID_0).generate_invoke_with_tip(2);
tx_generator.account_with_id_mut(ACCOUNT_ID_0).generate_invoke_with_tip(2);
let account0_invoke_nonce2 =
tx_generator.account_with_id(ACCOUNT_ID_0).generate_invoke_with_tip(3);
tx_generator.account_with_id_mut(ACCOUNT_ID_0).generate_invoke_with_tip(3);
let account1_invoke_nonce1 =
tx_generator.account_with_id(ACCOUNT_ID_1).generate_invoke_with_tip(4);
tx_generator.account_with_id_mut(ACCOUNT_ID_1).generate_invoke_with_tip(4);

vec![account0_invoke_nonce1, account0_invoke_nonce2, account1_invoke_nonce1]
}
Expand All @@ -180,7 +180,7 @@ fn create_account_txs(
n_txs: usize,
) -> Vec<RpcTransaction> {
(0..n_txs)
.map(|_| tx_generator.account_with_id(account_id).generate_invoke_with_tip(1))
.map(|_| tx_generator.account_with_id_mut(account_id).generate_invoke_with_tip(1))
.collect()
}

Expand Down

0 comments on commit da0ded7

Please sign in to comment.