Add more address balance transfer helpers to test tx builder

Author: lxfind

crates/sui-core/src/unit_tests/coin_deny_list_tests.rs

@@ -10,7 +10,7 @@ use move_core_types::identifier::Identifier;
 use move_core_types::language_storage::{StructTag, TypeTag};
 use std::sync::Arc;
 use sui_protocol_config::{Chain, PerObjectCongestionControlMode, ProtocolConfig, ProtocolVersion};
-use sui_test_transaction_builder::TestTransactionBuilder;
+use sui_test_transaction_builder::{FundSource, TestTransactionBuilder};
 use sui_types::base_types::{ObjectID, ObjectRef, SuiAddress, dbg_addr};
 use sui_types::crypto::{AccountKeyPair, get_account_key_pair};
 use sui_types::deny_list_v1::{CoinDenyCap, RegulatedCoinMetadata};
@@ -238,7 +238,11 @@ async fn test_regulated_coin_v2_funds_withdraw_deny() {
             env_gas_ref,
             env.authority.reference_gas_price_for_testing().unwrap(),
         )
-        .transfer_funds_to_address_balance(100_000_000, regulated_coin_type.clone(), denied_address)
+        .transfer_funds_to_address_balance(
+            FundSource::address_fund(),
+            vec![(100_000_000, denied_address)],
+            regulated_coin_type.clone(),
+        )
         .build_and_sign(&env.keypair);
         let effects = send_and_confirm_transaction_(&env.authority, None, tx, true)
             .await

crates/sui-e2e-tests/tests/address_balance_tests.rs

@@ -9,13 +9,13 @@ use sui_keys::keystore::AccountKeystore;
 use sui_macros::*;
 use sui_protocol_config::{ProtocolConfig, ProtocolVersion};
 use sui_sdk::wallet_context::WalletContext;
-use sui_test_transaction_builder::TestTransactionBuilder;
+use sui_test_transaction_builder::{FundSource, TestTransactionBuilder};
 use sui_types::{
     SUI_ACCUMULATOR_ROOT_OBJECT_ID, SUI_FRAMEWORK_PACKAGE_ID,
     accumulator_metadata::AccumulatorOwner,
     accumulator_root::{AccumulatorValue, U128},
     balance::Balance,
-    base_types::{ObjectID, ObjectRef, SuiAddress},
+    base_types::{ObjectID, ObjectRef, SuiAddress, dbg_addr},
     digests::{ChainIdentifier, CheckpointDigest},
     effects::{InputConsensusObject, TransactionEffectsAPI},
     gas::GasCostSummary,
@@ -364,7 +364,7 @@ async fn test_multiple_settlement_txns() {
         .map(|_| (1u64, SuiAddress::random_for_testing_only()))
         .collect::<Vec<_>>();
 
-    let tx = make_send_to_multi_account_tx(&amounts_and_recipients, sender, gas, rgp);
+    let tx = make_send_to_multi_account_tx(amounts_and_recipients.clone(), sender, gas, rgp);
 
     let res = test_cluster.sign_and_execute_transaction(&tx).await;
     let gas = res.effects.unwrap().gas_object().reference.to_object_ref();
@@ -644,37 +644,12 @@ fn withdraw_from_balance_tx_with_reservation(
     gas: ObjectRef,
     rgp: u64,
 ) -> TransactionData {
-    let mut builder = ProgrammableTransactionBuilder::new();
-
-    // Add withdraw reservation
-    let withdraw_arg = sui_types::transaction::FundsWithdrawalArg::balance_from_sender(
-        reservation_amount,
-        sui_types::type_input::TypeInput::from(sui_types::gas_coin::GAS::type_tag()),
-    );
-    let withdraw_arg = builder.funds_withdrawal(withdraw_arg).unwrap();
-
-    let amount_arg = builder.pure(U256::from(amount)).unwrap();
-
-    let split_withdraw_arg = builder.programmable_move_call(
-        SUI_FRAMEWORK_PACKAGE_ID,
-        Identifier::new("funds_accumulator").unwrap(),
-        Identifier::new("withdrawal_split").unwrap(),
-        vec!["0x2::balance::Balance<0x2::sui::SUI>".parse().unwrap()],
-        vec![withdraw_arg, amount_arg],
-    );
-
-    let coin = builder.programmable_move_call(
-        SUI_FRAMEWORK_PACKAGE_ID,
-        Identifier::new("coin").unwrap(),
-        Identifier::new("redeem_funds").unwrap(),
-        vec!["0x2::sui::SUI".parse().unwrap()],
-        vec![split_withdraw_arg],
-    );
-
-    builder.transfer_arg(sender, coin);
-
-    let tx = TransactionKind::ProgrammableTransaction(builder.finish());
-    TransactionData::new(tx, sender, gas, 10000000, rgp)
+    TestTransactionBuilder::new(sender, gas, rgp)
+        .transfer_sui_to_address_balance(
+            FundSource::address_fund_with_reservation(reservation_amount),
+            vec![(amount, dbg_addr(2))],
+        )
+        .build()
 }
 
 fn make_send_to_account_tx(
@@ -684,39 +659,18 @@ fn make_send_to_account_tx(
     gas: ObjectRef,
     rgp: u64,
 ) -> TransactionData {
-    make_send_to_multi_account_tx(&[(amount, recipient)], sender, gas, rgp)
+    make_send_to_multi_account_tx(vec![(amount, recipient)], sender, gas, rgp)
 }
 
 fn make_send_to_multi_account_tx(
-    amounts_and_recipients: &[(u64, SuiAddress)],
+    amounts_and_recipients: Vec<(u64, SuiAddress)>,
     sender: SuiAddress,
     gas: ObjectRef,
     rgp: u64,
 ) -> TransactionData {
-    let mut builder = ProgrammableTransactionBuilder::new();
-
-    for (amount, recipient) in amounts_and_recipients {
-        let amount_arg = builder.pure(*amount).unwrap();
-        let recipient_arg = builder.pure(recipient).unwrap();
-        let coin = builder.command(Command::SplitCoins(Argument::GasCoin, vec![amount_arg]));
-
-        let Argument::Result(coin_idx) = coin else {
-            panic!("coin is not a result");
-        };
-
-        let coin = Argument::NestedResult(coin_idx, 0);
-
-        builder.programmable_move_call(
-            SUI_FRAMEWORK_PACKAGE_ID,
-            Identifier::new("coin").unwrap(),
-            Identifier::new("send_funds").unwrap(),
-            vec!["0x2::sui::SUI".parse().unwrap()],
-            vec![coin, recipient_arg],
-        );
-    }
-
-    let tx = TransactionKind::ProgrammableTransaction(builder.finish());
-    TransactionData::new(tx, sender, gas, 10000000, rgp)
+    TestTransactionBuilder::new(sender, gas, rgp)
+        .transfer_sui_to_address_balance(FundSource::coin(gas), amounts_and_recipients)
+        .build()
 }
 
 #[sim_test]

crates/sui-test-transaction-builder/src/lib.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use move_core_types::ident_str;
+use move_core_types::u256::U256;
 use shared_crypto::intent::{Intent, IntentMessage};
 use std::path::PathBuf;
 use sui_genesis_builder::validator_info::GenesisValidatorMetadata;
@@ -10,23 +11,52 @@ use sui_sdk::rpc_types::{
     SuiObjectDataOptions, SuiTransactionBlockEffectsAPI, SuiTransactionBlockResponse,
 };
 use sui_sdk::wallet_context::WalletContext;
-use sui_types::SUI_RANDOMNESS_STATE_OBJECT_ID;
+use sui_types::balance::Balance;
 use sui_types::base_types::{FullObjectRef, ObjectID, ObjectRef, SequenceNumber, SuiAddress};
 use sui_types::crypto::{AccountKeyPair, Signature, Signer, get_key_pair};
 use sui_types::digests::TransactionDigest;
 use sui_types::gas_coin::GAS;
 use sui_types::multisig::{BitmapUnit, MultiSig, MultiSigPublicKey};
 use sui_types::multisig_legacy::{MultiSigLegacy, MultiSigPublicKeyLegacy};
 use sui_types::object::Owner;
+use sui_types::programmable_transaction_builder::ProgrammableTransactionBuilder;
 use sui_types::signature::GenericSignature;
 use sui_types::sui_system_state::SUI_SYSTEM_MODULE_NAME;
 use sui_types::transaction::{
-    CallArg, DEFAULT_VALIDATOR_GAS_PRICE, ObjectArg, ProgrammableTransaction,
-    SharedObjectMutability, TEST_ONLY_GAS_UNIT_FOR_HEAVY_COMPUTATION_STORAGE,
-    TEST_ONLY_GAS_UNIT_FOR_TRANSFER, Transaction, TransactionData,
+    Argument, CallArg, DEFAULT_VALIDATOR_GAS_PRICE, FundsWithdrawalArg, ObjectArg,
+    ProgrammableTransaction, SharedObjectMutability,
+    TEST_ONLY_GAS_UNIT_FOR_HEAVY_COMPUTATION_STORAGE, TEST_ONLY_GAS_UNIT_FOR_TRANSFER, Transaction,
+    TransactionData,
 };
+use sui_types::{Identifier, SUI_FRAMEWORK_PACKAGE_ID, SUI_RANDOMNESS_STATE_OBJECT_ID};
 use sui_types::{SUI_SYSTEM_PACKAGE_ID, TypeTag};
 
+#[derive(Clone)]
+pub enum FundSource {
+    Coin(ObjectRef),
+    AddressFund {
+        /// If None, it will be set the same as the total transfer amount.
+        reservation: Option<u64>,
+    },
+    // TODO: Add object fund source
+}
+
+impl FundSource {
+    pub fn coin(coin: ObjectRef) -> Self {
+        Self::Coin(coin)
+    }
+
+    pub fn address_fund() -> Self {
+        Self::AddressFund { reservation: None }
+    }
+
+    pub fn address_fund_with_reservation(reservation: u64) -> Self {
+        Self::AddressFund {
+            reservation: Some(reservation),
+        }
+    }
+}
+
 pub struct TestTransactionBuilder {
     test_data: TestTransactionData,
     sender: SuiAddress,
@@ -295,27 +325,31 @@ impl TestTransactionBuilder {
         self
     }
 
-    pub fn transfer_sui_to_address_balance(mut self, amount: u64, recipient: SuiAddress) -> Self {
+    pub fn transfer_sui_to_address_balance(
+        mut self,
+        source: FundSource,
+        amounts_and_recipients: Vec<(u64, SuiAddress)>,
+    ) -> Self {
         self.test_data =
             TestTransactionData::TransferFundsToAddressBalance(TransferFundsToAddressBalanceData {
-                amount,
+                source,
+                amounts_and_recipients,
                 type_arg: GAS::type_tag(),
-                recipient,
             });
         self
     }
 
     pub fn transfer_funds_to_address_balance(
         mut self,
-        amount: u64,
+        source: FundSource,
+        amounts_and_recipients: Vec<(u64, SuiAddress)>,
         type_arg: TypeTag,
-        recipient: SuiAddress,
     ) -> Self {
         self.test_data =
             TestTransactionData::TransferFundsToAddressBalance(TransferFundsToAddressBalanceData {
-                amount,
+                source,
+                amounts_and_recipients,
                 type_arg,
-                recipient,
             });
         self
     }
@@ -395,17 +429,87 @@ impl TestTransactionBuilder {
                 self.gas_price,
             ),
             TestTransactionData::TransferFundsToAddressBalance(data) => {
-                TransactionData::new_transfer_funds_to_address_balance(
-                    data.recipient,
+                let mut builder = ProgrammableTransactionBuilder::new();
+                let source = match data.source.clone() {
+                    FundSource::Coin(coin) => {
+                        if coin == self.gas_object {
+                            Argument::GasCoin
+                        } else {
+                            builder.obj(ObjectArg::ImmOrOwnedObject(coin)).unwrap()
+                        }
+                    }
+                    FundSource::AddressFund { reservation } => {
+                        let reservation = reservation.unwrap_or_else(|| {
+                            data.amounts_and_recipients
+                                .iter()
+                                .map(|(amount, _)| *amount)
+                                .sum::<u64>()
+                        });
+                        builder
+                            .funds_withdrawal(FundsWithdrawalArg::balance_from_sender(
+                                reservation,
+                                data.type_arg.clone().into(),
+                            ))
+                            .unwrap()
+                    }
+                };
+                for (amount, recipient) in data.amounts_and_recipients {
+                    let balance = match data.source.clone() {
+                        FundSource::Coin(_) => {
+                            let amount_arg = builder.pure(amount).unwrap();
+                            let coin = builder.programmable_move_call(
+                                SUI_FRAMEWORK_PACKAGE_ID,
+                                Identifier::new("coin").unwrap(),
+                                Identifier::new("split").unwrap(),
+                                vec![data.type_arg.clone()],
+                                vec![source, amount_arg],
+                            );
+                            builder.programmable_move_call(
+                                SUI_FRAMEWORK_PACKAGE_ID,
+                                Identifier::new("coin").unwrap(),
+                                Identifier::new("into_balance").unwrap(),
+                                vec![data.type_arg.clone()],
+                                vec![coin],
+                            )
+                        }
+                        FundSource::AddressFund { .. } => {
+                            let amount_arg = builder.pure(U256::from(amount)).unwrap();
+                            let split = builder.programmable_move_call(
+                                SUI_FRAMEWORK_PACKAGE_ID,
+                                Identifier::new("funds_accumulator").unwrap(),
+                                Identifier::new("withdrawal_split").unwrap(),
+                                vec![Balance::type_tag(data.type_arg.clone())],
+                                vec![source, amount_arg],
+                            );
+                            builder.programmable_move_call(
+                                SUI_FRAMEWORK_PACKAGE_ID,
+                                Identifier::new("balance").unwrap(),
+                                Identifier::new("redeem_funds").unwrap(),
+                                vec![data.type_arg.clone()],
+                                vec![split],
+                            )
+                        }
+                    };
+
+                    let recipient_arg = builder.pure(recipient).unwrap();
+                    builder.programmable_move_call(
+                        SUI_FRAMEWORK_PACKAGE_ID,
+                        Identifier::new("balance").unwrap(),
+                        Identifier::new("send_funds").unwrap(),
+                        vec![data.type_arg.clone()],
+                        vec![balance, recipient_arg],
+                    );
+                }
+                let pt = builder.finish();
+                TransactionData::new_programmable(
                     self.sender,
-                    data.amount,
-                    data.type_arg,
-                    self.gas_object,
-                    self.gas_budget
-                        .unwrap_or(self.gas_price * TEST_ONLY_GAS_UNIT_FOR_TRANSFER),
+                    vec![self.gas_object],
+                    pt,
+                    self.gas_budget.unwrap_or(
+                        self.gas_price * TEST_ONLY_GAS_UNIT_FOR_HEAVY_COMPUTATION_STORAGE,
+                    ),
                     self.gas_price,
                 )
-                .unwrap()
             }
             TestTransactionData::SplitCoin(data) => TransactionData::new_split_coin(
                 self.sender,
@@ -548,9 +652,9 @@ struct TransferSuiData {
 }
 
 struct TransferFundsToAddressBalanceData {
-    amount: u64,
+    source: FundSource,
+    amounts_and_recipients: Vec<(u64, SuiAddress)>,
     type_arg: TypeTag,
-    recipient: SuiAddress,
 }
 
 struct SplitCoinData {