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