program-test: Fix getting new blockhash post-warp (#20710)

banks-server/src/banks_server.rs

@@ -240,10 +240,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();

program-test/src/lib.rs

@@ -1038,7 +1038,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
@@ -1051,7 +1051,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();

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();
+}