From 90850fb74308f99468322fc05645cb06c23f47df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Tue, 17 Jan 2023 11:23:32 +0100
Subject: [PATCH] Adjust tests.

---
 programs/sbf/tests/programs.rs | 52 ++++++++++++++++++++++++----------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs
index f23f5111fc5539..2b587f23ffc07c 100644
--- a/programs/sbf/tests/programs.rs
+++ b/programs/sbf/tests/programs.rs
@@ -1822,14 +1822,11 @@ fn test_program_sbf_invoke_in_same_tx_as_deployment() {
     )
     .unwrap();
 
-    // Deployment is invisible to top-level-instructions but visible to CPI instructions
-    for (invoke_instruction, expected) in [
-        (
-            invoke_instruction,
-            Err(TransactionError::ProgramAccountNotFound),
-        ),
-        (indirect_invoke_instruction, Ok(())),
-    ] {
+    // Deployment is invisible to both top-level-instructions and CPI instructions
+    for (index, invoke_instruction) in [invoke_instruction, indirect_invoke_instruction]
+        .into_iter()
+        .enumerate()
+    {
         let mut instructions = deployment_instructions.clone();
         instructions.push(invoke_instruction);
         let tx = Transaction::new(
@@ -1837,11 +1834,18 @@ fn test_program_sbf_invoke_in_same_tx_as_deployment() {
             Message::new(&instructions, Some(&mint_keypair.pubkey())),
             bank.last_blockhash(),
         );
-        let results = execute_transactions(&bank, vec![tx]);
-        if let Err(err) = expected {
-            assert_eq!(results[0].as_ref().unwrap_err(), &err);
+        if index == 0 {
+            let results = execute_transactions(&bank, vec![tx]);
+            assert_eq!(
+                results[0].as_ref().unwrap_err(),
+                &TransactionError::ProgramAccountNotFound,
+            );
         } else {
-            assert!(results[0].is_ok());
+            let (result, _, _) = process_transaction_and_record_inner(&bank, tx);
+            assert_eq!(
+                result.unwrap_err(),
+                TransactionError::InstructionError(2, InstructionError::InvalidAccountData),
+            );
         }
     }
 }
@@ -1887,6 +1891,17 @@ fn test_program_sbf_invoke_in_same_tx_as_redeployment() {
         "solana_sbf_rust_invoke_and_return",
     );
 
+    // Deploy panic program
+    let panic_program_keypair = Keypair::new();
+    load_upgradeable_program(
+        &bank_client,
+        &mint_keypair,
+        &buffer_keypair,
+        &panic_program_keypair,
+        &authority_keypair,
+        "solana_sbf_rust_panic",
+    );
+
     let invoke_instruction =
         Instruction::new_with_bytes(program_id, &[0], vec![AccountMeta::new(clock::id(), false)]);
     let indirect_invoke_instruction = Instruction::new_with_bytes(
@@ -1899,6 +1914,7 @@ fn test_program_sbf_invoke_in_same_tx_as_redeployment() {
     );
 
     // Prepare redeployment
+    let buffer_keypair = Keypair::new();
     load_upgradeable_buffer(
         &bank_client,
         &mint_keypair,
@@ -1912,8 +1928,10 @@ fn test_program_sbf_invoke_in_same_tx_as_redeployment() {
         &authority_keypair.pubkey(),
         &mint_keypair.pubkey(),
     );
+    let panic_instruction =
+        Instruction::new_with_bytes(panic_program_keypair.pubkey(), &[], vec![]);
 
-    // Redeployment is visible to both top-level-instructions and CPI instructions
+    // Redeployment is invisible to both top-level-instructions and CPI instructions
     for invoke_instruction in [invoke_instruction, indirect_invoke_instruction] {
         // Call upgradeable program
         let result =
@@ -1922,7 +1940,11 @@ fn test_program_sbf_invoke_in_same_tx_as_redeployment() {
 
         // Upgrade the program and invoke in same tx
         let message = Message::new(
-            &[redeployment_instruction.clone(), invoke_instruction],
+            &[
+                redeployment_instruction.clone(),
+                invoke_instruction,
+                panic_instruction.clone(), // Make sure the TX fails, so we don't have to deploy another program
+            ],
             Some(&mint_keypair.pubkey()),
         );
         let tx = Transaction::new(
@@ -1933,7 +1955,7 @@ fn test_program_sbf_invoke_in_same_tx_as_redeployment() {
         let (result, _, _) = process_transaction_and_record_inner(&bank, tx);
         assert_eq!(
             result.unwrap_err(),
-            TransactionError::InstructionError(1, InstructionError::ProgramFailedToComplete),
+            TransactionError::InstructionError(2, InstructionError::ProgramFailedToComplete),
         );
     }
 }