Skip to main content

solana_program_test/
lib.rs

1#![cfg(feature = "agave-unstable-api")]
2//! The solana-program-test provides a BanksClient-based test framework SBF programs
3#![allow(clippy::arithmetic_side_effects)]
4
5// Export tokio for test clients
6pub use tokio;
7use {
8    agave_feature_set::{
9        FEATURE_NAMES, FeatureSet, increase_cpi_account_info_limit, raise_cpi_nesting_limit_to_8,
10    },
11    async_trait::async_trait,
12    base64::{Engine, prelude::BASE64_STANDARD},
13    chrono_humanize::{Accuracy, HumanTime, Tense},
14    log::*,
15    solana_account::{
16        Account, AccountSharedData, ReadableAccount, create_account_shared_data_for_test,
17        state_traits::StateMut,
18    },
19    solana_account_info::AccountInfo,
20    solana_accounts_db::accounts_db::ACCOUNTS_DB_CONFIG_FOR_TESTING,
21    solana_banks_client::start_client,
22    solana_banks_server::banks_server::start_local_server,
23    solana_clock::{Clock, Epoch, Slot},
24    solana_cluster_type::ClusterType,
25    solana_compute_budget::compute_budget::{ComputeBudget, SVMTransactionExecutionCost},
26    solana_epoch_rewards::EpochRewards,
27    solana_epoch_schedule::EpochSchedule,
28    solana_fee_calculator::{DEFAULT_TARGET_LAMPORTS_PER_SIGNATURE, FeeRateGovernor},
29    solana_genesis_config::GenesisConfig,
30    solana_hash::Hash,
31    solana_instruction::{
32        Instruction,
33        error::{InstructionError, UNSUPPORTED_SYSVAR},
34    },
35    solana_keypair::Keypair,
36    solana_native_token::LAMPORTS_PER_SOL,
37    solana_poh_config::PohConfig,
38    solana_program_binaries as programs,
39    solana_program_entrypoint::{SUCCESS, deserialize},
40    solana_program_error::{ProgramError, ProgramResult},
41    solana_program_runtime::{
42        invoke_context::BuiltinFunctionWithContext, loaded_programs::ProgramCacheEntry,
43        serialization::serialize_parameters, stable_log, sysvar_cache::SysvarCache,
44    },
45    solana_pubkey::Pubkey,
46    solana_rent::Rent,
47    solana_runtime::{
48        bank::Bank,
49        bank_forks::BankForks,
50        commitment::BlockCommitmentCache,
51        genesis_utils::{GenesisConfigInfo, create_genesis_config_with_leader_ex},
52        runtime_config::RuntimeConfig,
53    },
54    solana_signer::Signer,
55    solana_svm_log_collector::ic_msg,
56    solana_svm_timings::ExecuteTimings,
57    solana_sysvar::{SysvarSerialize, last_restart_slot::LastRestartSlot},
58    solana_sysvar_id::SysvarId,
59    solana_vote_program::vote_state::{VoteStateV4, VoteStateVersions},
60    std::{
61        cell::RefCell,
62        collections::{HashMap, HashSet},
63        fs::File,
64        io::{self, Read},
65        mem::transmute,
66        panic::AssertUnwindSafe,
67        path::{Path, PathBuf},
68        ptr,
69        sync::{
70            Arc, RwLock,
71            atomic::{AtomicBool, Ordering},
72        },
73        time::{Duration, Instant},
74    },
75    thiserror::Error,
76    tokio::task::JoinHandle,
77};
78// Export types so test clients can limit their solana crate dependencies
79pub use {
80    solana_banks_client::{BanksClient, BanksClientError},
81    solana_banks_interface::BanksTransactionResultWithMetadata,
82    solana_program_runtime::invoke_context::InvokeContext,
83    solana_sbpf::{
84        error::EbpfError,
85        vm::{EbpfVm, get_runtime_environment_key},
86    },
87    solana_transaction_context::IndexOfAccount,
88};
89
90/// Errors from the program test environment
91#[derive(Error, Debug, PartialEq, Eq)]
92pub enum ProgramTestError {
93    /// The chosen warp slot is not in the future, so warp is not performed
94    #[error("Warp slot not in the future")]
95    InvalidWarpSlot,
96}
97
98thread_local! {
99    static INVOKE_CONTEXT: RefCell<Option<usize>> = const { RefCell::new(None) };
100}
101fn set_invoke_context(new: &mut InvokeContext) {
102    INVOKE_CONTEXT.with(|invoke_context| unsafe {
103        invoke_context.replace(Some(transmute::<&mut InvokeContext, usize>(new)))
104    });
105}
106fn get_invoke_context<'a, 'b>() -> &'a mut InvokeContext<'b, 'b> {
107    let ptr = INVOKE_CONTEXT.with(|invoke_context| match *invoke_context.borrow() {
108        Some(val) => val,
109        None => panic!("Invoke context not set!"),
110    });
111    unsafe { &mut *ptr::with_exposed_provenance_mut(ptr) }
112}
113
114pub fn invoke_builtin_function(
115    builtin_function: solana_program_entrypoint::ProcessInstruction,
116    invoke_context: &mut InvokeContext,
117) -> Result<u64, Box<dyn std::error::Error>> {
118    set_invoke_context(invoke_context);
119
120    let transaction_context = &invoke_context.transaction_context;
121    let instruction_context = transaction_context.get_current_instruction_context()?;
122    let instruction_account_indices = 0..instruction_context.get_number_of_instruction_accounts();
123
124    // mock builtin program must consume units
125    invoke_context.consume_checked(1)?;
126
127    let log_collector = invoke_context.get_log_collector();
128    let program_id = instruction_context.get_program_key()?;
129    stable_log::program_invoke(
130        &log_collector,
131        program_id,
132        invoke_context.get_stack_height(),
133    );
134
135    // Copy indices_in_instruction into a HashSet to ensure there are no duplicates
136    let deduplicated_indices: HashSet<IndexOfAccount> = instruction_account_indices.collect();
137
138    // Serialize entrypoint parameters with SBF ABI
139    let (mut parameter_bytes, _regions, _account_lengths, _instruction_data_offset) =
140        serialize_parameters(
141            &instruction_context,
142            false, // There is no VM so stricter_abi_and_runtime_constraints can not be implemented here
143            false, // There is no VM so account_data_direct_mapping can not be implemented here
144        )?;
145
146    // Deserialize data back into instruction params
147    let (program_id, account_infos, input) =
148        unsafe { deserialize(&mut parameter_bytes.as_slice_mut()[0] as *mut u8) };
149
150    // Execute the program
151    match std::panic::catch_unwind(AssertUnwindSafe(|| {
152        builtin_function(program_id, &account_infos, input)
153    })) {
154        Ok(program_result) => {
155            program_result.map_err(|program_error| {
156                let err = InstructionError::from(u64::from(program_error));
157                stable_log::program_failure(&log_collector, program_id, &err);
158                let err: Box<dyn std::error::Error> = Box::new(err);
159                err
160            })?;
161        }
162        Err(_panic_error) => {
163            let err = InstructionError::ProgramFailedToComplete;
164            stable_log::program_failure(&log_collector, program_id, &err);
165            let err: Box<dyn std::error::Error> = Box::new(err);
166            Err(err)?;
167        }
168    };
169
170    stable_log::program_success(&log_collector, program_id);
171
172    // Lookup table for AccountInfo
173    let account_info_map: HashMap<_, _> = account_infos.into_iter().map(|a| (a.key, a)).collect();
174
175    // Re-fetch the instruction context. The previous reference may have been
176    // invalidated due to the `set_invoke_context` in a CPI.
177    let transaction_context = &invoke_context.transaction_context;
178    let instruction_context = transaction_context.get_current_instruction_context()?;
179
180    // Commit AccountInfo changes back into KeyedAccounts
181    for i in deduplicated_indices.into_iter() {
182        let mut borrowed_account = instruction_context.try_borrow_instruction_account(i)?;
183        if borrowed_account.is_writable() {
184            if let Some(account_info) = account_info_map.get(borrowed_account.get_key()) {
185                if borrowed_account.get_lamports() != account_info.lamports() {
186                    borrowed_account.set_lamports(account_info.lamports())?;
187                }
188
189                if borrowed_account
190                    .can_data_be_resized(account_info.data_len())
191                    .is_ok()
192                {
193                    borrowed_account.set_data_from_slice(&account_info.data.borrow())?;
194                }
195                if borrowed_account.get_owner() != account_info.owner {
196                    borrowed_account.set_owner(account_info.owner.as_ref())?;
197                }
198            }
199        }
200    }
201
202    Ok(0)
203}
204
205/// Converts a `solana-program`-style entrypoint into the runtime's entrypoint style, for
206/// use with `ProgramTest::add_program`
207#[macro_export]
208macro_rules! processor {
209    ($builtin_function:expr) => {
210        Some(|vm, _arg0, _arg1, _arg2, _arg3, _arg4| {
211            let vm = unsafe {
212                &mut *((vm as *mut u64).offset(-($crate::get_runtime_environment_key() as isize))
213                    as *mut $crate::EbpfVm<$crate::InvokeContext>)
214            };
215            vm.program_result =
216                $crate::invoke_builtin_function($builtin_function, vm.context_object_pointer)
217                    .map_err(|err| $crate::EbpfError::SyscallError(err))
218                    .into();
219        })
220    };
221}
222
223fn get_sysvar<T: Default + SysvarSerialize + Sized + serde::de::DeserializeOwned + Clone>(
224    sysvar: Result<Arc<T>, InstructionError>,
225    var_addr: *mut u8,
226) -> u64 {
227    let invoke_context = get_invoke_context();
228    if invoke_context
229        .consume_checked(invoke_context.get_execution_cost().sysvar_base_cost + T::size_of() as u64)
230        .is_err()
231    {
232        panic!("Exceeded compute budget");
233    }
234
235    match sysvar {
236        Ok(sysvar_data) => unsafe {
237            *(var_addr as *mut _ as *mut T) = T::clone(&sysvar_data);
238            SUCCESS
239        },
240        Err(_) => UNSUPPORTED_SYSVAR,
241    }
242}
243
244struct SyscallStubs {}
245
246impl SyscallStubs {
247    fn fetch_and_write_sysvar<T: SysvarSerialize>(
248        &self,
249        var_addr: *mut u8,
250        offset: u64,
251        length: u64,
252        fetch: impl FnOnce(&SysvarCache) -> Result<Arc<T>, InstructionError>,
253    ) -> u64 {
254        // Consume compute units for the syscall.
255        let invoke_context = get_invoke_context();
256        let SVMTransactionExecutionCost {
257            sysvar_base_cost,
258            cpi_bytes_per_unit,
259            mem_op_base_cost,
260            ..
261        } = *invoke_context.get_execution_cost();
262
263        let sysvar_id_cost = 32_u64.checked_div(cpi_bytes_per_unit).unwrap_or(0);
264        let sysvar_buf_cost = length.checked_div(cpi_bytes_per_unit).unwrap_or(0);
265
266        if invoke_context
267            .consume_checked(
268                sysvar_base_cost
269                    .saturating_add(sysvar_id_cost)
270                    .saturating_add(std::cmp::max(sysvar_buf_cost, mem_op_base_cost)),
271            )
272            .is_err()
273        {
274            panic!("Exceeded compute budget");
275        }
276
277        // Fetch the sysvar from the cache.
278        let Ok(sysvar) = fetch(get_invoke_context().get_sysvar_cache()) else {
279            return UNSUPPORTED_SYSVAR;
280        };
281
282        // Check that the requested length is not greater than
283        // the actual serialized length of the sysvar data.
284        let Ok(expected_length) = bincode::serialized_size(&sysvar) else {
285            return UNSUPPORTED_SYSVAR;
286        };
287
288        if offset.saturating_add(length) > expected_length {
289            return UNSUPPORTED_SYSVAR;
290        }
291
292        // Write only the requested slice [offset, offset + length).
293        if let Ok(serialized) = bincode::serialize(&sysvar) {
294            unsafe {
295                ptr::copy_nonoverlapping(
296                    serialized[offset as usize..].as_ptr(),
297                    var_addr,
298                    length as usize,
299                )
300            };
301            SUCCESS
302        } else {
303            UNSUPPORTED_SYSVAR
304        }
305    }
306}
307impl solana_sysvar::program_stubs::SyscallStubs for SyscallStubs {
308    fn sol_log(&self, message: &str) {
309        let invoke_context = get_invoke_context();
310        ic_msg!(invoke_context, "Program log: {}", message);
311    }
312
313    fn sol_invoke_signed(
314        &self,
315        instruction: &Instruction,
316        account_infos: &[AccountInfo],
317        signers_seeds: &[&[&[u8]]],
318    ) -> ProgramResult {
319        let invoke_context = get_invoke_context();
320        let log_collector = invoke_context.get_log_collector();
321        let transaction_context = &invoke_context.transaction_context;
322        let instruction_context = transaction_context
323            .get_current_instruction_context()
324            .unwrap();
325        let caller = instruction_context.get_program_key().unwrap();
326
327        stable_log::program_invoke(
328            &log_collector,
329            &instruction.program_id,
330            invoke_context.get_stack_height(),
331        );
332
333        let signers = signers_seeds
334            .iter()
335            .map(|seeds| Pubkey::create_program_address(seeds, caller).unwrap())
336            .collect::<Vec<_>>();
337
338        invoke_context
339            .prepare_next_cpi_instruction(instruction.clone(), &signers)
340            .unwrap();
341
342        // Copy caller's account_info modifications into invoke_context accounts
343        let transaction_context = &invoke_context.transaction_context;
344        let instruction_context = transaction_context
345            .get_current_instruction_context()
346            .unwrap();
347        let next_instruction_context = transaction_context.get_next_instruction_context().unwrap();
348        let next_instruction_accounts = next_instruction_context.instruction_accounts();
349        let mut account_indices = Vec::with_capacity(next_instruction_accounts.len());
350        for instruction_account in next_instruction_accounts.iter() {
351            let account_key = transaction_context
352                .get_key_of_account_at_index(instruction_account.index_in_transaction)
353                .unwrap();
354            let account_info_index = account_infos
355                .iter()
356                .position(|account_info| account_info.unsigned_key() == account_key)
357                .ok_or(InstructionError::MissingAccount)
358                .unwrap();
359            let account_info = &account_infos[account_info_index];
360            let index_in_caller = instruction_context
361                .get_index_of_account_in_instruction(instruction_account.index_in_transaction)
362                .unwrap();
363            let mut borrowed_account = instruction_context
364                .try_borrow_instruction_account(index_in_caller)
365                .unwrap();
366            if borrowed_account.get_lamports() != account_info.lamports() {
367                borrowed_account
368                    .set_lamports(account_info.lamports())
369                    .unwrap();
370            }
371            let account_info_data = account_info.try_borrow_data().unwrap();
372            // The redundant check helps to avoid the expensive data comparison if we can
373            match borrowed_account.can_data_be_resized(account_info_data.len()) {
374                Ok(()) => borrowed_account
375                    .set_data_from_slice(&account_info_data)
376                    .unwrap(),
377                Err(err) if borrowed_account.get_data() != *account_info_data => {
378                    panic!("{err:?}");
379                }
380                _ => {}
381            }
382            // Change the owner at the end so that we are allowed to change the lamports and data before
383            if borrowed_account.get_owner() != account_info.owner {
384                borrowed_account
385                    .set_owner(account_info.owner.as_ref())
386                    .unwrap();
387            }
388            if instruction_account.is_writable() {
389                account_indices
390                    .push((instruction_account.index_in_transaction, account_info_index));
391            }
392        }
393
394        let mut compute_units_consumed = 0;
395        invoke_context
396            .process_instruction(&mut compute_units_consumed, &mut ExecuteTimings::default())
397            .map_err(|err| ProgramError::try_from(err).unwrap_or_else(|err| panic!("{}", err)))?;
398
399        // Copy invoke_context accounts modifications into caller's account_info
400        let transaction_context = &invoke_context.transaction_context;
401        let instruction_context = transaction_context
402            .get_current_instruction_context()
403            .unwrap();
404        for (index_in_transaction, account_info_index) in account_indices.into_iter() {
405            let index_in_caller = instruction_context
406                .get_index_of_account_in_instruction(index_in_transaction)
407                .unwrap();
408            let borrowed_account = instruction_context
409                .try_borrow_instruction_account(index_in_caller)
410                .unwrap();
411            let account_info = &account_infos[account_info_index];
412            **account_info.try_borrow_mut_lamports().unwrap() = borrowed_account.get_lamports();
413            if account_info.owner != borrowed_account.get_owner() {
414                // TODO Figure out a better way to allow the System Program to set the account owner
415                #[allow(clippy::transmute_ptr_to_ptr)]
416                #[allow(mutable_transmutes)]
417                let account_info_mut =
418                    unsafe { transmute::<&Pubkey, &mut Pubkey>(account_info.owner) };
419                *account_info_mut = *borrowed_account.get_owner();
420            }
421
422            let new_data = borrowed_account.get_data();
423            let new_len = new_data.len();
424
425            // Resize account_info data
426            if account_info.data_len() != new_len {
427                account_info.resize(new_len)?;
428            }
429
430            // Clone the data
431            let mut data = account_info.try_borrow_mut_data()?;
432            data.clone_from_slice(new_data);
433        }
434
435        stable_log::program_success(&log_collector, &instruction.program_id);
436        Ok(())
437    }
438
439    fn sol_get_clock_sysvar(&self, var_addr: *mut u8) -> u64 {
440        get_sysvar(
441            get_invoke_context().get_sysvar_cache().get_clock(),
442            var_addr,
443        )
444    }
445
446    fn sol_get_epoch_schedule_sysvar(&self, var_addr: *mut u8) -> u64 {
447        get_sysvar(
448            get_invoke_context().get_sysvar_cache().get_epoch_schedule(),
449            var_addr,
450        )
451    }
452
453    fn sol_get_epoch_rewards_sysvar(&self, var_addr: *mut u8) -> u64 {
454        get_sysvar(
455            get_invoke_context().get_sysvar_cache().get_epoch_rewards(),
456            var_addr,
457        )
458    }
459
460    #[allow(deprecated)]
461    fn sol_get_fees_sysvar(&self, var_addr: *mut u8) -> u64 {
462        get_sysvar(get_invoke_context().get_sysvar_cache().get_fees(), var_addr)
463    }
464
465    fn sol_get_rent_sysvar(&self, var_addr: *mut u8) -> u64 {
466        get_sysvar(get_invoke_context().get_sysvar_cache().get_rent(), var_addr)
467    }
468
469    fn sol_get_last_restart_slot(&self, var_addr: *mut u8) -> u64 {
470        get_sysvar(
471            get_invoke_context()
472                .get_sysvar_cache()
473                .get_last_restart_slot(),
474            var_addr,
475        )
476    }
477
478    fn sol_get_return_data(&self) -> Option<(Pubkey, Vec<u8>)> {
479        let (program_id, data) = get_invoke_context().transaction_context.get_return_data();
480        Some((*program_id, data.to_vec()))
481    }
482
483    fn sol_set_return_data(&self, data: &[u8]) {
484        let invoke_context = get_invoke_context();
485        let transaction_context = &mut invoke_context.transaction_context;
486        let instruction_context = transaction_context
487            .get_current_instruction_context()
488            .unwrap();
489        let caller = *instruction_context.get_program_key().unwrap();
490        transaction_context
491            .set_return_data(caller, data.to_vec())
492            .unwrap();
493    }
494
495    fn sol_get_stack_height(&self) -> u64 {
496        let invoke_context = get_invoke_context();
497        invoke_context.get_stack_height().try_into().unwrap()
498    }
499
500    fn sol_get_sysvar(
501        &self,
502        sysvar_id_addr: *const u8,
503        var_addr: *mut u8,
504        offset: u64,
505        length: u64,
506    ) -> u64 {
507        let sysvar_id = unsafe { &*(sysvar_id_addr as *const Pubkey) };
508
509        match *sysvar_id {
510            id if id == Clock::id() => self.fetch_and_write_sysvar::<Clock>(
511                var_addr,
512                offset,
513                length,
514                SysvarCache::get_clock,
515            ),
516            id if id == EpochRewards::id() => self.fetch_and_write_sysvar::<EpochRewards>(
517                var_addr,
518                offset,
519                length,
520                SysvarCache::get_epoch_rewards,
521            ),
522            id if id == EpochSchedule::id() => self.fetch_and_write_sysvar::<EpochSchedule>(
523                var_addr,
524                offset,
525                length,
526                SysvarCache::get_epoch_schedule,
527            ),
528            id if id == LastRestartSlot::id() => self.fetch_and_write_sysvar::<LastRestartSlot>(
529                var_addr,
530                offset,
531                length,
532                SysvarCache::get_last_restart_slot,
533            ),
534            id if id == Rent::id() => {
535                self.fetch_and_write_sysvar::<Rent>(var_addr, offset, length, SysvarCache::get_rent)
536            }
537            _ => UNSUPPORTED_SYSVAR,
538        }
539    }
540}
541
542pub fn find_file(filename: &str) -> Option<PathBuf> {
543    for dir in default_shared_object_dirs() {
544        let candidate = dir.join(filename);
545        if candidate.exists() {
546            return Some(candidate);
547        }
548    }
549    None
550}
551
552fn default_shared_object_dirs() -> Vec<PathBuf> {
553    let mut search_path = vec![];
554    if let Ok(bpf_out_dir) = std::env::var("BPF_OUT_DIR") {
555        search_path.push(PathBuf::from(bpf_out_dir));
556    } else if let Ok(bpf_out_dir) = std::env::var("SBF_OUT_DIR") {
557        search_path.push(PathBuf::from(bpf_out_dir));
558    }
559    search_path.push(PathBuf::from("tests/fixtures"));
560    if let Ok(dir) = std::env::current_dir() {
561        search_path.push(dir);
562    }
563    trace!("SBF .so search path: {search_path:?}");
564    search_path
565}
566
567pub fn read_file<P: AsRef<Path>>(path: P) -> Vec<u8> {
568    let path = path.as_ref();
569    let mut file = File::open(path)
570        .unwrap_or_else(|err| panic!("Failed to open \"{}\": {}", path.display(), err));
571
572    let mut file_data = Vec::new();
573    file.read_to_end(&mut file_data)
574        .unwrap_or_else(|err| panic!("Failed to read \"{}\": {}", path.display(), err));
575    file_data
576}
577
578pub struct ProgramTest {
579    accounts: Vec<(Pubkey, AccountSharedData)>,
580    genesis_accounts: Vec<(Pubkey, AccountSharedData)>,
581    builtin_programs: Vec<(Pubkey, &'static str, ProgramCacheEntry)>,
582    compute_max_units: Option<u64>,
583    prefer_bpf: bool,
584    deactivate_feature_set: HashSet<Pubkey>,
585    transaction_account_lock_limit: Option<usize>,
586}
587
588impl Default for ProgramTest {
589    /// Initialize a new ProgramTest
590    ///
591    /// If the `BPF_OUT_DIR` environment variable is defined, BPF programs will be preferred over
592    /// over a native instruction processor.  The `ProgramTest::prefer_bpf()` method may be
593    /// used to override this preference at runtime.  `cargo test-bpf` will set `BPF_OUT_DIR`
594    /// automatically.
595    ///
596    /// SBF program shared objects and account data files are searched for in
597    /// * the value of the `BPF_OUT_DIR` environment variable
598    /// * the `tests/fixtures` sub-directory
599    /// * the current working directory
600    ///
601    fn default() -> Self {
602        agave_logger::setup_with_default(
603            "solana_sbpf::vm=debug,solana_runtime::message_processor=debug,\
604             solana_runtime::system_instruction_processor=trace,solana_program_test=info",
605        );
606        let prefer_bpf =
607            std::env::var("BPF_OUT_DIR").is_ok() || std::env::var("SBF_OUT_DIR").is_ok();
608
609        Self {
610            accounts: vec![],
611            genesis_accounts: vec![],
612            builtin_programs: vec![],
613            compute_max_units: None,
614            prefer_bpf,
615            deactivate_feature_set: HashSet::default(),
616            transaction_account_lock_limit: None,
617        }
618    }
619}
620
621impl ProgramTest {
622    /// Create a `ProgramTest`.
623    ///
624    /// This is a wrapper around [`default`] and [`add_program`]. See their documentation for more
625    /// details.
626    ///
627    /// [`default`]: #method.default
628    /// [`add_program`]: #method.add_program
629    pub fn new(
630        program_name: &'static str,
631        program_id: Pubkey,
632        builtin_function: Option<BuiltinFunctionWithContext>,
633    ) -> Self {
634        let mut me = Self::default();
635        me.add_program(program_name, program_id, builtin_function);
636        me
637    }
638
639    /// Override default SBF program selection
640    pub fn prefer_bpf(&mut self, prefer_bpf: bool) {
641        self.prefer_bpf = prefer_bpf;
642    }
643
644    /// Override the default maximum compute units
645    pub fn set_compute_max_units(&mut self, compute_max_units: u64) {
646        debug_assert!(
647            compute_max_units <= i64::MAX as u64,
648            "Compute unit limit must fit in `i64::MAX`"
649        );
650        self.compute_max_units = Some(compute_max_units);
651    }
652
653    /// Override the default transaction account lock limit
654    pub fn set_transaction_account_lock_limit(&mut self, transaction_account_lock_limit: usize) {
655        self.transaction_account_lock_limit = Some(transaction_account_lock_limit);
656    }
657
658    /// Add an account to the test environment's genesis config.
659    pub fn add_genesis_account(&mut self, address: Pubkey, account: Account) {
660        self.genesis_accounts
661            .push((address, AccountSharedData::from(account)));
662    }
663
664    /// Add an account to the test environment
665    pub fn add_account(&mut self, address: Pubkey, account: Account) {
666        self.accounts
667            .push((address, AccountSharedData::from(account)));
668    }
669
670    /// Add an account to the test environment with the account data in the provided `filename`
671    pub fn add_account_with_file_data(
672        &mut self,
673        address: Pubkey,
674        lamports: u64,
675        owner: Pubkey,
676        filename: &str,
677    ) {
678        self.add_account(
679            address,
680            Account {
681                lamports,
682                data: read_file(find_file(filename).unwrap_or_else(|| {
683                    panic!("Unable to locate {filename}");
684                })),
685                owner,
686                executable: false,
687                rent_epoch: 0,
688            },
689        );
690    }
691
692    /// Add an account to the test environment with the account data in the provided as a base 64
693    /// string
694    pub fn add_account_with_base64_data(
695        &mut self,
696        address: Pubkey,
697        lamports: u64,
698        owner: Pubkey,
699        data_base64: &str,
700    ) {
701        self.add_account(
702            address,
703            Account {
704                lamports,
705                data: BASE64_STANDARD
706                    .decode(data_base64)
707                    .unwrap_or_else(|err| panic!("Failed to base64 decode: {err}")),
708                owner,
709                executable: false,
710                rent_epoch: 0,
711            },
712        );
713    }
714
715    pub fn add_sysvar_account<S: SysvarSerialize>(&mut self, address: Pubkey, sysvar: &S) {
716        let account = create_account_shared_data_for_test(sysvar);
717        self.add_account(address, account.into());
718    }
719
720    /// Add a BPF Upgradeable program to the test environment's genesis config.
721    ///
722    /// When testing BPF programs using the program ID of a runtime builtin
723    /// program - such as Core BPF programs - the program accounts must be
724    /// added to the genesis config in order to make them available to the new
725    /// Bank as it's being initialized.
726    ///
727    /// The presence of these program accounts will cause Bank to skip adding
728    /// the builtin version of the program, allowing the provided BPF program
729    /// to be used at the designated program ID instead.
730    ///
731    /// See https://github.com/anza-xyz/agave/blob/c038908600b8a1b0080229dea015d7fc9939c418/runtime/src/bank.rs#L5109-L5126.
732    pub fn add_upgradeable_program_to_genesis(
733        &mut self,
734        program_name: &'static str,
735        program_id: &Pubkey,
736    ) {
737        let program_file = find_file(&format!("{program_name}.so")).unwrap_or_else(|| {
738            panic!("Program file data not available for {program_name} ({program_id})")
739        });
740        let elf = read_file(program_file);
741        let program_accounts =
742            programs::bpf_loader_upgradeable_program_accounts(program_id, &elf, &Rent::default());
743        for (address, account) in program_accounts {
744            self.add_genesis_account(address, account);
745        }
746    }
747
748    /// Add a SBF program to the test environment.
749    ///
750    /// `program_name` will also be used to locate the SBF shared object in the current or fixtures
751    /// directory.
752    ///
753    /// If `builtin_function` is provided, the natively built-program may be used instead of the
754    /// SBF shared object depending on the `BPF_OUT_DIR` environment variable.
755    pub fn add_program(
756        &mut self,
757        program_name: &'static str,
758        program_id: Pubkey,
759        builtin_function: Option<BuiltinFunctionWithContext>,
760    ) {
761        let add_bpf = |this: &mut ProgramTest, program_file: PathBuf| {
762            let data = read_file(&program_file);
763            info!(
764                "\"{}\" SBF program from {}{}",
765                program_name,
766                program_file.display(),
767                std::fs::metadata(&program_file)
768                    .map(|metadata| {
769                        metadata
770                            .modified()
771                            .map(|time| {
772                                format!(
773                                    ", modified {}",
774                                    HumanTime::from(time)
775                                        .to_text_en(Accuracy::Precise, Tense::Past)
776                                )
777                            })
778                            .ok()
779                    })
780                    .ok()
781                    .flatten()
782                    .unwrap_or_default()
783            );
784
785            this.add_account(
786                program_id,
787                Account {
788                    lamports: Rent::default().minimum_balance(data.len()).max(1),
789                    data,
790                    owner: solana_sdk_ids::bpf_loader::id(),
791                    executable: true,
792                    rent_epoch: 0,
793                },
794            );
795        };
796
797        let warn_invalid_program_name = || {
798            let valid_program_names = default_shared_object_dirs()
799                .iter()
800                .filter_map(|dir| dir.read_dir().ok())
801                .flat_map(|read_dir| {
802                    read_dir.filter_map(|entry| {
803                        let path = entry.ok()?.path();
804                        if !path.is_file() {
805                            return None;
806                        }
807                        match path.extension()?.to_str()? {
808                            "so" => Some(path.file_stem()?.to_os_string()),
809                            _ => None,
810                        }
811                    })
812                })
813                .collect::<Vec<_>>();
814
815            if valid_program_names.is_empty() {
816                // This should be unreachable as `test-bpf` should guarantee at least one shared
817                // object exists somewhere.
818                warn!("No SBF shared objects found.");
819                return;
820            }
821
822            warn!(
823                "Possible bogus program name. Ensure the program name ({program_name}) matches \
824                 one of the following recognizable program names:",
825            );
826            for name in valid_program_names {
827                warn!(" - {}", name.to_str().unwrap());
828            }
829        };
830
831        let program_file = find_file(&format!("{program_name}.so"));
832        match (self.prefer_bpf, program_file, builtin_function) {
833            // If SBF is preferred (i.e., `test-sbf` is invoked) and a BPF shared object exists,
834            // use that as the program data.
835            (true, Some(file), _) => add_bpf(self, file),
836
837            // If SBF is not required (i.e., we were invoked with `test`), use the provided
838            // processor function as is.
839            (false, _, Some(builtin_function)) => {
840                self.add_builtin_program(program_name, program_id, builtin_function)
841            }
842
843            // Invalid: `test-sbf` invocation with no matching SBF shared object.
844            (true, None, _) => {
845                warn_invalid_program_name();
846                panic!("Program file data not available for {program_name} ({program_id})");
847            }
848
849            // Invalid: regular `test` invocation without a processor.
850            (false, _, None) => {
851                panic!("Program processor not available for {program_name} ({program_id})");
852            }
853        }
854    }
855
856    /// Add a builtin program to the test environment.
857    ///
858    /// Note that builtin programs are responsible for their own `stable_log` output.
859    pub fn add_builtin_program(
860        &mut self,
861        program_name: &'static str,
862        program_id: Pubkey,
863        builtin_function: BuiltinFunctionWithContext,
864    ) {
865        info!("\"{program_name}\" builtin program");
866        self.builtin_programs.push((
867            program_id,
868            program_name,
869            ProgramCacheEntry::new_builtin(0, program_name.len(), builtin_function),
870        ));
871    }
872
873    /// Deactivate a runtime feature.
874    ///
875    /// Note that all features are activated by default.
876    pub fn deactivate_feature(&mut self, feature_id: Pubkey) {
877        self.deactivate_feature_set.insert(feature_id);
878    }
879
880    fn setup_bank(
881        &mut self,
882    ) -> (
883        Arc<RwLock<BankForks>>,
884        Arc<RwLock<BlockCommitmentCache>>,
885        Hash,
886        GenesisConfigInfo,
887    ) {
888        {
889            use std::sync::Once;
890            static ONCE: Once = Once::new();
891
892            ONCE.call_once(|| {
893                solana_sysvar::program_stubs::set_syscall_stubs(Box::new(SyscallStubs {}));
894            });
895        }
896
897        let rent = Rent::default();
898        let fee_rate_governor = FeeRateGovernor {
899            // Initialize with a non-zero fee
900            lamports_per_signature: DEFAULT_TARGET_LAMPORTS_PER_SIGNATURE / 2,
901            ..FeeRateGovernor::default()
902        };
903        let bootstrap_validator_pubkey = Pubkey::new_unique();
904        let bootstrap_validator_stake_lamports =
905            rent.minimum_balance(VoteStateV4::size_of()) + 1_000_000 * LAMPORTS_PER_SOL;
906
907        let mint_keypair = Keypair::new();
908        let voting_keypair = Keypair::new();
909
910        // Remove features tagged to deactivate
911        let mut feature_set = FeatureSet::all_enabled();
912        for deactivate_feature_pk in &self.deactivate_feature_set {
913            if FEATURE_NAMES.contains_key(deactivate_feature_pk) {
914                feature_set.deactivate(deactivate_feature_pk);
915            } else {
916                warn!(
917                    "Feature {deactivate_feature_pk:?} set for deactivation is not a known \
918                     Feature public key"
919                );
920            }
921        }
922
923        let mut genesis_config = create_genesis_config_with_leader_ex(
924            1_000_000 * LAMPORTS_PER_SOL,
925            &mint_keypair.pubkey(),
926            &bootstrap_validator_pubkey,
927            &voting_keypair.pubkey(),
928            &Pubkey::new_unique(),
929            None,
930            bootstrap_validator_stake_lamports,
931            890_880,
932            fee_rate_governor,
933            rent.clone(),
934            ClusterType::Development,
935            &feature_set,
936            std::mem::take(&mut self.genesis_accounts),
937        );
938
939        let target_tick_duration = Duration::from_micros(100);
940        genesis_config.poh_config = PohConfig::new_sleep(target_tick_duration);
941        debug!("Payer address: {}", mint_keypair.pubkey());
942        debug!("Genesis config: {genesis_config}");
943
944        let bank = Bank::new_from_genesis(
945            &genesis_config,
946            Arc::new(RuntimeConfig {
947                compute_budget: self.compute_max_units.map(|max_units| ComputeBudget {
948                    compute_unit_limit: max_units,
949                    ..ComputeBudget::new_with_defaults(
950                        genesis_config
951                            .accounts
952                            .contains_key(&raise_cpi_nesting_limit_to_8::id()),
953                        genesis_config
954                            .accounts
955                            .contains_key(&increase_cpi_account_info_limit::id()),
956                    )
957                }),
958                transaction_account_lock_limit: self.transaction_account_lock_limit,
959                ..RuntimeConfig::default()
960            }),
961            Vec::default(),
962            None,
963            ACCOUNTS_DB_CONFIG_FOR_TESTING,
964            None,
965            None,
966            Arc::default(),
967            None,
968            None,
969        );
970
971        // Add commonly-used SPL programs as a convenience to the user
972        for (program_id, account) in programs::spl_programs(&rent).iter() {
973            bank.store_account(program_id, account);
974        }
975
976        // Add migrated Core BPF programs.
977        for (program_id, account) in programs::core_bpf_programs(&rent, |feature_id| {
978            genesis_config.accounts.contains_key(feature_id)
979        })
980        .iter()
981        {
982            bank.store_account(program_id, account);
983        }
984
985        // User-supplied additional builtins
986        let mut builtin_programs = Vec::new();
987        std::mem::swap(&mut self.builtin_programs, &mut builtin_programs);
988        for (program_id, name, builtin) in builtin_programs.into_iter() {
989            bank.add_builtin(program_id, name, builtin);
990        }
991
992        for (address, account) in self.accounts.iter() {
993            if bank.get_account(address).is_some() {
994                info!("Overriding account at {address}");
995            }
996            bank.store_account(address, account);
997        }
998        bank.set_capitalization_for_tests(bank.calculate_capitalization_for_tests());
999        // Advance beyond slot 0 for a slightly more realistic test environment
1000        let bank = {
1001            let bank = Arc::new(bank);
1002            bank.fill_bank_with_ticks_for_tests();
1003            let bank = Bank::new_from_parent(bank.clone(), bank.leader_id(), bank.slot() + 1);
1004            debug!("Bank slot: {}", bank.slot());
1005            bank
1006        };
1007        let slot = bank.slot();
1008        let last_blockhash = bank.last_blockhash();
1009        let bank_forks = BankForks::new_rw_arc(bank);
1010        let block_commitment_cache = Arc::new(RwLock::new(
1011            BlockCommitmentCache::new_for_tests_with_slots(slot, slot),
1012        ));
1013
1014        (
1015            bank_forks,
1016            block_commitment_cache,
1017            last_blockhash,
1018            GenesisConfigInfo {
1019                genesis_config,
1020                mint_keypair,
1021                voting_keypair,
1022                validator_pubkey: bootstrap_validator_pubkey,
1023            },
1024        )
1025    }
1026
1027    pub async fn start(mut self) -> (BanksClient, Keypair, Hash) {
1028        let (bank_forks, block_commitment_cache, last_blockhash, gci) = self.setup_bank();
1029        let target_tick_duration = gci.genesis_config.poh_config.target_tick_duration;
1030        let target_slot_duration = target_tick_duration * gci.genesis_config.ticks_per_slot as u32;
1031        let transport = start_local_server(
1032            bank_forks.clone(),
1033            block_commitment_cache.clone(),
1034            target_tick_duration,
1035        )
1036        .await;
1037        let banks_client = start_client(transport)
1038            .await
1039            .unwrap_or_else(|err| panic!("Failed to start banks client: {err}"));
1040
1041        // Run a simulated PohService to provide the client with new blockhashes.  New blockhashes
1042        // are required when sending multiple otherwise identical transactions in series from a
1043        // test
1044        tokio::spawn(async move {
1045            loop {
1046                tokio::time::sleep(target_slot_duration).await;
1047                bank_forks
1048                    .read()
1049                    .unwrap()
1050                    .working_bank()
1051                    .register_unique_recent_blockhash_for_test();
1052            }
1053        });
1054
1055        (banks_client, gci.mint_keypair, last_blockhash)
1056    }
1057
1058    /// Start the test client
1059    ///
1060    /// Returns a `BanksClient` interface into the test environment as well as a payer `Keypair`
1061    /// with SOL for sending transactions
1062    pub async fn start_with_context(mut self) -> ProgramTestContext {
1063        let (bank_forks, block_commitment_cache, last_blockhash, gci) = self.setup_bank();
1064        let target_tick_duration = gci.genesis_config.poh_config.target_tick_duration;
1065        let transport = start_local_server(
1066            bank_forks.clone(),
1067            block_commitment_cache.clone(),
1068            target_tick_duration,
1069        )
1070        .await;
1071        let banks_client = start_client(transport)
1072            .await
1073            .unwrap_or_else(|err| panic!("Failed to start banks client: {err}"));
1074
1075        ProgramTestContext::new(
1076            bank_forks,
1077            block_commitment_cache,
1078            banks_client,
1079            last_blockhash,
1080            gci,
1081        )
1082    }
1083}
1084
1085#[async_trait]
1086pub trait ProgramTestBanksClientExt {
1087    /// Get a new latest blockhash, similar in spirit to RpcClient::get_latest_blockhash()
1088    async fn get_new_latest_blockhash(&mut self, blockhash: &Hash) -> io::Result<Hash>;
1089}
1090
1091#[async_trait]
1092impl ProgramTestBanksClientExt for BanksClient {
1093    async fn get_new_latest_blockhash(&mut self, blockhash: &Hash) -> io::Result<Hash> {
1094        let mut num_retries = 0;
1095        let start = Instant::now();
1096        while start.elapsed().as_secs() < 5 {
1097            let new_blockhash = self.get_latest_blockhash().await?;
1098            if new_blockhash != *blockhash {
1099                return Ok(new_blockhash);
1100            }
1101            debug!("Got same blockhash ({blockhash:?}), will retry...");
1102
1103            tokio::time::sleep(Duration::from_millis(200)).await;
1104            num_retries += 1;
1105        }
1106
1107        Err(io::Error::other(format!(
1108            "Unable to get new blockhash after {}ms (retried {} times), stuck at {}",
1109            start.elapsed().as_millis(),
1110            num_retries,
1111            blockhash
1112        )))
1113    }
1114}
1115
1116struct DroppableTask<T>(Arc<AtomicBool>, JoinHandle<T>);
1117
1118impl<T> Drop for DroppableTask<T> {
1119    fn drop(&mut self) {
1120        self.0.store(true, Ordering::Relaxed);
1121        trace!(
1122            "stopping task, which is currently {}",
1123            if self.1.is_finished() {
1124                "finished"
1125            } else {
1126                "running"
1127            }
1128        );
1129    }
1130}
1131
1132pub struct ProgramTestContext {
1133    pub banks_client: BanksClient,
1134    pub last_blockhash: Hash,
1135    pub payer: Keypair,
1136    genesis_config: GenesisConfig,
1137    bank_forks: Arc<RwLock<BankForks>>,
1138    block_commitment_cache: Arc<RwLock<BlockCommitmentCache>>,
1139    _bank_task: DroppableTask<()>,
1140}
1141
1142impl ProgramTestContext {
1143    fn new(
1144        bank_forks: Arc<RwLock<BankForks>>,
1145        block_commitment_cache: Arc<RwLock<BlockCommitmentCache>>,
1146        banks_client: BanksClient,
1147        last_blockhash: Hash,
1148        genesis_config_info: GenesisConfigInfo,
1149    ) -> Self {
1150        // Run a simulated PohService to provide the client with new blockhashes.  New blockhashes
1151        // are required when sending multiple otherwise identical transactions in series from a
1152        // test
1153        let running_bank_forks = bank_forks.clone();
1154        let target_tick_duration = genesis_config_info
1155            .genesis_config
1156            .poh_config
1157            .target_tick_duration;
1158        let target_slot_duration =
1159            target_tick_duration * genesis_config_info.genesis_config.ticks_per_slot as u32;
1160        let exit = Arc::new(AtomicBool::new(false));
1161        let bank_task = DroppableTask(
1162            exit.clone(),
1163            tokio::spawn(async move {
1164                loop {
1165                    if exit.load(Ordering::Relaxed) {
1166                        break;
1167                    }
1168                    tokio::time::sleep(target_slot_duration).await;
1169                    running_bank_forks
1170                        .read()
1171                        .unwrap()
1172                        .working_bank()
1173                        .register_unique_recent_blockhash_for_test();
1174                }
1175            }),
1176        );
1177
1178        Self {
1179            banks_client,
1180            last_blockhash,
1181            payer: genesis_config_info.mint_keypair,
1182            genesis_config: genesis_config_info.genesis_config,
1183            bank_forks,
1184            block_commitment_cache,
1185            _bank_task: bank_task,
1186        }
1187    }
1188
1189    pub fn genesis_config(&self) -> &GenesisConfig {
1190        &self.genesis_config
1191    }
1192
1193    /// Manually increment vote credits for the current epoch in the specified vote account to simulate validator voting activity
1194    pub fn increment_vote_account_credits(
1195        &mut self,
1196        vote_account_address: &Pubkey,
1197        number_of_credits: u64,
1198    ) {
1199        let bank_forks = self.bank_forks.read().unwrap();
1200        let bank = bank_forks.working_bank();
1201
1202        // generate some vote activity for rewards
1203        let mut vote_account = bank.get_account(vote_account_address).unwrap();
1204        let mut vote_state =
1205            VoteStateV4::deserialize(vote_account.data(), vote_account_address).unwrap();
1206
1207        let epoch = bank.epoch();
1208        // Inlined from vote program - maximum number of epoch credits to keep in history
1209        const MAX_EPOCH_CREDITS_HISTORY: usize = 64;
1210        for _ in 0..number_of_credits {
1211            // Inline increment_credits logic from vote program.
1212            let credits = 1;
1213
1214            // never seen a credit
1215            if vote_state.epoch_credits.is_empty() {
1216                vote_state.epoch_credits.push((epoch, 0, 0));
1217            } else if epoch != vote_state.epoch_credits.last().unwrap().0 {
1218                let (_, credits_val, prev_credits) = *vote_state.epoch_credits.last().unwrap();
1219
1220                if credits_val != prev_credits {
1221                    // if credits were earned previous epoch
1222                    // append entry at end of list for the new epoch
1223                    vote_state
1224                        .epoch_credits
1225                        .push((epoch, credits_val, credits_val));
1226                } else {
1227                    // else just move the current epoch
1228                    vote_state.epoch_credits.last_mut().unwrap().0 = epoch;
1229                }
1230
1231                // Remove too old epoch_credits
1232                if vote_state.epoch_credits.len() > MAX_EPOCH_CREDITS_HISTORY {
1233                    vote_state.epoch_credits.remove(0);
1234                }
1235            }
1236
1237            vote_state.epoch_credits.last_mut().unwrap().1 = vote_state
1238                .epoch_credits
1239                .last()
1240                .unwrap()
1241                .1
1242                .saturating_add(credits);
1243        }
1244        let versioned = VoteStateVersions::new_v4(vote_state);
1245        vote_account.set_state(&versioned).unwrap();
1246        bank.store_account(vote_account_address, &vote_account);
1247    }
1248
1249    /// Create or overwrite an account, subverting normal runtime checks.
1250    ///
1251    /// This method exists to make it easier to set up artificial situations
1252    /// that would be difficult to replicate by sending individual transactions.
1253    /// Beware that it can be used to create states that would not be reachable
1254    /// by sending transactions!
1255    pub fn set_account(&mut self, address: &Pubkey, account: &AccountSharedData) {
1256        let bank_forks = self.bank_forks.read().unwrap();
1257        let bank = bank_forks.working_bank();
1258        bank.store_account(address, account);
1259    }
1260
1261    /// Create or overwrite a sysvar, subverting normal runtime checks.
1262    ///
1263    /// This method exists to make it easier to set up artificial situations
1264    /// that would be difficult to replicate on a new test cluster. Beware
1265    /// that it can be used to create states that would not be reachable
1266    /// under normal conditions!
1267    pub fn set_sysvar<T: SysvarId + SysvarSerialize>(&self, sysvar: &T) {
1268        let bank_forks = self.bank_forks.read().unwrap();
1269        let bank = bank_forks.working_bank();
1270        bank.set_sysvar_for_tests(sysvar);
1271    }
1272
1273    /// Force the working bank ahead to a new slot
1274    pub fn warp_to_slot(&mut self, warp_slot: Slot) -> Result<(), ProgramTestError> {
1275        let mut bank_forks = self.bank_forks.write().unwrap();
1276        let bank = bank_forks.working_bank();
1277
1278        // Fill ticks until a new blockhash is recorded, otherwise retried transactions will have
1279        // the same signature
1280        bank.fill_bank_with_ticks_for_tests();
1281
1282        // Ensure that we are actually progressing forward
1283        let working_slot = bank.slot();
1284        if warp_slot <= working_slot {
1285            return Err(ProgramTestError::InvalidWarpSlot);
1286        }
1287
1288        // Warp ahead to one slot *before* the desired slot because the bank
1289        // from Bank::warp_from_parent() is frozen. If the desired slot is one
1290        // slot *after* the working_slot, no need to warp at all.
1291        let pre_warp_slot = warp_slot - 1;
1292        let warp_bank = if pre_warp_slot == working_slot {
1293            bank.freeze();
1294            bank
1295        } else {
1296            bank_forks
1297                .insert(Bank::warp_from_parent(
1298                    bank,
1299                    &Pubkey::default(),
1300                    pre_warp_slot,
1301                ))
1302                .clone_without_scheduler()
1303        };
1304
1305        bank_forks.set_root(
1306            pre_warp_slot,
1307            None, // snapshots are disabled
1308            Some(pre_warp_slot),
1309        );
1310
1311        // warp_bank is frozen so go forward to get unfrozen bank at warp_slot
1312        bank_forks.insert(Bank::new_from_parent(
1313            warp_bank,
1314            &Pubkey::default(),
1315            warp_slot,
1316        ));
1317
1318        // Update block commitment cache, otherwise banks server will poll at
1319        // the wrong slot
1320        let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
1321        // HACK: The root set here should be `pre_warp_slot`, but since we're
1322        // in a testing environment, the root bank never updates after a warp.
1323        // The ticking thread only updates the working bank, and never the root
1324        // bank.
1325        w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
1326
1327        let bank = bank_forks.working_bank();
1328        self.last_blockhash = bank.last_blockhash();
1329        Ok(())
1330    }
1331
1332    pub fn warp_to_epoch(&mut self, warp_epoch: Epoch) -> Result<(), ProgramTestError> {
1333        let warp_slot = self
1334            .genesis_config
1335            .epoch_schedule
1336            .get_first_slot_in_epoch(warp_epoch);
1337        self.warp_to_slot(warp_slot)
1338    }
1339
1340    /// warp forward one more slot and force reward interval end
1341    pub fn warp_forward_force_reward_interval_end(&mut self) -> Result<(), ProgramTestError> {
1342        let mut bank_forks = self.bank_forks.write().unwrap();
1343        let bank = bank_forks.working_bank();
1344
1345        // Fill ticks until a new blockhash is recorded, otherwise retried transactions will have
1346        // the same signature
1347        bank.fill_bank_with_ticks_for_tests();
1348        let pre_warp_slot = bank.slot();
1349
1350        bank_forks.set_root(
1351            pre_warp_slot,
1352            None, // snapshot_controller
1353            Some(pre_warp_slot),
1354        );
1355
1356        // warp_bank is frozen so go forward to get unfrozen bank at warp_slot
1357        let warp_slot = pre_warp_slot + 1;
1358        let mut warp_bank = Bank::new_from_parent(bank, &Pubkey::default(), warp_slot);
1359
1360        warp_bank.force_reward_interval_end_for_tests();
1361        bank_forks.insert(warp_bank);
1362
1363        // Update block commitment cache, otherwise banks server will poll at
1364        // the wrong slot
1365        let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
1366        // HACK: The root set here should be `pre_warp_slot`, but since we're
1367        // in a testing environment, the root bank never updates after a warp.
1368        // The ticking thread only updates the working bank, and never the root
1369        // bank.
1370        w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
1371
1372        let bank = bank_forks.working_bank();
1373        self.last_blockhash = bank.last_blockhash();
1374        Ok(())
1375    }
1376
1377    /// Get a new latest blockhash, similar in spirit to RpcClient::get_latest_blockhash()
1378    pub async fn get_new_latest_blockhash(&mut self) -> io::Result<Hash> {
1379        let blockhash = self
1380            .banks_client
1381            .get_new_latest_blockhash(&self.last_blockhash)
1382            .await?;
1383        self.last_blockhash = blockhash;
1384        Ok(blockhash)
1385    }
1386
1387    /// record a hard fork slot in working bank; should be in the past
1388    pub fn register_hard_fork(&mut self, hard_fork_slot: Slot) {
1389        self.bank_forks
1390            .read()
1391            .unwrap()
1392            .working_bank()
1393            .register_hard_fork(hard_fork_slot)
1394    }
1395}