harness: don't stub program account if it was provided (#162)

Author: buffalojoec

Makefile

@@ -82,6 +82,10 @@ clippy-fix:
 check-features:
 	@cargo hack check --feature-powerset --no-dev-deps
 
+build:
+	@$(MAKE) build-test-programs
+	@cargo build
+
 test:
 	@$(MAKE) build-test-programs
 	@cargo test --all-features

harness/tests/system_program.rs

@@ -1,6 +1,6 @@
 use {
     mollusk_svm::{result::Check, Mollusk},
-    solana_account::Account,
+    solana_account::{Account, WritableAccount},
     solana_instruction::error::InstructionError,
     solana_pubkey::Pubkey,
     solana_system_program::system_processor::DEFAULT_COMPUTE_UNITS,
@@ -103,3 +103,36 @@ fn test_transfer_bad_owner() {
 
     Mollusk::default().process_and_validate_instruction(&instruction, &accounts, &checks);
 }
+
+#[test]
+fn test_transfer_swap_program_account() {
+    let sender = Pubkey::new_unique();
+    let recipient = Pubkey::new_unique();
+
+    let base_lamports = 100_000_000u64;
+    let transfer_amount = 42_000u64;
+
+    let instruction =
+        solana_system_interface::instruction::transfer(&sender, &recipient, transfer_amount);
+
+    // Provide a custom program account instead of letting it be stubbed.
+    let mut program_account = Account::new(1_000_000, 0, &solana_sdk_ids::native_loader::id());
+    program_account.set_executable(true);
+
+    let accounts = [
+        (
+            sender,
+            Account::new(base_lamports, 0, &solana_sdk_ids::system_program::id()),
+        ),
+        (
+            recipient,
+            Account::new(base_lamports, 0, &solana_sdk_ids::system_program::id()),
+        ),
+        (solana_sdk_ids::system_program::id(), program_account),
+    ];
+
+    // The test verifies that providing a custom program account does not panic.
+    // Before the fix, Mollusk would always stub the program account, ignoring the
+    // provided one. Now it uses the provided account if available.
+    let _result = Mollusk::default().process_instruction(&instruction, &accounts);
+}

keys/src/accounts.rs

@@ -2,7 +2,7 @@
 
 use {
     crate::keys::KeyMap,
-    mollusk_svm_error::error::{MolluskError, MolluskPanic},
+    mollusk_svm_error::error::MolluskError,
     solana_account::{Account, AccountSharedData},
     solana_instruction::Instruction,
     solana_pubkey::Pubkey,
@@ -56,17 +56,22 @@ pub fn compile_transaction_accounts_for_instruction(
     key_map
         .keys()
         .map(|key| {
-            if let Some(stub_out_program_account) = &stub_out_program_account {
-                if instruction.program_id == *key {
-                    return (*key, stub_out_program_account().into());
-                }
-            }
             let account = accounts
                 .iter()
                 .find(|(k, _)| k == key)
-                .map(|(_, account)| AccountSharedData::from(account.clone()))
-                .or_panic_with(MolluskError::AccountMissing(key));
-            (*key, account)
+                .map(|(_, account)| AccountSharedData::from(account.clone()));
+
+            if let Some(account) = account {
+                (*key, account)
+            } else if let Some(stub_out_program_account) = &stub_out_program_account {
+                if instruction.program_id == *key {
+                    (*key, stub_out_program_account().into())
+                } else {
+                    panic!("{}", MolluskError::AccountMissing(key))
+                }
+            } else {
+                panic!("{}", MolluskError::AccountMissing(key))
+            }
         })
         .collect()
 }
@@ -80,17 +85,22 @@ pub fn compile_transaction_accounts(
     key_map
         .keys()
         .map(|key| {
-            if let Some(stub_out_program_account) = &stub_out_program_account {
-                if instructions.iter().any(|ix| ix.program_id == *key) {
-                    return (*key, stub_out_program_account().into());
-                }
-            }
             let account = accounts
                 .iter()
                 .find(|(k, _)| k == key)
-                .map(|(_, account)| AccountSharedData::from(account.clone()))
-                .or_panic_with(MolluskError::AccountMissing(key));
-            (*key, account)
+                .map(|(_, account)| AccountSharedData::from(account.clone()));
+
+            if let Some(account) = account {
+                (*key, account)
+            } else if let Some(stub_out_program_account) = &stub_out_program_account {
+                if instructions.iter().any(|ix| ix.program_id == *key) {
+                    (*key, stub_out_program_account().into())
+                } else {
+                    panic!("{}", MolluskError::AccountMissing(key))
+                }
+            } else {
+                panic!("{}", MolluskError::AccountMissing(key))
+            }
         })
         .collect()
 }