From f33c6511149402690f2b1b9cdbd393d5e017fc62 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Tue, 21 Sep 2021 17:04:16 +0000
Subject: [PATCH] Minor ProgramTest improvements (backport #20051) (#20054)

* Add stable_log output when a program is loaded as native code instead of BPF

(cherry picked from commit 34f5020457d93f8443e616f402c9f3b23916062e)

* Add ProgramTest::add_builtin_program()

This permits the unit testing of builtin programs in the ProgramTest environment

(cherry picked from commit 830ca369f1880f94dd6e27e438b605def37bcb41)

Co-authored-by: Michael Vines <mvines@gmail.com>
---
 program-test/src/lib.rs | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs
index f839014024..c1ba815b33 100644
--- a/program-test/src/lib.rs
+++ b/program-test/src/lib.rs
@@ -103,6 +103,9 @@ pub fn builtin_process_instruction(
 ) -> Result<(), InstructionError> {
     set_invoke_context(invoke_context);
 
+    let logger = invoke_context.get_logger();
+    stable_log::program_invoke(&logger, program_id, invoke_context.invoke_depth());
+
     let keyed_accounts = invoke_context.get_keyed_accounts()?;
 
     // Copy all the accounts into a HashMap to ensure there are no duplicates
@@ -151,7 +154,12 @@ pub fn builtin_process_instruction(
         .collect();
 
     // Execute the program
-    process_instruction(program_id, &account_infos, input).map_err(u64::from)?;
+    process_instruction(program_id, &account_infos, input).map_err(|err| {
+        let err = u64::from(err);
+        stable_log::program_failure(&logger, program_id, &err.into());
+        err
+    })?;
+    stable_log::program_success(&logger, program_id);
 
     // Commit AccountInfo changes back into KeyedAccounts
     for keyed_account in keyed_accounts {
@@ -706,6 +714,20 @@ impl ProgramTest {
         }
     }
 
+    /// Add a builtin program to the test environment.
+    ///
+    /// Note that builtin programs are responsible for their own `stable_log` output.
+    pub fn add_builtin_program(
+        &mut self,
+        program_name: &str,
+        program_id: Pubkey,
+        process_instruction: ProcessInstructionWithContext,
+    ) {
+        info!("\"{}\" builtin program", program_name);
+        self.builtins
+            .push(Builtin::new(program_name, program_id, process_instruction));
+    }
+
     fn setup_bank(
         &self,
     ) -> (