feat: add MultiAccountTransactionGenerator (#252)

Based on https://reviewable.io/reviews/starkware-libs/blockifier/1277
but with multiple acocunts and internally managed nonce.
(So far only supporting Invoke). @dorimedini-starkware @ArniStarkware

commit-id:eda74c35

Co-authored-by: Gilad Chase <[email protected]>

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;

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<AccountId, FeatureContract>,
+    // 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<Item = FeatureContract>) -> 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);

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)
+}