test_cluster/
addr_balance_test_env.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    collections::BTreeMap,
6    path::{Path, PathBuf},
7};
8
9use crate::{TestCluster, TestClusterBuilder};
10use move_core_types::identifier::Identifier;
11use sui_keys::keystore::AccountKeystore;
12use sui_protocol_config::{OverrideGuard, ProtocolConfig, ProtocolVersion};
13use sui_test_transaction_builder::{FundSource, TestTransactionBuilder};
14use sui_types::{
15    SUI_FRAMEWORK_PACKAGE_ID, TypeTag,
16    accumulator_metadata::get_accumulator_object_count,
17    accumulator_root::{AccumulatorValue, U128},
18    balance::Balance,
19    base_types::{FullObjectRef, ObjectID, ObjectRef, SequenceNumber, SuiAddress},
20    coin_reservation::ParsedObjectRefWithdrawal,
21    digests::{ChainIdentifier, TransactionDigest},
22    effects::{TransactionEffects, TransactionEffectsAPI},
23    error::SuiResult,
24    gas_coin::GAS,
25    object::Owner,
26    programmable_transaction_builder::ProgrammableTransactionBuilder,
27    storage::ChildObjectResolver,
28    transaction::{
29        CallArg, FundsWithdrawalArg, GasData, ObjectArg, TransactionData, TransactionDataV1,
30        TransactionExpiration, TransactionKind,
31    },
32};
33
34// TODO: Some of this code may be useful for tests other than address balance tests,
35// we might want to rename it and expand its usage.
36
37pub struct TestEnvBuilder {
38    num_validators: usize,
39    test_cluster_builder_cb: Option<Box<dyn Fn(TestClusterBuilder) -> TestClusterBuilder + Send>>,
40    proto_override_cb:
41        Option<Box<dyn Fn(ProtocolVersion, ProtocolConfig) -> ProtocolConfig + Send>>,
42}
43
44impl Default for TestEnvBuilder {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl TestEnvBuilder {
51    pub fn new() -> Self {
52        Self {
53            test_cluster_builder_cb: None,
54            proto_override_cb: None,
55            num_validators: 1,
56        }
57    }
58
59    #[allow(dead_code)]
60    pub fn with_num_validators(mut self, num_validators: usize) -> Self {
61        self.num_validators = num_validators;
62        self
63    }
64
65    pub fn with_proto_override_cb(
66        mut self,
67        cb: Box<dyn Fn(ProtocolVersion, ProtocolConfig) -> ProtocolConfig + Send>,
68    ) -> Self {
69        self.proto_override_cb = Some(cb);
70        self
71    }
72
73    pub fn with_test_cluster_builder_cb(
74        mut self,
75        cb: Box<dyn Fn(TestClusterBuilder) -> TestClusterBuilder + Send>,
76    ) -> Self {
77        self.test_cluster_builder_cb = Some(cb);
78        self
79    }
80
81    pub async fn build(self) -> TestEnv {
82        let _guard = self
83            .proto_override_cb
84            .map(ProtocolConfig::apply_overrides_for_testing);
85
86        let mut test_cluster_builder =
87            TestClusterBuilder::new().with_num_validators(self.num_validators);
88
89        if let Some(cb) = self.test_cluster_builder_cb {
90            test_cluster_builder = cb(test_cluster_builder);
91        }
92
93        let test_cluster = test_cluster_builder.build().await;
94
95        let chain_id = test_cluster.get_chain_identifier();
96        let rgp = test_cluster.get_reference_gas_price().await;
97
98        let mut test_env = TestEnv {
99            cluster: test_cluster,
100            _guard,
101            rgp,
102            chain_id,
103            gas_objects: BTreeMap::new(),
104        };
105
106        test_env.update_all_gas().await;
107        test_env
108    }
109}
110
111pub struct TestEnv {
112    pub cluster: TestCluster,
113    _guard: Option<OverrideGuard>,
114    pub rgp: u64,
115    pub chain_id: ChainIdentifier,
116    pub gas_objects: BTreeMap<SuiAddress, Vec<ObjectRef>>,
117}
118
119impl TestEnv {
120    pub async fn update_all_gas(&mut self) {
121        // load all gas objects
122        let addresses = self.cluster.wallet.config.keystore.addresses();
123        self.gas_objects.clear();
124
125        for address in addresses {
126            let gas: Vec<_> = self
127                .cluster
128                .wallet
129                .gas_objects(address)
130                .await
131                .unwrap()
132                .into_iter()
133                .map(|(_, obj)| obj.compute_object_reference())
134                .collect();
135
136            self.gas_objects.insert(address, gas);
137        }
138    }
139
140    pub async fn fund_one_address_balance(&mut self, address: SuiAddress, amount: u64) {
141        let gas = self.gas_objects[&address][0];
142        let tx = TestTransactionBuilder::new(address, gas, self.rgp)
143            .transfer_sui_to_address_balance(FundSource::coin(gas), vec![(amount, address)])
144            .build();
145        let (digest, effects) = self
146            .cluster
147            .sign_and_execute_transaction_directly(&tx)
148            .await
149            .unwrap();
150        // update the gas object we used.
151        self.gas_objects.get_mut(&address).unwrap()[0] = effects.gas_object().unwrap().0;
152        self.cluster.wait_for_tx_settlement(&[digest]).await;
153    }
154
155    #[allow(dead_code)]
156    pub async fn fund_all_address_balances(&mut self, amount: u64) {
157        let senders = self.gas_objects.keys().copied().collect::<Vec<_>>();
158        for sender in senders {
159            self.fund_one_address_balance(sender, amount).await;
160        }
161    }
162
163    pub fn get_sender(&self, index: usize) -> SuiAddress {
164        self.gas_objects.keys().copied().nth(index).unwrap()
165    }
166
167    pub fn get_sender_and_gas(&self, index: usize) -> (SuiAddress, ObjectRef) {
168        let sender = self.get_sender(index);
169        let gas = self.gas_objects[&sender][0];
170        (sender, gas)
171    }
172
173    pub fn get_sender_and_all_gas(&self, index: usize) -> (SuiAddress, Vec<ObjectRef>) {
174        let sender = self.get_sender(index);
175        let gas = self.gas_objects[&sender].clone();
176        (sender, gas)
177    }
178
179    pub fn get_all_senders(&self) -> Vec<SuiAddress> {
180        self.cluster.wallet.get_addresses()
181    }
182
183    pub fn get_gas_for_sender(&self, sender: SuiAddress) -> Vec<ObjectRef> {
184        self.gas_objects.get(&sender).unwrap().clone()
185    }
186
187    pub fn tx_builder(&self, sender: SuiAddress) -> TestTransactionBuilder {
188        let gas = self.gas_objects.get(&sender).unwrap()[0];
189        TestTransactionBuilder::new(sender, gas, self.rgp)
190    }
191
192    pub fn tx_builder_with_gas(
193        &self,
194        sender: SuiAddress,
195        gas: ObjectRef,
196    ) -> TestTransactionBuilder {
197        TestTransactionBuilder::new(sender, gas, self.rgp)
198    }
199
200    pub fn tx_builder_with_gas_objects(
201        &self,
202        sender: SuiAddress,
203        gas_objects: Vec<ObjectRef>,
204    ) -> TestTransactionBuilder {
205        TestTransactionBuilder::new_with_gas_objects(sender, gas_objects, self.rgp)
206    }
207
208    pub async fn exec_tx_directly(
209        &mut self,
210        tx: TransactionData,
211    ) -> SuiResult<(TransactionDigest, TransactionEffects)> {
212        let res = self
213            .cluster
214            .sign_and_execute_transaction_directly(&tx)
215            .await;
216        self.update_all_gas().await;
217        res
218    }
219
220    pub async fn setup_test_package(&mut self, path: impl AsRef<Path>) -> ObjectID {
221        let context = &mut self.cluster.wallet;
222        let (sender, gas_object) = context.get_one_gas_object().await.unwrap().unwrap();
223        let gas_price = context.get_reference_gas_price().await.unwrap();
224        let txn = context
225            .sign_transaction(
226                &TestTransactionBuilder::new(sender, gas_object, gas_price)
227                    .publish_async(path.as_ref().to_path_buf())
228                    .await
229                    .build(),
230            )
231            .await;
232        let resp = context.execute_transaction_must_succeed(txn).await;
233        let package_ref = resp.get_new_package_obj().unwrap();
234        self.update_all_gas().await;
235        package_ref.0
236    }
237
238    pub async fn setup_custom_coin(&mut self) -> (SuiAddress, TypeTag) {
239        let (publisher, package_id, _) = self.publish_coins_package().await;
240        let coin_a_type: TypeTag = format!("{}::coin_a::COIN_A", package_id).parse().unwrap();
241        (publisher, coin_a_type)
242    }
243
244    /// Publish the coins package and return (publisher, package_id, coin_type, treasury_cap_ref).
245    /// The MINTABLE_COIN TreasuryCap is unfrozen so new Coin objects can be minted.
246    pub async fn setup_mintable_coin(&mut self) -> (SuiAddress, ObjectID, TypeTag, ObjectRef) {
247        let (publisher, package_id, effects) = self.publish_coins_package().await;
248        let coin_type: TypeTag = format!("{}::mintable_coin::MINTABLE_COIN", package_id)
249            .parse()
250            .unwrap();
251        let treasury_cap_ref = self.find_created_treasury_cap(&effects).await;
252        (publisher, package_id, coin_type, treasury_cap_ref)
253    }
254
255    /// Publish the coins test package and return (publisher, package_id, effects).
256    async fn publish_coins_package(&mut self) -> (SuiAddress, ObjectID, TransactionEffects) {
257        let mut path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
258        path.extend(["..", "sui-e2e-tests", "tests", "data", "coins"]);
259        let (publisher, gas) = self
260            .cluster
261            .wallet
262            .get_one_gas_object()
263            .await
264            .unwrap()
265            .unwrap();
266        let tx = TestTransactionBuilder::new(publisher, gas, self.rgp)
267            .publish_async(path)
268            .await
269            .build();
270        let (_, effects) = self.exec_tx_directly(tx).await.unwrap();
271        assert!(
272            effects.status().is_ok(),
273            "Publish failed: {:?}",
274            effects.status()
275        );
276        let package_id = self.find_created_package(&effects).await;
277        (publisher, package_id, effects)
278    }
279
280    async fn find_created_package(&self, effects: &TransactionEffects) -> ObjectID {
281        for (obj_ref, owner) in effects.created() {
282            if owner.is_immutable()
283                && let Some(obj) = self
284                    .cluster
285                    .get_object_from_fullnode_store(&obj_ref.0)
286                    .await
287                && obj.is_package()
288            {
289                return obj_ref.0;
290            }
291        }
292        panic!("Package should exist among created objects");
293    }
294
295    async fn find_created_treasury_cap(&self, effects: &TransactionEffects) -> ObjectRef {
296        for (obj_ref, owner) in effects.created() {
297            if matches!(owner, Owner::AddressOwner(_))
298                && let Some(obj) = self
299                    .cluster
300                    .get_object_from_fullnode_store(&obj_ref.0)
301                    .await
302                && obj
303                    .data
304                    .try_as_move()
305                    .is_some_and(|m| m.type_().is_treasury_cap())
306            {
307                return obj_ref;
308            }
309        }
310        panic!("TreasuryCap should exist among created objects");
311    }
312
313    /// Mint a `Coin<T>` of the given amount to the recipient using the `TreasuryCap`.
314    /// Returns the updated `TreasuryCap` ref and the new `Coin` object ref.
315    pub async fn mint_coin(
316        &mut self,
317        publisher: SuiAddress,
318        package_id: ObjectID,
319        treasury_cap_ref: ObjectRef,
320        amount: u64,
321        recipient: SuiAddress,
322    ) -> (ObjectRef, ObjectRef) {
323        let tx = self
324            .tx_builder(publisher)
325            .move_call(
326                package_id,
327                "mintable_coin",
328                "mint_and_transfer",
329                vec![
330                    CallArg::Object(ObjectArg::ImmOrOwnedObject(treasury_cap_ref)),
331                    CallArg::Pure(bcs::to_bytes(&amount).unwrap()),
332                    CallArg::Pure(bcs::to_bytes(&recipient).unwrap()),
333                ],
334            )
335            .build();
336        let (_, effects) = self.exec_tx_directly(tx).await.unwrap();
337        assert!(
338            effects.status().is_ok(),
339            "Mint failed: {:?}",
340            effects.status()
341        );
342        let new_treasury_cap_ref = effects
343            .mutated()
344            .into_iter()
345            .find(|(obj_ref, _)| obj_ref.0 == treasury_cap_ref.0)
346            .unwrap()
347            .0;
348        let coin_ref = effects
349            .created()
350            .into_iter()
351            .find(|(_, owner)| matches!(owner, Owner::AddressOwner(addr) if *addr == recipient))
352            .unwrap()
353            .0;
354        (new_treasury_cap_ref, coin_ref)
355    }
356
357    pub fn encode_coin_reservation(
358        &self,
359        sender: SuiAddress,
360        epoch: u64,
361        amount: u64,
362    ) -> ObjectRef {
363        let accumulator_obj_id = get_sui_accumulator_object_id(sender);
364        ParsedObjectRefWithdrawal::new(accumulator_obj_id, epoch, amount)
365            .encode(SequenceNumber::new(), self.chain_id)
366    }
367
368    pub fn encode_coin_reservation_for_type(
369        &self,
370        sender: SuiAddress,
371        epoch: u64,
372        amount: u64,
373        coin_type: TypeTag,
374    ) -> ObjectRef {
375        let accumulator_obj_id = get_accumulator_object_id(sender, coin_type);
376        ParsedObjectRefWithdrawal::new(accumulator_obj_id, epoch, amount)
377            .encode(SequenceNumber::new(), self.chain_id)
378    }
379
380    /// Transfer a portion of a coin to one or more addresses.
381    pub async fn transfer_from_coin_to_address_balance(
382        &mut self,
383        sender: SuiAddress,
384        coin: ObjectRef,
385        amounts_and_recipients: Vec<(u64, SuiAddress)>,
386    ) -> SuiResult<(TransactionDigest, TransactionEffects)> {
387        let tx = self
388            .tx_builder(sender)
389            .transfer_sui_to_address_balance(FundSource::coin(coin), amounts_and_recipients)
390            .build();
391        let res = self.exec_tx_directly(tx).await;
392        self.update_all_gas().await;
393        res
394    }
395
396    /// Transfer the entire coin to a single address.
397    pub async fn transfer_coin_to_address_balance(
398        &mut self,
399        sender: SuiAddress,
400        coin: ObjectRef,
401        recipient: SuiAddress,
402    ) -> SuiResult<(TransactionDigest, TransactionEffects)> {
403        let tx = self
404            .tx_builder(sender)
405            .transfer(FullObjectRef::from_fastpath_ref(coin), recipient)
406            .build();
407        let res = self.exec_tx_directly(tx).await;
408        self.update_all_gas().await;
409        res
410    }
411
412    pub fn verify_accumulator_exists(&self, owner: SuiAddress, expected_balance: u64) {
413        self.cluster.fullnode_handle.sui_node.with(|node| {
414            let state = node.state();
415            let child_object_resolver = state.get_child_object_resolver().as_ref();
416            verify_accumulator_exists(child_object_resolver, owner, expected_balance);
417        });
418    }
419
420    /// Verify the accumulator object count after settlement.
421    pub fn verify_accumulator_object_count(&self, expected_object_count: u64) {
422        self.cluster.fullnode_handle.sui_node.with(|node| {
423            let state = node.state();
424
425            let object_count = get_accumulator_object_count(state.get_object_store().as_ref())
426                .expect("read cannot fail")
427                .expect("accumulator object count should exist after settlement");
428            assert_eq!(object_count, expected_object_count);
429        });
430    }
431
432    /// Get the balance of the owner's SUI address balance.
433    pub fn get_sui_balance_ab(&self, owner: SuiAddress) -> u64 {
434        self.get_balance_ab(owner, GAS::type_tag())
435    }
436
437    pub async fn get_coin_balance(&self, object_id: ObjectID) -> u64 {
438        self.cluster
439            .get_object_from_fullnode_store(&object_id)
440            .await
441            .expect("coin object should exist")
442            .data
443            .try_as_move()
444            .expect("should be a Move object")
445            .get_coin_value_unsafe()
446    }
447
448    /// Get the balance of the owner's address balance for a given coin type.
449    pub fn get_balance_ab(&self, owner: SuiAddress, coin_type: TypeTag) -> u64 {
450        let db_balance = self.cluster.fullnode_handle.sui_node.with({
451            let coin_type = coin_type.clone();
452            move |node| {
453                let state = node.state();
454                let child_object_resolver = state.get_child_object_resolver().as_ref();
455                get_balance(child_object_resolver, owner, coin_type)
456            }
457        });
458
459        let client = self.cluster.grpc_client();
460        // Check that the rpc balance agrees with the db balance, on a best-effort basis.
461        tokio::task::spawn(async move {
462            match client
463                .get_balance(owner, &coin_type.to_canonical_string(true).parse().unwrap())
464                .await
465            {
466                Ok(rpc_balance) => {
467                    assert_eq!(db_balance, rpc_balance.address_balance());
468                }
469                Err(e) => {
470                    // this usually just means the cluster shut down first before the rpc
471                    // completed.
472                    tracing::info!("Failed to verify balance via gRPC: {e}");
473                }
474            }
475        });
476
477        db_balance
478    }
479
480    /// Get the total balance of SUI owned by the address (including address balance and coins).
481    pub async fn get_sui_balance(&self, owner: SuiAddress) -> u64 {
482        self.get_balance_for_coin_type(owner, GAS::type_tag()).await
483    }
484
485    /// Get the total balance of a given coin type owned by the address (including address balance and coins).
486    pub async fn get_balance_for_coin_type(&self, owner: SuiAddress, coin_type: TypeTag) -> u64 {
487        let client = self.cluster.grpc_client();
488        let rpc_balance = client
489            .get_balance(owner, &coin_type.to_canonical_string(true).parse().unwrap())
490            .await
491            .unwrap();
492        rpc_balance.balance()
493    }
494
495    pub fn verify_accumulator_removed(&self, owner: SuiAddress) {
496        self.cluster.fullnode_handle.sui_node.with(|node| {
497            let state = node.state();
498            let child_object_resolver = state.get_child_object_resolver().as_ref();
499            let sui_coin_type = Balance::type_tag(GAS::type_tag());
500            assert!(
501                !AccumulatorValue::exists(child_object_resolver, None, owner, &sui_coin_type)
502                    .unwrap(),
503                "Accumulator value should have been removed"
504            );
505        });
506    }
507
508    pub async fn trigger_reconfiguration(&self) {
509        self.cluster.trigger_reconfiguration().await;
510    }
511
512    pub fn create_gasless_transaction(
513        &self,
514        amount: u64,
515        token_type: TypeTag,
516        sender: SuiAddress,
517        recipient: SuiAddress,
518        nonce: u32,
519        epoch: u64,
520    ) -> TransactionData {
521        let mut builder = ProgrammableTransactionBuilder::new();
522        let withdraw_arg = FundsWithdrawalArg::balance_from_sender(amount, token_type.clone());
523        let withdraw_arg = builder.funds_withdrawal(withdraw_arg).unwrap();
524        let balance = builder.programmable_move_call(
525            SUI_FRAMEWORK_PACKAGE_ID,
526            Identifier::new("balance").unwrap(),
527            Identifier::new("redeem_funds").unwrap(),
528            vec![token_type.clone()],
529            vec![withdraw_arg],
530        );
531        let recipient_arg = builder.pure(recipient).unwrap();
532        builder.programmable_move_call(
533            SUI_FRAMEWORK_PACKAGE_ID,
534            Identifier::new("balance").unwrap(),
535            Identifier::new("send_funds").unwrap(),
536            vec![token_type],
537            vec![balance, recipient_arg],
538        );
539        let tx_kind = TransactionKind::ProgrammableTransaction(builder.finish());
540        self.gasless_transaction_data(tx_kind, sender, nonce, epoch)
541    }
542
543    pub fn gasless_transaction_data(
544        &self,
545        tx_kind: TransactionKind,
546        sender: SuiAddress,
547        nonce: u32,
548        epoch: u64,
549    ) -> TransactionData {
550        TransactionData::V1(TransactionDataV1 {
551            kind: tx_kind,
552            sender,
553            gas_data: GasData {
554                payment: vec![],
555                owner: sender,
556                price: 0,
557                budget: 0,
558            },
559            expiration: TransactionExpiration::ValidDuring {
560                min_epoch: Some(epoch),
561                max_epoch: Some(epoch),
562                min_timestamp: None,
563                max_timestamp: None,
564                chain: self.chain_id,
565                nonce,
566            },
567        })
568    }
569
570    /// Publishes the `object_balance` example package, creates an owned vault object,
571    /// and funds it with the given amount. Returns (package_id, vault_id).
572    pub async fn setup_funded_object_balance_vault(&mut self, amount: u64) -> (ObjectID, ObjectID) {
573        let sender = self.get_sender(0);
574
575        let tx = self
576            .tx_builder(sender)
577            .publish_examples("object_balance")
578            .await
579            .build();
580        let (_, effects) = self.exec_tx_directly(tx).await.unwrap();
581        let package_id = effects
582            .created()
583            .into_iter()
584            .find(|(_, owner)| owner.is_immutable())
585            .unwrap()
586            .0
587            .0;
588
589        let tx = self
590            .tx_builder(sender)
591            .move_call(package_id, "object_balance", "new_owned", vec![])
592            .build();
593        let (_, effects) = self.exec_tx_directly(tx).await.unwrap();
594        let vault_id = effects.created().into_iter().next().unwrap().0.0;
595
596        let tx = self
597            .tx_builder(sender)
598            .transfer_sui_to_address_balance(
599                FundSource::coin(self.get_sender_and_gas(0).1),
600                vec![(amount, vault_id.into())],
601            )
602            .build();
603        self.exec_tx_directly(tx).await.unwrap();
604        self.trigger_reconfiguration().await;
605
606        (package_id, vault_id)
607    }
608
609    /// Publish the trusted_coin package and return (package_id, coin_type, treasury_cap).
610    pub async fn publish_trusted_coin(
611        &mut self,
612        sender: SuiAddress,
613    ) -> (ObjectID, TypeTag, ObjectRef) {
614        let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
615        path.pop();
616        path.extend(["sui-e2e-tests", "tests", "rpc", "data", "trusted_coin"]);
617
618        let tx = self.tx_builder(sender).publish_async(path).await.build();
619        let (_, effects) = self.exec_tx_directly(tx).await.unwrap();
620
621        let package_id = effects.published_packages().into_iter().next().unwrap();
622        let coin_type: TypeTag = format!("{}::trusted_coin::TRUSTED_COIN", package_id)
623            .parse()
624            .unwrap();
625
626        // Find the treasury cap by checking object type
627        let mut treasury_cap = None;
628        for (obj_ref, owner) in effects.created() {
629            if owner.is_address_owned() {
630                let object = self
631                    .cluster
632                    .fullnode_handle
633                    .sui_node
634                    .with_async(
635                        |node| async move { node.state().get_object(&obj_ref.0).await.unwrap() },
636                    )
637                    .await;
638                if object.type_().unwrap().name().as_str() == "TreasuryCap" {
639                    treasury_cap = Some(obj_ref);
640                    break;
641                }
642            }
643        }
644
645        (
646            package_id,
647            coin_type,
648            treasury_cap.expect("Treasury cap not found"),
649        )
650    }
651
652    /// Mint a trusted coin. Returns (coin_ref, updated_treasury_cap).
653    pub async fn mint_trusted_coin(
654        &mut self,
655        sender: SuiAddress,
656        package_id: ObjectID,
657        treasury_cap: ObjectRef,
658        amount: u64,
659    ) -> (ObjectRef, ObjectRef) {
660        let tx = self
661            .tx_builder(sender)
662            .move_call(
663                package_id,
664                "trusted_coin",
665                "mint",
666                vec![
667                    CallArg::Object(ObjectArg::ImmOrOwnedObject(treasury_cap)),
668                    CallArg::Pure(bcs::to_bytes(&amount).unwrap()),
669                ],
670            )
671            .build();
672        let (_, effects) = self.exec_tx_directly(tx).await.unwrap();
673
674        let coin_ref = effects
675            .created()
676            .iter()
677            .find(|(_, owner)| owner.is_address_owned())
678            .unwrap()
679            .0;
680
681        let new_treasury_cap = effects
682            .mutated()
683            .iter()
684            .find(|(obj_ref, _)| obj_ref.0 == treasury_cap.0)
685            .map(|(obj_ref, _)| *obj_ref)
686            .unwrap();
687
688        (coin_ref, new_treasury_cap)
689    }
690
691    /// Transfer a coin to recipient.
692    pub async fn transfer_coin(
693        &mut self,
694        sender: SuiAddress,
695        coin: ObjectRef,
696        recipient: SuiAddress,
697    ) {
698        let tx = self
699            .tx_builder(sender)
700            .transfer(FullObjectRef::from_fastpath_ref(coin), recipient)
701            .build();
702        self.exec_tx_directly(tx).await.unwrap();
703    }
704
705    /// Transfer SUI from sender's gas to recipient.
706    pub async fn transfer_sui(&mut self, sender: SuiAddress, recipient: SuiAddress, amount: u64) {
707        let tx = self
708            .tx_builder(sender)
709            .transfer_sui(Some(amount), recipient)
710            .build();
711        self.exec_tx_directly(tx).await.unwrap();
712    }
713
714    /// Transfer SUI from sender's gas to recipient's address balance.
715    pub async fn transfer_sui_to_address_balance(
716        &mut self,
717        sender: SuiAddress,
718        recipient: SuiAddress,
719        amount: u64,
720    ) {
721        let gas = self.gas_objects[&sender][0];
722        let tx = TestTransactionBuilder::new(sender, gas, self.rgp)
723            .transfer_sui_to_address_balance(FundSource::coin(gas), vec![(amount, recipient)])
724            .build();
725        self.exec_tx_directly(tx).await.unwrap();
726    }
727
728    /// Convenience: publish trusted_coin and set up coins for recipient per config.
729    pub async fn publish_trusted_coin_and_setup(
730        &mut self,
731        funder: SuiAddress,
732        recipient: SuiAddress,
733        config: &CoinTypeConfig,
734        coin_amount: u64,
735    ) -> (ObjectID, TypeTag) {
736        let (package_id, coin_type, mut treasury_cap) = self.publish_trusted_coin(funder).await;
737
738        for _ in 0..config.real_coins {
739            let (coin, new_cap) = self
740                .mint_trusted_coin(funder, package_id, treasury_cap, coin_amount)
741                .await;
742            treasury_cap = new_cap;
743            self.transfer_coin(funder, coin, recipient).await;
744        }
745
746        if config.has_address_balance {
747            let (coin, _) = self
748                .mint_trusted_coin(funder, package_id, treasury_cap, coin_amount)
749                .await;
750            let tx = self
751                .tx_builder(funder)
752                .transfer_funds_to_address_balance(
753                    FundSource::Coin(coin),
754                    vec![(coin_amount, recipient)],
755                    coin_type.clone(),
756                )
757                .build();
758            self.exec_tx_directly(tx).await.unwrap();
759        }
760
761        (package_id, coin_type)
762    }
763
764    /// Legacy: publish trusted_coin with one real coin and address balance for sender.
765    pub async fn publish_and_mint_trusted_coin(
766        &mut self,
767        sender: SuiAddress,
768        amount: u64,
769    ) -> (ObjectID, TypeTag) {
770        let config = CoinTypeConfig {
771            real_coins: 1,
772            has_address_balance: true,
773        };
774        self.publish_trusted_coin_and_setup(sender, sender, &config, amount)
775            .await
776    }
777}
778
779/// Configuration for a single coin type in a test scenario.
780#[derive(Clone, Debug)]
781pub struct CoinTypeConfig {
782    /// Number of real coins to create for this type.
783    pub real_coins: usize,
784    /// Whether to create an address balance (fake coin) for this type.
785    pub has_address_balance: bool,
786}
787
788pub fn get_sui_accumulator_object_id(sender: SuiAddress) -> ObjectID {
789    get_accumulator_object_id(sender, GAS::type_tag())
790}
791
792pub fn get_accumulator_object_id(sender: SuiAddress, coin_type: TypeTag) -> ObjectID {
793    *AccumulatorValue::get_field_id(sender, &Balance::type_tag(coin_type))
794        .unwrap()
795        .inner()
796}
797
798pub fn get_balance(
799    child_object_resolver: &dyn ChildObjectResolver,
800    owner: SuiAddress,
801    coin_type: TypeTag,
802) -> u64 {
803    sui_core::accumulators::balances::get_balance(owner, child_object_resolver, coin_type).unwrap()
804}
805
806pub fn get_sui_balance(child_object_resolver: &dyn ChildObjectResolver, owner: SuiAddress) -> u64 {
807    get_balance(child_object_resolver, owner, GAS::type_tag())
808}
809
810pub fn verify_accumulator_exists(
811    child_object_resolver: &dyn ChildObjectResolver,
812    owner: SuiAddress,
813    expected_balance: u64,
814) {
815    let sui_coin_type = Balance::type_tag(GAS::type_tag());
816
817    assert!(
818        AccumulatorValue::exists(child_object_resolver, None, owner, &sui_coin_type).unwrap(),
819        "Accumulator value should have been created"
820    );
821
822    let accumulator_object =
823        AccumulatorValue::load_object(child_object_resolver, None, owner, &sui_coin_type)
824            .expect("read cannot fail")
825            .expect("accumulator should exist");
826
827    assert!(
828        accumulator_object
829            .data
830            .try_as_move()
831            .unwrap()
832            .type_()
833            .is_efficient_representation()
834    );
835
836    let accumulator_value =
837        AccumulatorValue::load(child_object_resolver, None, owner, &sui_coin_type)
838            .expect("read cannot fail")
839            .expect("accumulator should exist");
840
841    assert_eq!(
842        accumulator_value,
843        AccumulatorValue::U128(U128 {
844            value: expected_balance as u128
845        }),
846        "Accumulator value should be {expected_balance}"
847    );
848}
test_cluster/addr_balance_test_env.rs

test_cluster/
addr_balance_test_env.rs
