Skip to main content

sui_cluster_test/test_case/
staking_test.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{TestCaseImpl, TestContext};
5use async_trait::async_trait;
6use sui_test_transaction_builder::TestTransactionBuilder;
7use sui_types::base_types::SuiAddress;
8use sui_types::effects::TransactionEffectsAPI;
9use sui_types::gas_coin::{GAS, GasCoin};
10use sui_types::governance::StakedSui;
11use tracing::info;
12
13pub struct StakingTest;
14
15#[async_trait]
16impl TestCaseImpl for StakingTest {
17    fn name(&self) -> &'static str {
18        "Staking"
19    }
20
21    fn description(&self) -> &'static str {
22        "Stake SUI with a validator and verify the created StakedSui object"
23    }
24
25    async fn run(&self, ctx: &mut TestContext) -> Result<(), anyhow::Error> {
26        info!("Testing staking workflow");
27
28        let sender = ctx.get_wallet_address();
29        // Fund two coins: one to stake, one to pay for gas. Both object refs are
30        // supplied explicitly (from the faucet response), keeping transaction
31        // construction deterministic.
32        let coins = ctx.get_sui_from_faucet(Some(2)).await;
33        let stake_coin_id = *coins[0].id();
34        let gas_coin_id = *coins[1].id();
35
36        // Pick a validator from the gRPC system-state summary and remember its
37        // staking pool so we can correlate the created StakedSui object.
38        let system_state = ctx.get_latest_sui_system_state().await;
39        let validator = system_state
40            .active_validators
41            .first()
42            .expect("Should have at least one active validator");
43        let validator_addr = validator.sui_address;
44        let validator_pool_id = validator.staking_pool_id;
45        info!("Staking to validator: {validator_addr} (pool {validator_pool_id})");
46
47        let gas_price = ctx.get_reference_gas_price().await;
48        let gas_ref = ctx.current_object_ref(gas_coin_id).await;
49        let stake_ref = ctx.current_object_ref(stake_coin_id).await;
50
51        // Snapshot the staker's indexed SUI balance and owned `Coin<SUI>` count
52        // (StateService) before staking, so we can verify the post-tx accounting.
53        let sui_balance_before = Self::sui_balance(ctx, sender).await;
54        let sui_coin_count_before = Self::sui_coin_count(ctx, sender).await;
55
56        let data = TestTransactionBuilder::new(sender, gas_ref, gas_price)
57            .call_staking(stake_ref, validator_addr)
58            .build();
59        let response = ctx.sign_and_execute(data, "staking transaction").await;
60
61        // The staker's SUI balance change from the executed transaction effects.
62        let sender_sdk: sui_sdk_types::Address = sender.into();
63        let sui_type: sui_sdk_types::TypeTag =
64            sui_types::sui_sdk_types_conversions::type_tag_core_to_sdk(GAS::type_tag()).unwrap();
65        let sui_balance_change = response
66            .balance_changes
67            .iter()
68            .find(|b| b.address == sender_sdk && b.coin_type == sui_type)
69            .map(|b| b.amount)
70            .expect("staking should produce a SUI balance change for the staker");
71
72        // Identify the created StakedSui object from the effects, then fetch and
73        // decode it natively (`LedgerService`).
74        let created = response.effects.created();
75        let mut staked_sui = None;
76        for (obj_ref, _owner) in &created {
77            let object = ctx.get_grpc_client().get_object(obj_ref.0).await?;
78            if let Ok(stake) = StakedSui::try_from(&object) {
79                staked_sui = Some((object, stake));
80                break;
81            }
82        }
83        let (object, stake) = staked_sui.expect("Staking should create a StakedSui object");
84
85        // Owner is the staker.
86        assert_eq!(
87            object.owner(),
88            &sui_types::object::Owner::AddressOwner(sender),
89            "StakedSui should be owned by the staker",
90        );
91        // Principal matches the staked coin value.
92        assert_eq!(
93            stake.principal(),
94            coins[0].value(),
95            "StakedSui principal should equal the staked coin value",
96        );
97        // Staking pool matches the chosen validator.
98        assert_eq!(
99            stake.pool_id(),
100            validator_pool_id,
101            "StakedSui pool should match the validator's staking pool",
102        );
103        // A newly requested stake activates in the epoch AFTER the one the
104        // transaction executed in — it must not be reported as active now.
105        // Derive the execution epoch from the transaction's own checkpoint
106        // (a pre-execution system-state read races epoch boundaries).
107        let checkpoint_seq = response
108            .checkpoint
109            .expect("waited execution should carry a checkpoint");
110        let execution_epoch = ctx
111            .get_grpc_client()
112            .get_checkpoint_summary(checkpoint_seq)
113            .await?
114            .data()
115            .epoch;
116        assert_eq!(
117            stake.activation_epoch(),
118            execution_epoch + 1,
119            "Newly requested stake should activate in the epoch after execution",
120        );
121        info!(
122            "Staking verified: StakedSui {} principal {} pool {} activates epoch {}",
123            stake.id(),
124            stake.principal(),
125            stake.pool_id(),
126            stake.activation_epoch(),
127        );
128
129        // Balance / coin-count accounting (StateService), matching the old
130        // CoinIndex staking flow:
131        //  - the staked `Coin<SUI>` object is consumed, so the owned SUI coin
132        //    count drops by exactly one (the gas coin remains);
133        //  - the indexed SUI balance changes by exactly the effects-reported
134        //    balance change (principal staked + gas spent).
135        let sui_coin_count_after = Self::sui_coin_count(ctx, sender).await;
136        assert_eq!(
137            sui_coin_count_after,
138            sui_coin_count_before - 1,
139            "staking should consume exactly one Coin<SUI> object",
140        );
141        let sui_balance_after = Self::sui_balance(ctx, sender).await;
142        assert_eq!(
143            sui_balance_after,
144            (sui_balance_before as i128 + sui_balance_change) as u128,
145            "post-stake SUI balance should equal pre-stake balance plus the effects balance change",
146        );
147
148        Ok(())
149    }
150}
151
152impl StakingTest {
153    /// Indexed total SUI balance (`StateService::GetBalance`, keyed on the inner
154    /// coin type `0x2::sui::SUI`).
155    async fn sui_balance(ctx: &TestContext, owner: SuiAddress) -> u128 {
156        ctx.get_grpc_client()
157            .get_balance(owner, &GAS::type_())
158            .await
159            .unwrap()
160            .balance
161            .unwrap_or_default() as u128
162    }
163
164    /// Count of owned `Coin<SUI>` objects (`StateService::ListOwnedObjects`
165    /// filtered by the object type `Coin<0x2::sui::SUI>`). Stream errors are
166    /// propagated (not counted as objects), so an RPC or decoding failure fails
167    /// the test rather than inflating the count.
168    async fn sui_coin_count(ctx: &TestContext, owner: SuiAddress) -> usize {
169        use futures::TryStreamExt;
170        let client = ctx.get_grpc_client();
171        let objects: Vec<_> = client
172            .list_owned_objects(owner, Some(GasCoin::type_()))
173            .try_collect()
174            .await
175            .expect("failed to enumerate owned Coin<SUI> objects");
176        objects.len()
177    }
178}