comments, one more test case

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

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

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);

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

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<SuiSystemState>();
+        let mut test = test_scenario::begin(VALIDATOR_ADDR_1); // val 1 has 100 staked sui
+        let mut sui_system = test.take_shared<SuiSystemState>();
+        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<StakedSui>(STAKER_ADDR_1);
+        let staked_sui = test.take_from_address<StakedSui>(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<SuiSystemState>();
+        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<StakedSui>(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<StakedSui>(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();
     }
 }

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]