diff --git a/crates/gateway/src/stateful_transaction_validator_test.rs b/crates/gateway/src/stateful_transaction_validator_test.rs index 57787e58..2800c553 100644 --- a/crates/gateway/src/stateful_transaction_validator_test.rs +++ b/crates/gateway/src/stateful_transaction_validator_test.rs @@ -7,6 +7,7 @@ use mempool_test_utils::starknet_api_test_utils::{ VALID_L1_GAS_MAX_PRICE_PER_UNIT, }; use num_bigint::BigUint; +use pretty_assertions::assert_eq; use rstest::rstest; use starknet_api::felt; use starknet_api::rpc_transaction::RPCTransaction; 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 956c31fb..6217bab9 100644 --- a/crates/mempool_test_utils/src/starknet_api_test_utils.rs +++ b/crates/mempool_test_utils/src/starknet_api_test_utils.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::env; use std::fs::File; use std::path::Path; @@ -107,6 +108,121 @@ pub fn declare_tx() -> RPCTransaction { )) } +// Convenience method for generating a single invoke transaction with trivial fields. +// For multiple, nonce-incrementing transactions, use the transaction generator directly. +pub fn invoke_tx(cairo_version: CairoVersion) -> RPCTransaction { + let default_account = FeatureContract::AccountWithoutValidations(cairo_version); + + MultiAccountTransactionGenerator::new_for_account_contracts([default_account]) + .account_with_id(0) + .generate_default() +} + +// TODO: when moving this to Starknet API crate, move this const into a module alongside +// MultiAcconutTransactionGenerator. +type AccountId = u16; + +/// Manages transaction generation for multiple pre-funded accounts, internally bumping nonces +/// as needed. +/// +/// **Currently supports:** +/// - Single contract type +/// - Only supports invokes, which are all a trivial method in the contract type. +/// +/// # Example +/// +/// ```ignore +/// use starknet_gateway::invoke_tx_args; +/// use starknet_gateway::starknet_api_test_utils::MultiAccountTransactionGenerator; +/// +/// let mut tx_generator = MultiAccountTransactionGenerator::new(2); // Initialize with 2 accounts. +/// let account_0_tx_with_nonce_0 = tx_generator.account(0).generate_default(); +/// let account_1_tx_with_nonce_0 = tx_generator.account(1).generate_default(); +/// let account_0_tx_with_nonce_1 = tx_generator.account(0).generate_default(); +/// ``` +// Note: when moving this to starknet api crate, see if blockifier's +// [blockifier::transaction::test_utils::FaultyAccountTxCreatorArgs] can be made to use this. +pub struct MultiAccountTransactionGenerator { + // Invariant: coupled with nonce_manager. + account_contracts: HashMap, + // Invariant: nonces managed internally thorugh `generate` API. + nonce_manager: NonceManager, +} + +impl MultiAccountTransactionGenerator { + pub fn new(n_accounts: usize) -> Self { + let default_account_contract = + FeatureContract::AccountWithoutValidations(CairoVersion::Cairo1); + let accounts = std::iter::repeat(default_account_contract).take(n_accounts); + Self::new_for_account_contracts(accounts) + } + + pub fn new_for_account_contracts(accounts: impl IntoIterator) -> Self { + let enumerated_accounts = (0..).zip(accounts); + let account_contracts = enumerated_accounts.collect(); + + Self { account_contracts, nonce_manager: NonceManager::default() } + } + + pub fn account_with_id(&mut self, account_id: AccountId) -> AccountTransactionGenerator<'_> { + AccountTransactionGenerator { account_id, generator: self } + } +} + +/// Manages transaction generation for a single account. +/// Supports faulty transaction generation via [AccountTransactionGenerator::generate_raw]. +/// +/// This struct provides methods to generate both default and fully customized transactions, +/// with room for future extensions. +/// +/// TODO: add more transaction generation methods as needed. +pub struct AccountTransactionGenerator<'a> { + account_id: AccountId, + generator: &'a mut MultiAccountTransactionGenerator, +} + +impl<'a> AccountTransactionGenerator<'a> { + /// Generate a valid `RPCTransaction` with default parameters. + pub fn generate_default(&mut self) -> RPCTransaction { + let invoke_args = invoke_tx_args!( + sender_address: self.sender_address(), + resource_bounds: executable_resource_bounds_mapping(), + nonce: self.next_nonce(), + calldata: create_trivial_calldata(self.test_contract_address()), + ); + external_invoke_tx(invoke_args) + } + + // TODO: support more contracts, instead of this hardcoded type. + pub fn test_contract_address(&mut self) -> ContractAddress { + let cairo_version = self.generator.account_contracts[&self.account_id].cairo_version(); + FeatureContract::TestContract(cairo_version).get_instance_address(self.account_id) + } + + /// Generates an `RPCTransaction` with fully custom parameters. + /// + /// Caller must manually handle bumping nonce and fetching the correct sender address via + /// [AccountTransactionGenerator::nonce] and [AccountTransactionGenerator::sender_address]. + /// See [AccountTransactionGenerator::generate_default] to have these filled up by default. + /// + /// Note: This is a best effort attempt to make the API more useful; amend or add new methods + /// as needed. + pub fn generate_raw(&mut self, invoke_tx_args: InvokeTxArgs) -> RPCTransaction { + external_invoke_tx(invoke_tx_args) + } + + pub fn sender_address(&mut self) -> ContractAddress { + let account_id = self.account_id; + self.generator.account_contracts[&account_id].get_instance_address(account_id) + } + + /// Retrieves the nonce for the current account, and __increments__ it internally. + pub fn next_nonce(&mut self) -> Nonce { + let sender_address = self.sender_address(); + self.generator.nonce_manager.next(sender_address) + } +} + // TODO(Ayelet, 28/5/2025): Try unifying the macros. // TODO(Ayelet, 28/5/2025): Consider moving the macros StarkNet API. #[macro_export] @@ -346,20 +462,6 @@ pub fn external_tx_to_json(tx: &RPCTransaction) -> String { to_string_pretty(&tx_json).expect("Failed to serialize transaction") } -pub fn invoke_tx(cairo_version: CairoVersion) -> RPCTransaction { - let test_contract = FeatureContract::TestContract(cairo_version); - let account_contract = FeatureContract::AccountWithoutValidations(cairo_version); - let sender_address = account_contract.get_instance_address(0); - let mut nonce_manager = NonceManager::default(); - - external_invoke_tx(invoke_tx_args!( - resource_bounds: executable_resource_bounds_mapping(), - nonce : nonce_manager.next(sender_address), - sender_address, - calldata: create_trivial_calldata(test_contract.get_instance_address(0)) - )) -} - // TODO(Yael 18/6/2024): Get a final decision from product whether to support Cairo0. pub fn deploy_account_tx() -> RPCTransaction { let account_contract = FeatureContract::AccountWithoutValidations(CairoVersion::Cairo1); diff --git a/crates/tests-integration/src/integration_test_utils.rs b/crates/tests-integration/src/integration_test_utils.rs index 11c1f10d..58f9eeb9 100644 --- a/crates/tests-integration/src/integration_test_utils.rs +++ b/crates/tests-integration/src/integration_test_utils.rs @@ -1,7 +1,9 @@ use std::net::SocketAddr; use axum::body::Body; -use mempool_test_utils::starknet_api_test_utils::external_tx_to_json; +use mempool_test_utils::starknet_api_test_utils::{ + external_tx_to_json, MultiAccountTransactionGenerator, +}; use reqwest::{Client, Response}; use starknet_api::rpc_transaction::RPCTransaction; use starknet_api::transaction::TransactionHash; @@ -13,6 +15,8 @@ use starknet_gateway::errors::GatewayError; use starknet_mempool_node::config::MempoolNodeConfig; use tokio::net::TcpListener; +use crate::integration_test_setup::IntegrationTestSetup; + async fn create_gateway_config() -> GatewayConfig { let stateless_tx_validator_config = StatelessTransactionValidatorConfig { validate_non_zero_l1_gas_fee: true, @@ -95,3 +99,13 @@ pub async fn get_available_socket() -> SocketAddr { .local_addr() .expect("Failed to get local address") } + +/// Use to create a tx generator with _pre-funded_ accounts, alongside a mocked test setup. +pub async fn setup_with_tx_generation( + n_accounts: usize, +) -> (IntegrationTestSetup, MultiAccountTransactionGenerator) { + let integration_test_setup = IntegrationTestSetup::new(n_accounts).await; + let tx_generator = MultiAccountTransactionGenerator::new(n_accounts); + + (integration_test_setup, tx_generator) +}