From b66c9539c217fa33d325ad787edef847aa002fe8 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 15 Oct 2021 16:13:09 +0000
Subject: [PATCH] program-test: Fix getting new blockhash post-warp (#20710)
 (#20723)

(cherry picked from commit 0419e6c22e267f123bff59cb0a81f961cbdaf443)

Co-authored-by: Jon Cinque <jon.cinque@gmail.com>
---
 banks-server/src/banks_server.rs |  5 +----
 program-test/src/lib.rs          |  8 ++++++--
 program-test/tests/warp.rs       | 33 +++++++++++++++++++++++++++++++-
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/banks-server/src/banks_server.rs b/banks-server/src/banks_server.rs
index 17d2674610..ee896ad202 100644
--- a/banks-server/src/banks_server.rs
+++ b/banks-server/src/banks_server.rs
@@ -236,10 +236,7 @@ impl Banks for BanksServer {
 
         let blockhash = &transaction.message.recent_blockhash;
         let last_valid_block_height = self
-            .bank_forks
-            .read()
-            .unwrap()
-            .root_bank()
+            .bank(commitment)
             .get_blockhash_last_valid_block_height(blockhash)
             .unwrap();
         let signature = transaction.signatures.get(0).cloned().unwrap_or_default();
diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs
index c1ba815b33..ee730c2c89 100644
--- a/program-test/src/lib.rs
+++ b/program-test/src/lib.rs
@@ -1036,7 +1036,7 @@ impl ProgramTestContext {
         bank_forks.set_root(
             pre_warp_slot,
             &solana_runtime::accounts_background_service::AbsRequestSender::default(),
-            Some(warp_slot),
+            Some(pre_warp_slot),
         );
 
         // warp bank is frozen, so go forward one slot from it
@@ -1049,7 +1049,11 @@ impl ProgramTestContext {
         // Update block commitment cache, otherwise banks server will poll at
         // the wrong slot
         let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
-        w_block_commitment_cache.set_all_slots(pre_warp_slot, warp_slot);
+        // HACK: The root set here should be `pre_warp_slot`, but since we're
+        // in a testing environment, the root bank never updates after a warp.
+        // The ticking thread only updates the working bank, and never the root
+        // bank.
+        w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
 
         let bank = bank_forks.working_bank();
         self.last_blockhash = bank.last_blockhash();
diff --git a/program-test/tests/warp.rs b/program-test/tests/warp.rs
index bd6b365e88..0279146826 100644
--- a/program-test/tests/warp.rs
+++ b/program-test/tests/warp.rs
@@ -2,7 +2,9 @@
 use {
     bincode::deserialize,
     solana_banks_client::BanksClient,
-    solana_program_test::{processor, ProgramTest, ProgramTestContext, ProgramTestError},
+    solana_program_test::{
+        processor, ProgramTest, ProgramTestBanksClientExt, ProgramTestContext, ProgramTestError,
+    },
     solana_sdk::{
         account_info::{next_account_info, AccountInfo},
         clock::Clock,
@@ -413,3 +415,32 @@ async fn stake_merge_immediately_after_activation() {
         .await
         .unwrap();
 }
+
+#[tokio::test]
+async fn get_blockhash_post_warp() {
+    let program_test = ProgramTest::default();
+    let mut context = program_test.start_with_context().await;
+
+    let new_blockhash = context
+        .banks_client
+        .get_new_blockhash(&context.last_blockhash)
+        .await
+        .unwrap()
+        .0;
+    let mut tx = Transaction::new_with_payer(&[], Some(&context.payer.pubkey()));
+    tx.sign(&[&context.payer], new_blockhash);
+    context.banks_client.process_transaction(tx).await.unwrap();
+
+    context.warp_to_slot(10).unwrap();
+
+    let new_blockhash = context
+        .banks_client
+        .get_new_blockhash(&context.last_blockhash)
+        .await
+        .unwrap()
+        .0;
+
+    let mut tx = Transaction::new_with_payer(&[], Some(&context.payer.pubkey()));
+    tx.sign(&[&context.payer], new_blockhash);
+    context.banks_client.process_transaction(tx).await.unwrap();
+}