Fix tests

runtime/src/bank/partitioned_epoch_rewards/distribution.rs

@@ -211,8 +211,10 @@ mod tests {
         let total_rewards = 1_000_000_000;
         bank.create_epoch_rewards_sysvar(total_rewards, 0, 42);
         let pre_epoch_rewards_account = bank.get_account(&sysvar::epoch_rewards::id()).unwrap();
-        // Expected balance is the starting balance (1) + total_rewards
-        assert_eq!(pre_epoch_rewards_account.lamports(), 1 + total_rewards);
+        let expected_balance =
+            bank.get_minimum_balance_for_rent_exemption(pre_epoch_rewards_account.data().len());
+        // Expected balance is the sysvar rent-exempt balance
+        assert_eq!(pre_epoch_rewards_account.lamports(), expected_balance);
 
         // Set up a partition of rewards to distribute
         let expected_num = 100;
@@ -231,23 +233,18 @@ mod tests {
         bank.distribute_epoch_rewards_in_partition(&all_rewards, 0);
         let post_cap = bank.capitalization();
         let post_epoch_rewards_account = bank.get_account(&sysvar::epoch_rewards::id()).unwrap();
-        // Expected balance is the starting balance (1) + total_rewards
-        let expected_epoch_rewards_sysvar_lamports_remaining =
-            1 + total_rewards - rewards_to_distribute;
 
-        // Assert that epoch rewards sysvar lamports decreases by the distributed rewards
-        assert_eq!(
-            post_epoch_rewards_account.lamports(),
-            expected_epoch_rewards_sysvar_lamports_remaining
-        );
+        // Assert that epoch rewards sysvar lamports balance does not change
+        assert_eq!(post_epoch_rewards_account.lamports(), expected_balance);
 
         let epoch_rewards: sysvar::epoch_rewards::EpochRewards =
             from_account(&post_epoch_rewards_account).unwrap();
         assert_eq!(epoch_rewards.total_rewards, total_rewards);
         assert_eq!(epoch_rewards.distributed_rewards, rewards_to_distribute,);
 
-        // Assert that the bank total capital didn't change
-        assert_eq!(pre_cap, post_cap);
+        // Assert that the bank total capital changed by the amount of rewards
+        // distributed
+        assert_eq!(pre_cap + rewards_to_distribute, post_cap);
     }
 
     /// Test partitioned credits and reward history updates of epoch rewards do cover all the rewards

runtime/src/bank/partitioned_epoch_rewards/mod.rs

@@ -443,19 +443,33 @@ mod tests {
                 } else {
                     assert_eq!(post_cap, pre_cap);
                 }
-            } else if slot == num_slots_in_epoch + 1 || slot == num_slots_in_epoch + 2 {
-                // 1. when curr_slot == num_slots_in_epoch + 1, the 2nd block of epoch 1, reward distribution should happen in this block.
-                // however, all stake rewards are paid at the this block therefore reward_status should have transitioned to inactive. And since
-                // rewards are transferred from epoch_rewards sysvar to stake accounts. The cap should stay the same.
-                // 2. when curr_slot == num_slots_in_epoch+2, the 3rd block of epoch 1. reward distribution should have already completed. Therefore,
-                // reward_status should stay inactive and cap should stay the same.
+            } else if slot == num_slots_in_epoch + 1 {
+                // 1. when curr_slot == num_slots_in_epoch + 1, the 2nd block of
+                // epoch 1, reward distribution should happen in this block.
+                // however, all stake rewards are paid at the this block
+                // therefore reward_status should have transitioned to inactive.
+                // The cap should increase accordingly.
                 assert_matches!(
                     curr_bank.get_reward_interval(),
                     RewardInterval::OutsideInterval
                 );
 
-                assert_eq!(post_cap, pre_cap);
+                let account = curr_bank
+                    .get_account(&solana_sdk::sysvar::epoch_rewards::id())
+                    .unwrap();
+                let epoch_rewards: solana_sdk::sysvar::epoch_rewards::EpochRewards =
+                    solana_sdk::account::from_account(&account).unwrap();
+                assert_eq!(post_cap, pre_cap + epoch_rewards.distributed_rewards);
             } else {
+                // 2. when curr_slot == num_slots_in_epoch+2, the 3rd block of
+                // epoch 1 (or any other slot). reward distribution should have
+                // already completed. Therefore, reward_status should stay
+                // inactive and cap should stay the same.
+                assert_matches!(
+                    curr_bank.get_reward_interval(),
+                    RewardInterval::OutsideInterval
+                );
+
                 // slot is not in rewards, cap should not change
                 assert_eq!(post_cap, pre_cap);
             }
@@ -528,6 +542,14 @@ mod tests {
         // simulate block progress
         for slot in 2..=num_slots_in_epoch + 3 {
             let pre_cap = previous_bank.capitalization();
+
+            let pre_sysvar_account = previous_bank
+                .get_account(&solana_sdk::sysvar::epoch_rewards::id())
+                .unwrap_or_default();
+            let pre_epoch_rewards: solana_sdk::sysvar::epoch_rewards::EpochRewards =
+                solana_sdk::account::from_account(&pre_sysvar_account).unwrap_or_default();
+            let pre_distributed_rewards = pre_epoch_rewards.distributed_rewards;
+
             let curr_bank = Bank::new_from_parent(previous_bank, &Pubkey::default(), slot);
             let post_cap = curr_bank.capitalization();
 
@@ -566,27 +588,53 @@ mod tests {
                 // cap should increase because of new epoch rewards
                 assert!(post_cap > pre_cap);
             } else if slot == num_slots_in_epoch + 1 {
-                // When curr_slot == num_slots_in_epoch + 1, the 2nd block of epoch 1, reward distribution should happen in this block.
-                // however, since rewards are transferred from epoch_rewards sysvar to stake accounts. The cap should stay the same.
+                // When curr_slot == num_slots_in_epoch + 1, the 2nd block of
+                // epoch 1, reward distribution should happen in this block. The
+                // cap should increase accordingly.
                 assert_matches!(
                     curr_bank.get_reward_interval(),
                     RewardInterval::InsideInterval
                 );
 
-                assert_eq!(post_cap, pre_cap);
-            } else if slot == num_slots_in_epoch + 2 || slot == num_slots_in_epoch + 3 {
-                // 1. when curr_slot == num_slots_in_epoch + 2, the 3nd block of epoch 1, reward distribution should happen in this block.
-                // however, all stake rewards are paid at the this block therefore reward_status should have transitioned to inactive. And since
-                // rewards are transferred from epoch_rewards sysvar to stake accounts. The cap should stay the same.
-                // 2. when curr_slot == num_slots_in_epoch+2, the 3rd block of epoch 1. reward distribution should have already completed. Therefore,
-                // reward_status should stay inactive and cap should stay the same.
+                let account = curr_bank
+                    .get_account(&solana_sdk::sysvar::epoch_rewards::id())
+                    .unwrap();
+                let epoch_rewards: solana_sdk::sysvar::epoch_rewards::EpochRewards =
+                    solana_sdk::account::from_account(&account).unwrap();
+                assert_eq!(
+                    post_cap,
+                    pre_cap + epoch_rewards.distributed_rewards - pre_distributed_rewards
+                );
+            } else if slot == num_slots_in_epoch + 2 {
+                // When curr_slot == num_slots_in_epoch + 2, the 3nd block of
+                // epoch 1, reward distribution should happen in this block.
+                // however, all stake rewards are paid at the this block
+                // therefore reward_status should have transitioned to inactive.
+                // The cap should increase accordingly.
                 assert_matches!(
                     curr_bank.get_reward_interval(),
                     RewardInterval::OutsideInterval
                 );
 
-                assert_eq!(post_cap, pre_cap);
+                let account = curr_bank
+                    .get_account(&solana_sdk::sysvar::epoch_rewards::id())
+                    .unwrap();
+                let epoch_rewards: solana_sdk::sysvar::epoch_rewards::EpochRewards =
+                    solana_sdk::account::from_account(&account).unwrap();
+                assert_eq!(
+                    post_cap,
+                    pre_cap + epoch_rewards.distributed_rewards - pre_distributed_rewards
+                );
             } else {
+                // When curr_slot == num_slots_in_epoch + 3, the 4th block of
+                // epoch 1 (or any other slot). reward distribution should have
+                // already completed. Therefore, reward_status should stay
+                // inactive and cap should stay the same.
+                assert_matches!(
+                    curr_bank.get_reward_interval(),
+                    RewardInterval::OutsideInterval
+                );
+
                 // slot is not in rewards, cap should not change
                 assert_eq!(post_cap, pre_cap);
             }

runtime/src/bank/partitioned_epoch_rewards/sysvar.rs

@@ -133,18 +133,17 @@ mod tests {
 
         bank.create_epoch_rewards_sysvar(total_rewards, 10, 42);
         let account = bank.get_account(&sysvar::epoch_rewards::id()).unwrap();
-        // Expected balance is the starting balance (1) + total_rewards -
-        // distributed_rewards (10)
-        assert_eq!(account.lamports(), 1 + total_rewards - 10);
+        let expected_balance = bank.get_minimum_balance_for_rent_exemption(account.data().len());
+        // Expected balance is the sysvar rent-exempt balance
+        assert_eq!(account.lamports(), expected_balance);
         let epoch_rewards: sysvar::epoch_rewards::EpochRewards = from_account(&account).unwrap();
         assert_eq!(epoch_rewards, expected_epoch_rewards);
 
         // make a distribution from epoch rewards sysvar
         bank.update_epoch_rewards_sysvar(10);
         let account = bank.get_account(&sysvar::epoch_rewards::id()).unwrap();
-        // Expected balance is the starting balance (1) + total_rewards -
-        // distributed_rewards (10)
-        assert_eq!(account.lamports(), 1 + total_rewards - 20);
+        // Balance should not change
+        assert_eq!(account.lamports(), expected_balance);
         let epoch_rewards: sysvar::epoch_rewards::EpochRewards = from_account(&account).unwrap();
         let expected_epoch_rewards = sysvar::epoch_rewards::EpochRewards {
             distribution_starting_block_height: 42,