Update account helpers (#21)

buffalojoec · web-flow · commit be0350a3f43a · 2024-10-09T13:55:58.000+08:00
diff --git a/harness/src/lib.rs b/harness/src/lib.rs
@@ -83,7 +83,7 @@ impl Default for Mollusk {
              solana_runtime::message_processor=debug,\
              solana_runtime::system_instruction_processor=trace",
         );
-        let (program_id, program_account) = program::system_program();
+        let (program_id, program_account) = program::keyed_account_for_system_program();
         Self {
             compute_budget: ComputeBudget::default(),
             feature_set: FeatureSet::all_enabled(),
@@ -105,7 +105,7 @@ impl Mollusk {
     pub fn new(program_id: &Pubkey, program_name: &'static str) -> Self {
         let mut mollusk = Self {
             program_id: *program_id,
-            program_account: program::program_account(program_id),
+            program_account: program::create_program_account_loader_v3(program_id),
             ..Default::default()
         };
         mollusk.add_program(program_id, program_name);
diff --git a/harness/src/program.rs b/harness/src/program.rs
@@ -115,7 +115,11 @@ static BUILTINS: &[Builtin] = &[
     /* ... */
 ];
 
-fn builtin_program_account(program_id: &Pubkey, name: &str) -> (Pubkey, AccountSharedData) {
+/// Create a key and account for a builtin program.
+pub fn create_keyed_account_for_builtin_program(
+    program_id: &Pubkey,
+    name: &str,
+) -> (Pubkey, AccountSharedData) {
     let data = name.as_bytes().to_vec();
     let lamports = Rent::default().minimum_balance(data.len());
     let account = AccountSharedData::from(Account {
@@ -129,19 +133,24 @@ fn builtin_program_account(program_id: &Pubkey, name: &str) -> (Pubkey, AccountS
 }
 
 /// Get the key and account for the system program.
-pub fn system_program() -> (Pubkey, AccountSharedData) {
-    builtin_program_account(&BUILTINS[0].program_id, BUILTINS[0].name)
+pub fn keyed_account_for_system_program() -> (Pubkey, AccountSharedData) {
+    create_keyed_account_for_builtin_program(&BUILTINS[0].program_id, BUILTINS[0].name)
 }
 
-/// Get the key and account for the BPF Loader Upgradeable program.
-pub fn bpf_loader_upgradeable_program() -> (Pubkey, AccountSharedData) {
-    builtin_program_account(&BUILTINS[1].program_id, BUILTINS[1].name)
+/// Get the key and account for the BPF Loader v2 program.
+pub fn keyed_account_for_bpf_loader_v2_program() -> (Pubkey, AccountSharedData) {
+    create_keyed_account_for_builtin_program(&BUILTINS[1].program_id, BUILTINS[1].name)
+}
+
+/// Get the key and account for the BPF Loader v3 (Upgradeable) program.
+pub fn keyed_account_for_bpf_loader_v3_program() -> (Pubkey, AccountSharedData) {
+    create_keyed_account_for_builtin_program(&BUILTINS[1].program_id, BUILTINS[1].name)
 }
 
 /* ... */
 
 /// Create a BPF Loader 2 program account.
-pub fn program_account_loader_2(elf: &[u8]) -> AccountSharedData {
+pub fn create_program_account_loader_v2(elf: &[u8]) -> AccountSharedData {
     let lamports = Rent::default().minimum_balance(elf.len());
     AccountSharedData::from(Account {
         lamports,
@@ -152,8 +161,8 @@ pub fn program_account_loader_2(elf: &[u8]) -> AccountSharedData {
     })
 }
 
-/// Create a BPF Loader Upgradeable program account.
-pub fn program_account(program_id: &Pubkey) -> AccountSharedData {
+/// Create a BPF Loader v3 (Upgradeable) program account.
+pub fn create_program_account_loader_v3(program_id: &Pubkey) -> AccountSharedData {
     let programdata_address =
         Pubkey::find_program_address(&[program_id.as_ref()], &bpf_loader_upgradeable::id()).0;
     let data = bincode::serialize(&UpgradeableLoaderState::Program {
@@ -170,8 +179,8 @@ pub fn program_account(program_id: &Pubkey) -> AccountSharedData {
     })
 }
 
-/// Create a BPF Loader Upgradeable program data account.
-pub fn program_data_account(elf: &[u8]) -> AccountSharedData {
+/// Create a BPF Loader v3 (Upgradeable) program data account.
+pub fn create_program_data_account_loader_v3(elf: &[u8]) -> AccountSharedData {
     let data = {
         let elf_offset = UpgradeableLoaderState::size_of_programdata_metadata();
         let data_len = elf_offset + elf.len();
@@ -197,10 +206,16 @@ pub fn program_data_account(elf: &[u8]) -> AccountSharedData {
     })
 }
 
-/// Create a BPF Loader Upgradeable program and program data account.
+/// Create a BPF Loader v3 (Upgradeable) program and program data account.
 ///
 /// Returns a tuple, where the first element is the program account and the
 /// second element is the program data account.
-pub fn program_accounts(program_id: &Pubkey, elf: &[u8]) -> (AccountSharedData, AccountSharedData) {
-    (program_account(program_id), program_data_account(elf))
+pub fn create_program_account_pair_loader_v3(
+    program_id: &Pubkey,
+    elf: &[u8],
+) -> (AccountSharedData, AccountSharedData) {
+    (
+        create_program_account_loader_v3(program_id),
+        create_program_data_account_loader_v3(elf),
+    )
 }
diff --git a/harness/src/sysvar.rs b/harness/src/sysvar.rs
@@ -3,14 +3,16 @@
 use {
     solana_program_runtime::sysvar_cache::SysvarCache,
     solana_sdk::{
+        account::AccountSharedData,
         clock::{Clock, Slot},
         epoch_rewards::EpochRewards,
         epoch_schedule::EpochSchedule,
         hash::Hash,
+        pubkey::Pubkey,
         rent::Rent,
         slot_hashes::SlotHashes,
         stake_history::StakeHistory,
-        sysvar::{last_restart_slot::LastRestartSlot, SysvarId},
+        sysvar::{self, last_restart_slot::LastRestartSlot, Sysvar, SysvarId},
     },
 };
 
@@ -29,6 +31,50 @@ pub struct Sysvars {
 }
 
 impl Sysvars {
+    fn sysvar_account<T: SysvarId + Sysvar>(&self, sysvar: &T) -> (Pubkey, AccountSharedData) {
+        let data = bincode::serialize::<T>(sysvar).unwrap();
+        let space = data.len();
+        let lamports = self.rent.minimum_balance(space);
+        let mut account = AccountSharedData::new(lamports, space, &sysvar::id());
+        account.set_data_from_slice(&data);
+        (T::id(), account)
+    }
+
+    /// Get the key and account for the clock sysvar.
+    pub fn keyed_account_for_clock_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.clock)
+    }
+
+    /// Get the key and account for the epoch rewards sysvar.
+    pub fn keyed_account_for_epoch_rewards_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.epoch_rewards)
+    }
+
+    /// Get the key and account for the epoch schedule sysvar.
+    pub fn keyed_account_for_epoch_schedule_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.epoch_schedule)
+    }
+
+    /// Get the key and account for the last restart slot sysvar.
+    pub fn keyed_account_for_last_restart_slot_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.last_restart_slot)
+    }
+
+    /// Get the key and account for the rent sysvar.
+    pub fn keyed_account_for_rent_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.rent)
+    }
+
+    /// Get the key and account for the slot hashes sysvar.
+    pub fn keyed_account_for_slot_hashes_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.slot_hashes)
+    }
+
+    /// Get the key and account for the stake history sysvar.
+    pub fn keyed_account_for_stake_history_sysvar(&self) -> (Pubkey, AccountSharedData) {
+        self.sysvar_account(&self.stake_history)
+    }
+
     /// Warp the test environment to a slot by updating sysvars.
     pub fn warp_to_slot(&mut self, slot: Slot) {
         // First update `Clock`.
diff --git a/harness/tests/bpf_program.rs b/harness/tests/bpf_program.rs
@@ -1,6 +1,6 @@
 use {
     mollusk_svm::{
-        program::{program_account, system_program},
+        program::{create_program_account_loader_v3, keyed_account_for_system_program},
         result::Check,
         Mollusk,
     },
@@ -128,7 +128,7 @@ fn test_transfer() {
             &[
                 (payer, payer_account.clone()),
                 (recipient, recipient_account.clone()),
-                system_program(),
+                keyed_account_for_system_program(),
             ],
             &[
                 Check::err(ProgramError::MissingRequiredSignature),
@@ -144,7 +144,7 @@ fn test_transfer() {
             &[
                 (payer, AccountSharedData::default()),
                 (recipient, recipient_account.clone()),
-                system_program(),
+                keyed_account_for_system_program(),
             ],
             &[
                 Check::err(ProgramError::Custom(
@@ -161,7 +161,7 @@ fn test_transfer() {
         &[
             (payer, payer_account.clone()),
             (recipient, recipient_account.clone()),
-            system_program(),
+            keyed_account_for_system_program(),
         ],
         &[
             Check::success(),
@@ -207,7 +207,7 @@ fn test_close_account() {
             &[
                 (key, account.clone()),
                 (incinerator::id(), AccountSharedData::default()),
-                system_program(),
+                keyed_account_for_system_program(),
             ],
             &[
                 Check::err(ProgramError::MissingRequiredSignature),
@@ -222,7 +222,7 @@ fn test_close_account() {
         &[
             (key, account.clone()),
             (incinerator::id(), AccountSharedData::default()),
-            system_program(),
+            keyed_account_for_system_program(),
         ],
         &[
             Check::success(),
@@ -287,7 +287,7 @@ fn test_cpi() {
                 (key, account.clone()),
                 (
                     cpi_target_program_id,
-                    program_account(&cpi_target_program_id),
+                    create_program_account_loader_v3(&cpi_target_program_id),
                 ),
             ],
             &[
@@ -312,7 +312,7 @@ fn test_cpi() {
                 (key, account.clone()),
                 (
                     cpi_target_program_id,
-                    program_account(&cpi_target_program_id),
+                    create_program_account_loader_v3(&cpi_target_program_id),
                 ),
             ],
             &[
@@ -336,7 +336,7 @@ fn test_cpi() {
                 (key, account.clone()),
                 (
                     cpi_target_program_id,
-                    program_account(&cpi_target_program_id),
+                    create_program_account_loader_v3(&cpi_target_program_id),
                 ),
             ],
             &[
@@ -353,7 +353,7 @@ fn test_cpi() {
             (key, account.clone()),
             (
                 cpi_target_program_id,
-                program_account(&cpi_target_program_id),
+                create_program_account_loader_v3(&cpi_target_program_id),
             ),
         ],
         &[
