Update test to new StakeProgram functionality

runtime/src/bank/partitioned_epoch_rewards/mod.rs

@@ -189,10 +189,14 @@ mod tests {
         },
         solana_program_runtime::runtime_config::RuntimeConfig,
         solana_sdk::{
+            account::Account,
             epoch_schedule::EpochSchedule,
             native_token::LAMPORTS_PER_SOL,
             signature::Signer,
+            signer::keypair::Keypair,
+            stake::instruction::StakeError,
             system_transaction,
+            transaction::Transaction,
             vote::state::{VoteStateVersions, MAX_LOCKOUT_HISTORY},
         },
         solana_vote_program::{vote_state, vote_transaction},
@@ -642,28 +646,54 @@ mod tests {
         }
     }
 
-    /// Test that program execution that involves stake accounts should fail during reward period.
-    /// Any programs, which result in stake account changes, will throw `ProgramExecutionTemporarilyRestricted` error when
-    /// in reward period.
+    /// Test that program execution that attempts to mutate a stake account
+    /// incorrectly should fail during reward period. A credit should succeed,
+    /// but a withdrawal shoudl fail.
     #[test]
     fn test_program_execution_restricted_for_stake_account_in_reward_period() {
-        use solana_sdk::transaction::TransactionError::ProgramExecutionTemporarilyRestricted;
+        use solana_sdk::transaction::TransactionError::InstructionError;
 
         let validator_vote_keypairs = ValidatorVoteKeypairs::new_rand();
         let validator_keypairs = vec![&validator_vote_keypairs];
-        let GenesisConfigInfo { genesis_config, .. } = create_genesis_config_with_vote_accounts(
+        let GenesisConfigInfo {
+            mut genesis_config,
+            mint_keypair,
+            ..
+        } = create_genesis_config_with_vote_accounts(
             1_000_000_000,
             &validator_keypairs,
             vec![1_000_000_000; 1],
         );
 
-        let node_key = &validator_keypairs[0].node_keypair;
-        let stake_key = &validator_keypairs[0].stake_keypair;
+        // Add stake account to try to mutate
+        let vote_key = validator_keypairs[0].vote_keypair.pubkey();
+        let vote_account = genesis_config
+            .accounts
+            .iter()
+            .find(|(&address, _)| address == vote_key)
+            .map(|(_, account)| account)
+            .unwrap()
+            .clone();
+
+        let new_stake_signer = Keypair::new();
+        let new_stake_address = new_stake_signer.pubkey();
+        let new_stake_account = Account::from(solana_stake_program::stake_state::create_account(
+            &new_stake_address,
+            &vote_key,
+            &vote_account.into(),
+            &genesis_config.rent,
+            2_000_000_000,
+        ));
+        genesis_config
+            .accounts
+            .extend(vec![(new_stake_address, new_stake_account)]);
 
         let (mut previous_bank, bank_forks) = Bank::new_with_bank_forks_for_tests(&genesis_config);
         let num_slots_in_epoch = previous_bank.get_slots_in_epoch(previous_bank.epoch());
         assert_eq!(num_slots_in_epoch, 32);
 
+        let transfer_amount = 5_000;
+
         for slot in 1..=num_slots_in_epoch + 2 {
             let bank = new_bank_from_parent_with_bank_forks(
                 bank_forks.as_ref(),
@@ -685,25 +715,46 @@ mod tests {
             );
             bank.process_transaction(&vote).unwrap();
 
-            // Insert a transfer transaction from node account to stake account
-            let tx = system_transaction::transfer(
-                node_key,
-                &stake_key.pubkey(),
-                1,
+            // Insert a transfer transaction from the mint to new stake account
+            let system_tx = system_transaction::transfer(
+                &mint_keypair,
+                &new_stake_address,
+                transfer_amount,
+                bank.last_blockhash(),
+            );
+            let system_result = bank.process_transaction(&system_tx);
+
+            // Credits should always succeed
+            assert!(system_result.is_ok());
+
+            // Attempt to withdraw from new stake account to the mint
+            let stake_ix = solana_sdk::stake::instruction::withdraw(
+                &new_stake_address,
+                &new_stake_address,
+                &mint_keypair.pubkey(),
+                transfer_amount,
+                None,
+            );
+            let stake_tx = Transaction::new_signed_with_payer(
+                &[stake_ix],
+                Some(&mint_keypair.pubkey()),
+                &[&mint_keypair, &new_stake_signer],
                 bank.last_blockhash(),
             );
-            let r = bank.process_transaction(&tx);
+            let stake_result = bank.process_transaction(&stake_tx);
 
             if slot == num_slots_in_epoch {
-                // When the bank is at the beginning of the new epoch, i.e. slot 32,
-                // ProgramExecutionTemporarilyRestricted should be thrown for the transfer transaction.
+                // When the bank is at the beginning of the new epoch, i.e. slot
+                // 32, StakeError::EpochRewardsActive should be thrown for
+                // actions like StakeInstruction::Withdraw
                 assert_eq!(
-                    r,
-                    Err(ProgramExecutionTemporarilyRestricted { account_index: 1 })
+                    stake_result,
+                    Err(InstructionError(0, StakeError::EpochRewardsActive.into()))
                 );
             } else {
-                // When the bank is outside of reward interval, the transfer transaction should not be affected and will succeed.
-                assert!(r.is_ok());
+                // When the bank is outside of reward interval, the withdraw
+                // transaction should not be affected and will succeed.
+                assert!(stake_result.is_ok());
             }
 
             // Push a dummy blockhash, so that the latest_blockhash() for the transfer transaction in each