diff --git a/crates/sui-framework/packages/sui-framework/sources/tx_context.move b/crates/sui-framework/packages/sui-framework/sources/tx_context.move index 1fdef9ff83a81..bcb88c3842728 100644 --- a/crates/sui-framework/packages/sui-framework/sources/tx_context.move +++ b/crates/sui-framework/packages/sui-framework/sources/tx_context.move @@ -135,6 +135,11 @@ public fun increment_epoch_number(self: &mut TxContext) { self.epoch = self.epoch + 1 } +#[test_only] +public fun set_epoch_number_for_testing(self: &mut TxContext, epoch: u64) { + self.epoch = epoch +} + #[test_only] public fun increment_epoch_timestamp(self: &mut TxContext, delta_ms: u64) { self.epoch_timestamp_ms = self.epoch_timestamp_ms + delta_ms diff --git a/crates/sui-framework/packages/sui-system/sources/staking_pool.move b/crates/sui-framework/packages/sui-system/sources/staking_pool.move index 0274e8e061435..0b75801200c90 100644 --- a/crates/sui-framework/packages/sui-system/sources/staking_pool.move +++ b/crates/sui-framework/packages/sui-system/sources/staking_pool.move @@ -687,37 +687,31 @@ module sui_system::staking_pool { #[test_only] public use fun fungible_staked_sui_data_principal_value as FungibleStakedSuiData.principal_value; - // NEVER remove `#[test_only]` #[test_only] - public(package) fun process_pending_stake_withdraw_test_only(pool: &mut StakingPool) { + public(package) fun process_pending_stake_withdraw_for_testing(pool: &mut StakingPool) { pool.process_pending_stake_withdraw() } - // NEVER remove `#[test_only]` #[test_only] - public(package) fun increase_pending_pool_token_withdraw_test_only(pool: &mut StakingPool, delta: u64) { + public(package) fun increase_pending_pool_token_withdraw_for_testing(pool: &mut StakingPool, delta: u64) { pool.pending_pool_token_withdraw = pool.pending_pool_token_withdraw + delta } - // NEVER remove `#[test_only]` #[test_only] - public(package) fun increase_pending_total_sui_withdraw_test_only(pool: &mut StakingPool, delta: u64) { + public(package) fun increase_pending_total_sui_withdraw_for_testing(pool: &mut StakingPool, delta: u64) { pool.pending_total_sui_withdraw = pool.pending_total_sui_withdraw + delta } - // NEVER remove `#[test_only]` #[test_only] public(package) fun pending_total_sui_withdraw(pool: &StakingPool): u64 { pool.pending_total_sui_withdraw } - // NEVER remove `#[test_only]` #[test_only] public(package) fun pool_token_balance(pool: &StakingPool): u64 { pool.pool_token_balance } - // NEVER remove `#[test_only]` #[test_only] public(package) fun pending_pool_token_withdraw(pool: &StakingPool): u64 { pool.pending_pool_token_withdraw diff --git a/crates/sui-framework/packages/sui-system/sources/sui_system.move b/crates/sui-framework/packages/sui-system/sources/sui_system.move index a3d594222d28d..cd50f67ec8e99 100644 --- a/crates/sui-framework/packages/sui-system/sources/sui_system.move +++ b/crates/sui-framework/packages/sui-system/sources/sui_system.move @@ -698,7 +698,6 @@ module sui_system::sui_system { self.set_epoch_for_testing(epoch_num) } - // NEVER remove `#[test_only]` #[test_only] public(package) fun increment_epoch_for_testing(wrapper: &mut SuiSystemState) { let self = load_system_state_mut(wrapper); diff --git a/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move b/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move index b79bf42e38339..c9dd40dd1d90c 100644 --- a/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move +++ b/crates/sui-framework/packages/sui-system/sources/sui_system_state_inner.move @@ -1121,13 +1121,11 @@ module sui_system::sui_system_state_inner { self.stake_subsidy.get_distribution_counter() } - // NEVER remove `#[test_only]` #[test_only] public(package) fun set_epoch_for_testing(self: &mut SuiSystemStateInnerV2, epoch_num: u64) { self.epoch = epoch_num } - // NEVER remove `#[test_only]` #[test_only] public(package) fun increment_epoch_for_testing(self: &mut SuiSystemStateInnerV2) { self.epoch = self.epoch + 1 diff --git a/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move b/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move index a809c749b1cc5..0d41892a49c2b 100644 --- a/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move +++ b/crates/sui-framework/packages/sui-system/tests/rewards_distribution_tests.move @@ -631,58 +631,197 @@ module sui_system::rewards_distribution_tests { test.end(); } + // This test triggers both sui balance and pool token to fall short but no underflow happens. #[test] - fun test_process_pending_stake_withdraw_no_underflow_in_safe_mode() { + fun test_process_pending_stake_withdraw_no_underflow_in_safe_mode_1() { let start_epoch: u64 = 1; let epoch_start_time = 100000000000; set_up_sui_system_state(); - let mut scenario_val = test_scenario::begin(VALIDATOR_ADDR_1); // val 1 has 100 staked sui - let mut sui_system = scenario_val.take_shared(); + let mut test = test_scenario::begin(VALIDATOR_ADDR_1); // val 1 has 100 staked sui + let mut sui_system = test.take_shared(); + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + assert!(pool.sui_balance() == 100 * MIST_PER_SUI); + assert!(pool.pool_token_balance() == 100 * MIST_PER_SUI); - start_epoch.do!(|_| scenario_val.ctx().increment_epoch_number()); // epoch 1, entering safe mode + start_epoch.do!(|_| test.ctx().increment_epoch_number()); // epoch 1, entering safe mode sui_system.set_epoch_for_testing(start_epoch); // staker 1 stakes 101 sui in safe mode - scenario_val.next_tx(STAKER_ADDR_1); - sui_system.request_add_stake(coin::mint_for_testing(101 * MIST_PER_SUI, scenario_val.ctx()), VALIDATOR_ADDR_1, scenario_val.ctx()); - scenario_val.next_tx(STAKER_ADDR_1); + test.next_tx(STAKER_ADDR_1); + sui_system.request_add_stake(coin::mint_for_testing(101 * MIST_PER_SUI, test.ctx()), VALIDATOR_ADDR_1, test.ctx()); + test.next_tx(STAKER_ADDR_1); + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + assert!(pool.sui_balance() == 100 * MIST_PER_SUI); + assert!(pool.pool_token_balance() == 100 * MIST_PER_SUI); + assert!(pool.pending_stake_amount() == 101 * MIST_PER_SUI); - scenario_val.ctx().increment_epoch_number(); // epoch 2: still in safe mode + test.ctx().increment_epoch_number(); // epoch 2: still in safe mode // There is no need to update `sui_system.set_epoch_for_testing` because // it's `ctx.epoch()` being used here sui_system.increment_epoch_for_testing(); - let staked_sui = scenario_val.take_from_address(STAKER_ADDR_1); + let staked_sui = test.take_from_address(STAKER_ADDR_1); // staker 1 unstakes in safe mode - sui_system.request_withdraw_stake(staked_sui, scenario_val.ctx()); + sui_system.request_withdraw_stake(staked_sui, test.ctx()); let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); let pool = val.get_staking_pool_ref(); - assert!(pool.sui_balance() == 100 * MIST_PER_SUI, 0); - assert!(pool.pool_token_balance() == 100 * MIST_PER_SUI, 0); + assert!(pool.sui_balance() == 100 * MIST_PER_SUI); + // Note `pending_pool_token_withdraw > pool_token_balance` + assert!(pool.pool_token_balance() == 100 * MIST_PER_SUI); // FIXME: these 3 values will be fixed in the next PR - assert!(pool.pending_stake_amount() == 101 * MIST_PER_SUI, 0); - assert!(pool.pending_total_sui_withdraw() == 101 * MIST_PER_SUI, 0); - assert!(pool.pending_pool_token_withdraw() == 101 * MIST_PER_SUI, 0); + assert!(pool.pending_stake_amount() == 101 * MIST_PER_SUI); + assert!(pool.pending_total_sui_withdraw() == 101 * MIST_PER_SUI); + assert!(pool.pending_pool_token_withdraw() == 101 * MIST_PER_SUI); // epoch 3: exiting safe mode // There is no underflow here sui_system .inner_mut_for_testing() - .advance_epoch(start_epoch + 2, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time, scenario_val.ctx()) + .advance_epoch(start_epoch + 2, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time, test.ctx()) .destroy_for_testing(); // balance returned from `advance_epoch` - scenario_val.next_tx(VALIDATOR_ADDR_1); + test.next_tx(VALIDATOR_ADDR_1); let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); let pool = val.get_staking_pool_ref(); - assert!(pool.pending_stake_amount() == 0 * MIST_PER_SUI, 0); - assert!(pool.pending_total_sui_withdraw() == 0 * MIST_PER_SUI, 0); - assert!(pool.pending_pool_token_withdraw() == 0 * MIST_PER_SUI, 0); + assert!(pool.pending_stake_amount() == 0 * MIST_PER_SUI); + assert!(pool.pending_total_sui_withdraw() == 0 * MIST_PER_SUI); + assert!(pool.pending_pool_token_withdraw() == 0 * MIST_PER_SUI); // FIXME: these 2 values will be fixed in the next PR - assert!(pool.sui_balance() == 101 * MIST_PER_SUI, 0); - assert!(pool.pool_token_balance() == 101 * MIST_PER_SUI, 0); + assert!(pool.sui_balance() == 101 * MIST_PER_SUI); + assert!(pool.pool_token_balance() == 101 * MIST_PER_SUI); test_scenario::return_shared(sui_system); - scenario_val.end(); + test.end(); + } + + // This test triggers pool token to fall short but no underflow happens. + #[test] + fun test_process_pending_stake_withdraw_no_underflow_in_safe_mode_2() { + let mut epoch: u64 = 1; + let epoch_start_time = 100000000000; + + // Epoch 0: genesis + set_up_sui_system_state(); // val 1 has 100 staked sui from this + let mut test = test_scenario::begin(VALIDATOR_ADDR_1); + let mut sui_system = test.take_shared(); + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + assert!(pool.sui_balance() == 100 * MIST_PER_SUI); + assert!(pool.pool_token_balance() == 100 * MIST_PER_SUI); + + // Epoch 1: + test.ctx().set_epoch_number_for_testing(0); + assert!(test.ctx().epoch() == 0); + sui_system.set_epoch_for_testing(0); + sui_system + .inner_mut_for_testing() + .advance_epoch(epoch, 65, balance::zero(), balance::create_for_testing(100 * MIST_PER_SUI), 0, 0, 0, 0, epoch_start_time, test.ctx()) + .destroy_for_testing(); // balance returned from `advance_epoch` + test.ctx().increment_epoch_number(); + assert!(test.ctx().epoch() == epoch); + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + let exchange_rate = pool.exchange_rates()[1]; + assert!(exchange_rate.sui_amount() == 125 * MIST_PER_SUI); + assert!(exchange_rate.pool_token_amount() == 100 * MIST_PER_SUI); + + // Epoch 2: entering safe mode + epoch = epoch + 1; // 2 + test.ctx().increment_epoch_number(); + assert!(test.ctx().epoch() == epoch); + sui_system.set_epoch_for_testing(epoch); + + // staker 1 stakes 101 sui in epoch 2 + test.next_tx(STAKER_ADDR_1); + sui_system.request_add_stake(coin::mint_for_testing(100 * MIST_PER_SUI, test.ctx()), VALIDATOR_ADDR_1, test.ctx()); + test.next_tx(STAKER_ADDR_1); + + // Epoch 3: still in safe mode + epoch = epoch + 1; // 3 + test.ctx().increment_epoch_number(); + assert!(test.ctx().epoch() == epoch); + // There is no need to update `sui_system.set_epoch_for_testing` because + // it's `ctx.epoch()` being used here + sui_system.increment_epoch_for_testing(); + + // Epoch 4: still in safe mode + epoch = epoch + 1; // 4 + test.ctx().increment_epoch_number(); + assert!(test.ctx().epoch() == epoch); + // There is no need to update `sui_system.set_epoch_for_testing` because + // it's `ctx.epoch()` being used here + sui_system.increment_epoch_for_testing(); + + // Epoch 5: now out of safe mode + epoch = epoch + 1; // 5 + sui_system + .inner_mut_for_testing() + .advance_epoch(epoch, 65, balance::zero(), balance::create_for_testing(100 * MIST_PER_SUI), 0, 0, 0, 0, epoch_start_time, test.ctx()) + .destroy_for_testing(); // balance returned from `advance_epoch` + test.ctx().increment_epoch_number(); + assert!(test.ctx().epoch() == epoch); + + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + let exchange_rate = pool.exchange_rates()[5]; + assert!(exchange_rate.sui_amount() == 250 * MIST_PER_SUI); + // old pool token balance / old pool sui balance * (pool sui balance + pending stake) + // 100 / 150 * (150 + 100) = 166666666666 + assert!(exchange_rate.pool_token_amount() == 166666666666); + assert!(pool.sui_balance() == 250 * MIST_PER_SUI); + assert!(pool.pool_token_balance() == 166666666666); + + // staker 1 unstakes + test.next_tx(STAKER_ADDR_1); + let staked_sui = test.take_from_address(STAKER_ADDR_1); + assert!(staked_sui.stake_activation_epoch() == 3); + sui_system.request_withdraw_stake(staked_sui, test.ctx()); + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + assert!(pool.pending_total_sui_withdraw() == 120 * MIST_PER_SUI); // 100 pricinpal + 20 rewards + // Pool's exchange rate for epoch 2 is missing because of safe mode + // So it fallbacks to epoch 1's exchange rate: PoolTokenExchangeRate { + // sui_amount: 125000000000, + // pool_token_amount: 100000000000 + // } + // pending pool token to withdraw: 100 / 125 * 100 = 80 + assert!(pool.pending_pool_token_withdraw() == 80 * MIST_PER_SUI); + + // Epoch 6: + epoch = epoch + 1; // 6 + test.ctx().increment_epoch_number(); // epoch 5 exit safe mode + assert!(test.ctx().epoch() == epoch); + sui_system + .inner_mut_for_testing() + .advance_epoch(epoch, 65, balance::zero(), balance::create_for_testing(100 * MIST_PER_SUI), 0, 0, 0, 0, epoch_start_time, test.ctx()) + .destroy_for_testing(); // balance returned from `advance_epoch` + + assert!(test.ctx().epoch() == epoch); + + // Validator unstakes + test.next_tx(VALIDATOR_ADDR_1); + let staked_sui = test.take_from_address(VALIDATOR_ADDR_1); + sui_system.request_withdraw_stake(staked_sui, test.ctx()); + let val = sui_system.active_validator_by_address(VALIDATOR_ADDR_1); + let pool = val.get_staking_pool_ref(); + // insufficient pool token balance + assert!(pool.pool_token_balance() < pool.pending_pool_token_withdraw()); + test.next_tx(VALIDATOR_ADDR_1); + + // Epoch 7: + // No underflow should happen + epoch = epoch + 1; // 7 + test.ctx().increment_epoch_number(); + assert!(test.ctx().epoch() == epoch); + sui_system + .inner_mut_for_testing() + .advance_epoch(epoch, 65, balance::zero(), balance::zero(), 0, 0, 0, 0, epoch_start_time, test.ctx()) + .destroy_for_testing(); // balance returned from `advance_epoch` + + test_scenario::return_shared(sui_system); + test.end(); } } diff --git a/crates/sui-framework/packages/sui-system/tests/staking_pool.move b/crates/sui-framework/packages/sui-system/tests/staking_pool.move index 242f8a265bfd9..62236bab5d7e2 100644 --- a/crates/sui-framework/packages/sui-system/tests/staking_pool.move +++ b/crates/sui-framework/packages/sui-system/tests/staking_pool.move @@ -117,27 +117,27 @@ module sui_system::staking_pool_tests { #[test] fun test_process_pending_stake_withdraw_no_underflow() { - let mut scenario = test_scenario::begin(@0x0); - let mut staking_pool = staking_pool::new(scenario.ctx()); + let mut test = test_scenario::begin(@0x0); + let mut staking_pool = staking_pool::new(test.ctx()); staking_pool.activate_staking_pool(0); let sui = balance::create_for_testing(1_000_000_000); - let staked_sui_1 = staking_pool.request_add_stake(sui, scenario.ctx().epoch() + 1, scenario.ctx()); - assert!(distribute_rewards_and_advance_epoch(&mut staking_pool, &mut scenario, 0) == 1, 0); + let staked_sui_1 = staking_pool.request_add_stake(sui, test.ctx().epoch() + 1, test.ctx()); + assert!(distribute_rewards_and_advance_epoch(&mut staking_pool, &mut test, 0) == 1); - staking_pool.increase_pending_pool_token_withdraw_test_only(1_000_000_000); - staking_pool.increase_pending_total_sui_withdraw_test_only(1_000_000_000); + staking_pool.increase_pending_pool_token_withdraw_for_testing(1_000_000_000); + staking_pool.increase_pending_total_sui_withdraw_for_testing(1_000_000_000); - staking_pool.process_pending_stake_withdraw_test_only(); + staking_pool.process_pending_stake_withdraw_for_testing(); - assert!(staking_pool.sui_balance() == 0, 0); - assert!(staking_pool.pending_total_sui_withdraw() == 0, 0); - assert!(staking_pool.pool_token_balance() == 0, 0); - assert!(staking_pool.pending_pool_token_withdraw() == 0, 0); + assert!(staking_pool.sui_balance() == 0); + assert!(staking_pool.pending_total_sui_withdraw() == 0); + assert!(staking_pool.pool_token_balance() == 0); + assert!(staking_pool.pending_pool_token_withdraw() == 0); sui::test_utils::destroy(staking_pool); sui::test_utils::destroy(staked_sui_1); - scenario.end(); + test.end(); } #[test]