1#![allow(clippy::integer_arithmetic)]
3
4pub use tokio;
6use {
7 async_trait::async_trait,
8 chrono_humanize::{Accuracy, HumanTime, Tense},
9 log::*,
10 solana_bpf_loader_program::serialization::serialize_parameters,
11 solana_program_runtime::{
12 compute_budget::ComputeBudget, ic_msg, invoke_context::ProcessInstructionWithContext,
13 stable_log, timings::ExecuteTimings,
14 },
15 solana_runtime::{
16 bank::Bank,
17 bank_forks::BankForks,
18 builtins::Builtin,
19 commitment::BlockCommitmentCache,
20 genesis_utils::{create_genesis_config_with_leader_ex, GenesisConfigInfo},
21 },
22 solana_sdk::{
23 account::{Account, AccountSharedData, ReadableAccount},
24 account_info::AccountInfo,
25 clock::Slot,
26 entrypoint::{deserialize, ProgramResult, SUCCESS},
27 feature_set::FEATURE_NAMES,
28 fee_calculator::{FeeCalculator, FeeRateGovernor, DEFAULT_TARGET_LAMPORTS_PER_SIGNATURE},
29 genesis_config::{ClusterType, GenesisConfig},
30 hash::Hash,
31 instruction::{Instruction, InstructionError},
32 native_token::sol_to_lamports,
33 poh_config::PohConfig,
34 program_error::{ProgramError, ACCOUNT_BORROW_FAILED, UNSUPPORTED_SYSVAR},
35 pubkey::Pubkey,
36 rent::Rent,
37 signature::{Keypair, Signer},
38 sysvar::{Sysvar, SysvarId},
39 },
40 solana_vote_program::vote_state::{VoteState, VoteStateVersions},
41 std::{
42 cell::RefCell,
43 collections::{HashMap, HashSet},
44 convert::TryFrom,
45 fs::File,
46 io::{self, Read},
47 mem::transmute,
48 path::{Path, PathBuf},
49 sync::{
50 atomic::{AtomicBool, Ordering},
51 Arc, RwLock,
52 },
53 time::{Duration, Instant},
54 },
55 thiserror::Error,
56 tokio::task::JoinHandle,
57};
58pub use solana_program_runtime::invoke_context::InvokeContext;
60
61use crate::banks_client::{start_client, BanksClient};
62
63use crate::banks_server::start_local_server;
64
65#[derive(Error, Debug, PartialEq, Eq)]
67pub enum ProgramTestError {
68 #[error("Warp slot not in the future")]
70 InvalidWarpSlot,
71}
72
73thread_local! {
74 static INVOKE_CONTEXT: RefCell<Option<usize>> = RefCell::new(None);
75}
76fn set_invoke_context(new: &mut InvokeContext) {
77 INVOKE_CONTEXT
78 .with(|invoke_context| unsafe { invoke_context.replace(Some(transmute::<_, usize>(new))) });
79}
80fn get_invoke_context<'a, 'b>() -> &'a mut InvokeContext<'b> {
81 let ptr = INVOKE_CONTEXT.with(|invoke_context| match *invoke_context.borrow() {
82 Some(val) => val,
83 None => panic!("Invoke context not set!"),
84 });
85 unsafe { transmute::<usize, &mut InvokeContext>(ptr) }
86}
87
88pub fn builtin_process_instruction(
89 process_instruction: solana_sdk::entrypoint::ProcessInstruction,
90 _first_instruction_account: usize,
91 invoke_context: &mut InvokeContext,
92) -> Result<(), InstructionError> {
93 set_invoke_context(invoke_context);
94
95 let transaction_context = &invoke_context.transaction_context;
96 let instruction_context = transaction_context.get_current_instruction_context()?;
97 let instruction_data = instruction_context.get_instruction_data();
98 let instruction_account_indices = 0..instruction_context.get_number_of_instruction_accounts();
99
100 let log_collector = invoke_context.get_log_collector();
101 let program_id = instruction_context.get_last_program_key(transaction_context)?;
102 stable_log::program_invoke(
103 &log_collector,
104 program_id,
105 invoke_context.get_stack_height(),
106 );
107
108 let deduplicated_indices: HashSet<usize> = instruction_account_indices.collect();
110
111 let (mut parameter_bytes, _account_lengths) = serialize_parameters(
113 invoke_context.transaction_context,
114 invoke_context
115 .transaction_context
116 .get_current_instruction_context()?,
117 true,
118 )?;
119
120 let (program_id, account_infos, _input) =
122 unsafe { deserialize(&mut parameter_bytes.as_slice_mut()[0] as *mut u8) };
123
124 process_instruction(program_id, &account_infos, instruction_data).map_err(|err| {
126 let err = u64::from(err);
127 stable_log::program_failure(&log_collector, program_id, &err.into());
128 err
129 })?;
130 stable_log::program_success(&log_collector, program_id);
131
132 let account_info_map: HashMap<_, _> = account_infos.into_iter().map(|a| (a.key, a)).collect();
134
135 let transaction_context = &invoke_context.transaction_context;
138 let instruction_context = transaction_context.get_current_instruction_context()?;
139
140 for i in deduplicated_indices.into_iter() {
142 let mut borrowed_account =
143 instruction_context.try_borrow_instruction_account(transaction_context, i)?;
144 if borrowed_account.is_writable() {
145 if let Some(account_info) = account_info_map.get(borrowed_account.get_key()) {
146 if borrowed_account.get_lamports() != account_info.lamports() {
147 borrowed_account.set_lamports(account_info.lamports())?;
148 }
149
150 if borrowed_account
151 .can_data_be_resized(account_info.data_len())
152 .is_ok()
153 && borrowed_account.can_data_be_changed().is_ok()
154 {
155 borrowed_account.set_data(&account_info.data.borrow())?;
156 }
157 if borrowed_account.get_owner() != account_info.owner {
158 borrowed_account.set_owner(account_info.owner.as_ref())?;
159 }
160 }
161 }
162 }
163
164 Ok(())
165}
166
167
168fn get_sysvar<T: Default + Sysvar + Sized + serde::de::DeserializeOwned + Clone>(
169 sysvar: Result<Arc<T>, InstructionError>,
170 var_addr: *mut u8,
171) -> u64 {
172 let invoke_context = get_invoke_context();
173 if invoke_context
174 .get_compute_meter()
175 .try_borrow_mut()
176 .map_err(|_| ACCOUNT_BORROW_FAILED)
177 .unwrap()
178 .consume(invoke_context.get_compute_budget().sysvar_base_cost + T::size_of() as u64)
179 .is_err()
180 {
181 panic!("Exceeded compute budget");
182 }
183
184 match sysvar {
185 Ok(sysvar_data) => unsafe {
186 *(var_addr as *mut _ as *mut T) = T::clone(&sysvar_data);
187 SUCCESS
188 },
189 Err(_) => UNSUPPORTED_SYSVAR,
190 }
191}
192
193struct SyscallStubs {}
194impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
195 fn sol_log(&self, message: &str) {
196 let invoke_context = get_invoke_context();
197 ic_msg!(invoke_context, "Program log: {}", message);
198 }
199
200 fn sol_log_compute_units(&self) {
201 let invoke_context = get_invoke_context();
202 let compute_meter = invoke_context.get_compute_meter();
203 let compute_units = compute_meter.borrow().get_remaining();
204 ic_msg!(
205 invoke_context,
206 "Program consumed {} compute units",
207 compute_units
208 );
209 }
210
211 fn sol_invoke_signed(
212 &self,
213 instruction: &Instruction,
214 account_infos: &[AccountInfo],
215 signers_seeds: &[&[&[u8]]],
216 ) -> ProgramResult {
217 let invoke_context = get_invoke_context();
218 let log_collector = invoke_context.get_log_collector();
219 let transaction_context = &invoke_context.transaction_context;
220 let instruction_context = transaction_context
221 .get_current_instruction_context()
222 .unwrap();
223 let caller = instruction_context
224 .get_last_program_key(transaction_context)
225 .unwrap();
226
227 stable_log::program_invoke(
228 &log_collector,
229 &instruction.program_id,
230 invoke_context.get_stack_height(),
231 );
232
233 let signers = signers_seeds
234 .iter()
235 .map(|seeds| Pubkey::create_program_address(seeds, caller).unwrap())
236 .collect::<Vec<_>>();
237
238 let (instruction_accounts, program_indices) = invoke_context
239 .prepare_instruction(instruction, &signers)
240 .unwrap();
241
242 let transaction_context = &invoke_context.transaction_context;
244 let instruction_context = transaction_context
245 .get_current_instruction_context()
246 .unwrap();
247 let mut account_indices = Vec::with_capacity(instruction_accounts.len());
248 for instruction_account in instruction_accounts.iter() {
249 let account_key = transaction_context
250 .get_key_of_account_at_index(instruction_account.index_in_transaction)
251 .unwrap();
252 let account_info_index = account_infos
253 .iter()
254 .position(|account_info| account_info.unsigned_key() == account_key)
255 .ok_or(InstructionError::MissingAccount)
256 .unwrap();
257 let account_info = &account_infos[account_info_index];
258 let mut borrowed_account = instruction_context
259 .try_borrow_instruction_account(
260 transaction_context,
261 instruction_account.index_in_caller,
262 )
263 .unwrap();
264 if borrowed_account.get_lamports() != account_info.lamports() {
265 borrowed_account
266 .set_lamports(account_info.lamports())
267 .unwrap();
268 }
269 let account_info_data = account_info.try_borrow_data().unwrap();
270 match borrowed_account
272 .can_data_be_resized(account_info_data.len())
273 .and_then(|_| borrowed_account.can_data_be_changed())
274 {
275 Ok(()) => borrowed_account.set_data(&account_info_data).unwrap(),
276 Err(err) if borrowed_account.get_data() != *account_info_data => {
277 panic!("{:?}", err);
278 }
279 _ => {}
280 }
281 if borrowed_account.is_executable() != account_info.executable {
282 borrowed_account
283 .set_executable(account_info.executable)
284 .unwrap();
285 }
286 if borrowed_account.get_owner() != account_info.owner {
288 borrowed_account
289 .set_owner(account_info.owner.as_ref())
290 .unwrap();
291 }
292 drop(borrowed_account);
293 let account = transaction_context
294 .get_account_at_index(instruction_account.index_in_transaction)
295 .unwrap()
296 .borrow();
297 assert_eq!(account.rent_epoch(), account_info.rent_epoch);
298 if instruction_account.is_writable {
299 account_indices.push((instruction_account.index_in_caller, account_info_index));
300 }
301 }
302
303 let mut compute_units_consumed = 0;
304 invoke_context
305 .process_instruction(
306 &instruction.data,
307 &instruction_accounts,
308 &program_indices,
309 &mut compute_units_consumed,
310 &mut ExecuteTimings::default(),
311 )
312 .map_err(|err| ProgramError::try_from(err).unwrap_or_else(|err| panic!("{}", err)))?;
313
314 let transaction_context = &invoke_context.transaction_context;
316 let instruction_context = transaction_context
317 .get_current_instruction_context()
318 .unwrap();
319 for (index_in_caller, account_info_index) in account_indices.into_iter() {
320 let borrowed_account = instruction_context
321 .try_borrow_instruction_account(transaction_context, index_in_caller)
322 .unwrap();
323 let account_info = &account_infos[account_info_index];
324 **account_info.try_borrow_mut_lamports().unwrap() = borrowed_account.get_lamports();
325 if account_info.owner != borrowed_account.get_owner() {
326 #[allow(clippy::transmute_ptr_to_ptr)]
328 #[allow(mutable_transmutes)]
329 let account_info_mut =
330 unsafe { transmute::<&Pubkey, &mut Pubkey>(account_info.owner) };
331 *account_info_mut = *borrowed_account.get_owner();
332 }
333
334 let new_data = borrowed_account.get_data();
335 let new_len = new_data.len();
336
337 if account_info.data_len() < new_len {
339 account_info.realloc(new_len, false)?;
340 }
341
342 let mut data = account_info.try_borrow_mut_data()?;
344 data.clone_from_slice(new_data);
345 }
346
347 stable_log::program_success(&log_collector, &instruction.program_id);
348 Ok(())
349 }
350
351 fn sol_get_clock_sysvar(&self, var_addr: *mut u8) -> u64 {
352 get_sysvar(
353 get_invoke_context().get_sysvar_cache().get_clock(),
354 var_addr,
355 )
356 }
357
358 fn sol_get_epoch_schedule_sysvar(&self, var_addr: *mut u8) -> u64 {
359 get_sysvar(
360 get_invoke_context().get_sysvar_cache().get_epoch_schedule(),
361 var_addr,
362 )
363 }
364
365 #[allow(deprecated)]
366 fn sol_get_fees_sysvar(&self, var_addr: *mut u8) -> u64 {
367 get_sysvar(get_invoke_context().get_sysvar_cache().get_fees(), var_addr)
368 }
369
370 fn sol_get_rent_sysvar(&self, var_addr: *mut u8) -> u64 {
371 get_sysvar(get_invoke_context().get_sysvar_cache().get_rent(), var_addr)
372 }
373
374 fn sol_get_return_data(&self) -> Option<(Pubkey, Vec<u8>)> {
375 let (program_id, data) = get_invoke_context().transaction_context.get_return_data();
376 Some((*program_id, data.to_vec()))
377 }
378
379 fn sol_set_return_data(&self, data: &[u8]) {
380 let invoke_context = get_invoke_context();
381 let transaction_context = &mut invoke_context.transaction_context;
382 let instruction_context = transaction_context
383 .get_current_instruction_context()
384 .unwrap();
385 let caller = *instruction_context
386 .get_last_program_key(transaction_context)
387 .unwrap();
388 transaction_context
389 .set_return_data(caller, data.to_vec())
390 .unwrap();
391 }
392}
393
394pub fn find_file(filename: &str) -> Option<PathBuf> {
395 for dir in default_shared_object_dirs() {
396 let candidate = dir.join(filename);
397 if candidate.exists() {
398 return Some(candidate);
399 }
400 }
401 None
402}
403
404fn default_shared_object_dirs() -> Vec<PathBuf> {
405 let mut search_path = vec![];
406 if let Ok(bpf_out_dir) = std::env::var("BPF_OUT_DIR") {
407 search_path.push(PathBuf::from(bpf_out_dir));
408 } else if let Ok(bpf_out_dir) = std::env::var("SBF_OUT_DIR") {
409 search_path.push(PathBuf::from(bpf_out_dir));
410 }
411 search_path.push(PathBuf::from("tests/fixtures"));
412 if let Ok(dir) = std::env::current_dir() {
413 search_path.push(dir);
414 }
415 trace!("BPF .so search path: {:?}", search_path);
416 search_path
417}
418
419pub fn read_file<P: AsRef<Path>>(path: P) -> Vec<u8> {
420 let path = path.as_ref();
421 let mut file = File::open(path)
422 .unwrap_or_else(|err| panic!("Failed to open \"{}\": {}", path.display(), err));
423
424 let mut file_data = Vec::new();
425 file.read_to_end(&mut file_data)
426 .unwrap_or_else(|err| panic!("Failed to read \"{}\": {}", path.display(), err));
427 file_data
428}
429
430pub struct ProgramTest {
431 accounts: Vec<(Pubkey, AccountSharedData)>,
432 builtins: Vec<Builtin>,
433 compute_max_units: Option<u64>,
434 prefer_bpf: bool,
435 use_bpf_jit: bool,
436 deactivate_feature_set: HashSet<Pubkey>,
437}
438
439impl Default for ProgramTest {
440 fn default() -> Self {
453 solana_logger::setup_with_default(
454 "solana_rbpf::vm=debug,\
455 solana_runtime::message_processor=debug,\
456 solana_runtime::system_instruction_processor=trace,\
457 solana_program_test=info",
458 );
459 let prefer_bpf =
460 std::env::var("BPF_OUT_DIR").is_ok() || std::env::var("SBF_OUT_DIR").is_ok();
461
462 Self {
463 accounts: vec![],
464 builtins: vec![],
465 compute_max_units: None,
466 prefer_bpf,
467 use_bpf_jit: false,
468 deactivate_feature_set: HashSet::default(),
469 }
470 }
471}
472
473impl ProgramTest {
474 pub fn new(
482 program_name: &str,
483 program_id: Pubkey,
484 process_instruction: Option<ProcessInstructionWithContext>,
485 ) -> Self {
486 let mut me = Self::default();
487 me.add_program(program_name, program_id, process_instruction);
488 me
489 }
490
491 pub fn prefer_bpf(&mut self, prefer_bpf: bool) {
493 self.prefer_bpf = prefer_bpf;
494 }
495
496 pub fn set_compute_max_units(&mut self, compute_max_units: u64) {
498 self.compute_max_units = Some(compute_max_units);
499 }
500
501 #[allow(deprecated)]
503 #[deprecated(since = "1.8.0", note = "please use `set_compute_max_units` instead")]
504 pub fn set_bpf_compute_max_units(&mut self, bpf_compute_max_units: u64) {
505 self.compute_max_units = Some(bpf_compute_max_units);
506 }
507
508 pub fn use_bpf_jit(&mut self, use_bpf_jit: bool) {
510 self.use_bpf_jit = use_bpf_jit;
511 }
512
513 pub fn add_account(&mut self, address: Pubkey, account: Account) {
515 self.accounts
516 .push((address, AccountSharedData::from(account)));
517 }
518
519 pub fn add_account_with_file_data(
521 &mut self,
522 address: Pubkey,
523 lamports: u64,
524 owner: Pubkey,
525 filename: &str,
526 ) {
527 self.add_account(
528 address,
529 Account {
530 lamports,
531 data: read_file(find_file(filename).unwrap_or_else(|| {
532 panic!("Unable to locate {}", filename);
533 })),
534 owner,
535 executable: false,
536 rent_epoch: 0,
537 },
538 );
539 }
540
541 pub fn add_account_with_base64_data(
544 &mut self,
545 address: Pubkey,
546 lamports: u64,
547 owner: Pubkey,
548 data_base64: &str,
549 ) {
550 self.add_account(
551 address,
552 Account {
553 lamports,
554 data: base64::decode(data_base64)
555 .unwrap_or_else(|err| panic!("Failed to base64 decode: {}", err)),
556 owner,
557 executable: false,
558 rent_epoch: 0,
559 },
560 );
561 }
562
563 pub fn add_program(
571 &mut self,
572 program_name: &str,
573 program_id: Pubkey,
574 process_instruction: Option<ProcessInstructionWithContext>,
575 ) {
576 let add_bpf = |this: &mut ProgramTest, program_file: PathBuf| {
577 let data = read_file(&program_file);
578 info!(
579 "\"{}\" BPF program from {}{}",
580 program_name,
581 program_file.display(),
582 std::fs::metadata(&program_file)
583 .map(|metadata| {
584 metadata
585 .modified()
586 .map(|time| {
587 format!(
588 ", modified {}",
589 HumanTime::from(time)
590 .to_text_en(Accuracy::Precise, Tense::Past)
591 )
592 })
593 .ok()
594 })
595 .ok()
596 .flatten()
597 .unwrap_or_default()
598 );
599
600 this.add_account(
601 program_id,
602 Account {
603 lamports: Rent::default().minimum_balance(data.len()).min(1),
604 data,
605 owner: solana_sdk::bpf_loader::id(),
606 executable: true,
607 rent_epoch: 0,
608 },
609 );
610 };
611
612 let add_native = |this: &mut ProgramTest, process_fn: ProcessInstructionWithContext| {
613 info!("\"{}\" program loaded as native code", program_name);
614 this.builtins
615 .push(Builtin::new(program_name, program_id, process_fn));
616 };
617
618 let warn_invalid_program_name = || {
619 let valid_program_names = default_shared_object_dirs()
620 .iter()
621 .filter_map(|dir| dir.read_dir().ok())
622 .flat_map(|read_dir| {
623 read_dir.filter_map(|entry| {
624 let path = entry.ok()?.path();
625 if !path.is_file() {
626 return None;
627 }
628 match path.extension()?.to_str()? {
629 "so" => Some(path.file_stem()?.to_os_string()),
630 _ => None,
631 }
632 })
633 })
634 .collect::<Vec<_>>();
635
636 if valid_program_names.is_empty() {
637 warn!("No BPF shared objects found.");
640 return;
641 }
642
643 warn!(
644 "Possible bogus program name. Ensure the program name ({}) \
645 matches one of the following recognizable program names:",
646 program_name,
647 );
648 for name in valid_program_names {
649 warn!(" - {}", name.to_str().unwrap());
650 }
651 };
652
653 let program_file = find_file(&format!("{}.so", program_name));
654 match (self.prefer_bpf, program_file, process_instruction) {
655 (true, Some(file), _) => add_bpf(self, file),
658
659 (false, _, Some(process)) => add_native(self, process),
664
665 (true, None, _) => {
667 warn_invalid_program_name();
668 panic!(
669 "Program file data not available for {} ({})",
670 program_name, program_id
671 );
672 }
673
674 (false, _, None) => {
676 panic!(
677 "Program processor not available for {} ({})",
678 program_name, program_id
679 );
680 }
681 }
682 }
683
684 pub fn add_builtin_program(
688 &mut self,
689 program_name: &str,
690 program_id: Pubkey,
691 process_instruction: ProcessInstructionWithContext,
692 ) {
693 info!("\"{}\" builtin program", program_name);
694 self.builtins
695 .push(Builtin::new(program_name, program_id, process_instruction));
696 }
697
698 pub fn deactivate_feature(&mut self, feature_id: Pubkey) {
702 self.deactivate_feature_set.insert(feature_id);
703 }
704
705 fn setup_bank(
706 &self,
707 ) -> (
708 Arc<RwLock<BankForks>>,
709 Arc<RwLock<BlockCommitmentCache>>,
710 Hash,
711 GenesisConfigInfo,
712 ) {
713 {
714 use std::sync::Once;
715 static ONCE: Once = Once::new();
716
717 ONCE.call_once(|| {
718 solana_sdk::program_stubs::set_syscall_stubs(Box::new(SyscallStubs {}));
719 });
720 }
721
722 let rent = Rent::default();
723 let fee_rate_governor = FeeRateGovernor {
724 lamports_per_signature: DEFAULT_TARGET_LAMPORTS_PER_SIGNATURE / 2,
726 ..FeeRateGovernor::default()
727 };
728 let bootstrap_validator_pubkey = Pubkey::new_unique();
729 let bootstrap_validator_stake_lamports =
730 rent.minimum_balance(VoteState::size_of()) + sol_to_lamports(1_000_000.0);
731
732 let mint_keypair = Keypair::new();
733 let voting_keypair = Keypair::new();
734
735 let mut genesis_config = create_genesis_config_with_leader_ex(
736 sol_to_lamports(1_000_000.0),
737 &mint_keypair.pubkey(),
738 &bootstrap_validator_pubkey,
739 &voting_keypair.pubkey(),
740 &Pubkey::new_unique(),
741 bootstrap_validator_stake_lamports,
742 42,
743 fee_rate_governor,
744 rent,
745 ClusterType::Development,
746 vec![],
747 );
748
749 for deactivate_feature_pk in &self.deactivate_feature_set {
751 if FEATURE_NAMES.contains_key(deactivate_feature_pk) {
752 match genesis_config.accounts.remove(deactivate_feature_pk) {
753 Some(_) => debug!("Feature for {:?} deactivated", deactivate_feature_pk),
754 None => warn!(
755 "Feature {:?} set for deactivation not found in genesis_config account list, ignored.",
756 deactivate_feature_pk
757 ),
758 }
759 } else {
760 warn!(
761 "Feature {:?} set for deactivation is not a known Feature public key",
762 deactivate_feature_pk
763 );
764 }
765 }
766
767 let target_tick_duration = Duration::from_micros(100);
768 genesis_config.poh_config = PohConfig::new_sleep(target_tick_duration);
769 debug!("Payer address: {}", mint_keypair.pubkey());
770 debug!("Genesis config: {}", genesis_config);
771
772 let mut bank = Bank::new_for_tests(&genesis_config);
773
774 macro_rules! add_builtin {
776 ($b:expr) => {
777 bank.add_builtin(&$b.0, &$b.1, $b.2)
778 };
779 }
780 add_builtin!(solana_bpf_loader_deprecated_program!());
781 if self.use_bpf_jit {
782 add_builtin!(solana_bpf_loader_program_with_jit!());
783 add_builtin!(solana_bpf_loader_upgradeable_program_with_jit!());
784 } else {
785 add_builtin!(solana_bpf_loader_program!());
786 add_builtin!(solana_bpf_loader_upgradeable_program!());
787 }
788
789 for (program_id, account) in crate::programs::spl_programs(&Rent::default()).iter() {
791 bank.store_account(program_id, account);
792 }
793
794 for builtin in self.builtins.iter() {
796 bank.add_builtin(
797 &builtin.name,
798 &builtin.id,
799 builtin.process_instruction_with_context,
800 );
801 }
802
803 for (address, account) in self.accounts.iter() {
804 if bank.get_account(address).is_some() {
805 info!("Overriding account at {}", address);
806 }
807 bank.store_account(address, account);
808 }
809 bank.set_capitalization();
810 if let Some(max_units) = self.compute_max_units {
811 bank.set_compute_budget(Some(ComputeBudget {
812 compute_unit_limit: max_units,
813 ..ComputeBudget::default()
814 }));
815 }
816 let bank = {
818 let bank = Arc::new(bank);
819 bank.fill_bank_with_ticks_for_tests();
820 let bank = Bank::new_from_parent(&bank, bank.collector_id(), bank.slot() + 1);
821 debug!("Bank slot: {}", bank.slot());
822 bank
823 };
824 let slot = bank.slot();
825 let last_blockhash = bank.last_blockhash();
826 let bank_forks = Arc::new(RwLock::new(BankForks::new(bank)));
827 let block_commitment_cache = Arc::new(RwLock::new(
828 BlockCommitmentCache::new_for_tests_with_slots(slot, slot),
829 ));
830
831 (
832 bank_forks,
833 block_commitment_cache,
834 last_blockhash,
835 GenesisConfigInfo {
836 genesis_config,
837 mint_keypair,
838 voting_keypair,
839 validator_pubkey: bootstrap_validator_pubkey,
840 },
841 )
842 }
843
844 pub async fn start(self) -> (BanksClient, Keypair, Hash) {
845 let (bank_forks, block_commitment_cache, last_blockhash, gci) = self.setup_bank();
846 let target_tick_duration = gci.genesis_config.poh_config.target_tick_duration;
847 let target_slot_duration = target_tick_duration * gci.genesis_config.ticks_per_slot as u32;
848 let transport = start_local_server(
849 bank_forks.clone(),
850 block_commitment_cache.clone(),
851 target_tick_duration,
852 )
853 .await;
854 let banks_client = start_client(transport)
855 .await
856 .unwrap_or_else(|err| panic!("Failed to start banks client: {}", err));
857
858 tokio::spawn(async move {
862 loop {
863 tokio::time::sleep(target_slot_duration).await;
864 bank_forks
865 .read()
866 .unwrap()
867 .working_bank()
868 .register_recent_blockhash(&Hash::new_unique());
869 }
870 });
871
872 (banks_client, gci.mint_keypair, last_blockhash)
873 }
874
875 pub async fn start_with_context(self) -> ProgramTestContext {
880 let (bank_forks, block_commitment_cache, last_blockhash, gci) = self.setup_bank();
881 let target_tick_duration = gci.genesis_config.poh_config.target_tick_duration;
882 let transport = start_local_server(
883 bank_forks.clone(),
884 block_commitment_cache.clone(),
885 target_tick_duration,
886 )
887 .await;
888 let banks_client = start_client(transport)
889 .await
890 .unwrap_or_else(|err| panic!("Failed to start banks client: {}", err));
891
892 ProgramTestContext::new(
893 bank_forks,
894 block_commitment_cache,
895 banks_client,
896 last_blockhash,
897 gci,
898 )
899 }
900}
901
902#[async_trait]
903pub trait ProgramTestBanksClientExt {
904 #[deprecated(
908 since = "1.9.0",
909 note = "Please use `get_new_latest_blockhash `instead"
910 )]
911 async fn get_new_blockhash(&mut self, blockhash: &Hash) -> io::Result<(Hash, FeeCalculator)>;
912 async fn get_new_latest_blockhash(&mut self, blockhash: &Hash) -> io::Result<Hash>;
914}
915
916#[async_trait]
917impl ProgramTestBanksClientExt for BanksClient {
918 async fn get_new_blockhash(&mut self, blockhash: &Hash) -> io::Result<(Hash, FeeCalculator)> {
919 let mut num_retries = 0;
920 let start = Instant::now();
921 while start.elapsed().as_secs() < 5 {
922 #[allow(deprecated)]
923 if let Ok((fee_calculator, new_blockhash, _slot)) = self.get_fees().await {
924 if new_blockhash != *blockhash {
925 return Ok((new_blockhash, fee_calculator));
926 }
927 }
928 debug!("Got same blockhash ({:?}), will retry...", blockhash);
929
930 tokio::time::sleep(Duration::from_millis(200)).await;
931 num_retries += 1;
932 }
933
934 Err(io::Error::new(
935 io::ErrorKind::Other,
936 format!(
937 "Unable to get new blockhash after {}ms (retried {} times), stuck at {}",
938 start.elapsed().as_millis(),
939 num_retries,
940 blockhash
941 ),
942 ))
943 }
944
945 async fn get_new_latest_blockhash(&mut self, blockhash: &Hash) -> io::Result<Hash> {
946 let mut num_retries = 0;
947 let start = Instant::now();
948 while start.elapsed().as_secs() < 5 {
949 let new_blockhash = self.get_latest_blockhash().await?;
950 if new_blockhash != *blockhash {
951 return Ok(new_blockhash);
952 }
953 debug!("Got same blockhash ({:?}), will retry...", blockhash);
954
955 tokio::time::sleep(Duration::from_millis(200)).await;
956 num_retries += 1;
957 }
958
959 Err(io::Error::new(
960 io::ErrorKind::Other,
961 format!(
962 "Unable to get new blockhash after {}ms (retried {} times), stuck at {}",
963 start.elapsed().as_millis(),
964 num_retries,
965 blockhash
966 ),
967 ))
968 }
969}
970
971struct DroppableTask<T>(Arc<AtomicBool>, JoinHandle<T>);
972
973impl<T> Drop for DroppableTask<T> {
974 fn drop(&mut self) {
975 self.0.store(true, Ordering::Relaxed);
976 }
977}
978
979pub struct ProgramTestContext {
980 pub banks_client: BanksClient,
981 pub last_blockhash: Hash,
982 pub payer: Keypair,
983 genesis_config: GenesisConfig,
984 bank_forks: Arc<RwLock<BankForks>>,
985 block_commitment_cache: Arc<RwLock<BlockCommitmentCache>>,
986 _bank_task: DroppableTask<()>,
987}
988
989impl ProgramTestContext {
990 fn new(
991 bank_forks: Arc<RwLock<BankForks>>,
992 block_commitment_cache: Arc<RwLock<BlockCommitmentCache>>,
993 banks_client: BanksClient,
994 last_blockhash: Hash,
995 genesis_config_info: GenesisConfigInfo,
996 ) -> Self {
997 let running_bank_forks = bank_forks.clone();
1001 let target_tick_duration = genesis_config_info
1002 .genesis_config
1003 .poh_config
1004 .target_tick_duration;
1005 let target_slot_duration =
1006 target_tick_duration * genesis_config_info.genesis_config.ticks_per_slot as u32;
1007 let exit = Arc::new(AtomicBool::new(false));
1008 let bank_task = DroppableTask(
1009 exit.clone(),
1010 tokio::spawn(async move {
1011 loop {
1012 if exit.load(Ordering::Relaxed) {
1013 break;
1014 }
1015 tokio::time::sleep(target_slot_duration).await;
1016 running_bank_forks
1017 .read()
1018 .unwrap()
1019 .working_bank()
1020 .register_recent_blockhash(&Hash::new_unique());
1021 }
1022 }),
1023 );
1024
1025 Self {
1026 banks_client,
1027 last_blockhash,
1028 payer: genesis_config_info.mint_keypair,
1029 genesis_config: genesis_config_info.genesis_config,
1030 bank_forks,
1031 block_commitment_cache,
1032 _bank_task: bank_task,
1033 }
1034 }
1035
1036 pub fn genesis_config(&self) -> &GenesisConfig {
1037 &self.genesis_config
1038 }
1039
1040 pub fn increment_vote_account_credits(
1042 &mut self,
1043 vote_account_address: &Pubkey,
1044 number_of_credits: u64,
1045 ) {
1046 let bank_forks = self.bank_forks.read().unwrap();
1047 let bank = bank_forks.working_bank();
1048
1049 let mut vote_account = bank.get_account(vote_account_address).unwrap();
1051 let mut vote_state = VoteState::from(&vote_account).unwrap();
1052
1053 let epoch = bank.epoch();
1054 for _ in 0..number_of_credits {
1055 vote_state.increment_credits(epoch, 1);
1056 }
1057 let versioned = VoteStateVersions::new_current(vote_state);
1058 VoteState::to(&versioned, &mut vote_account).unwrap();
1059 bank.store_account(vote_account_address, &vote_account);
1060 }
1061
1062 pub fn set_account(&mut self, address: &Pubkey, account: &AccountSharedData) {
1069 let bank_forks = self.bank_forks.read().unwrap();
1070 let bank = bank_forks.working_bank();
1071 bank.store_account(address, account);
1072 }
1073
1074 pub fn set_sysvar<T: SysvarId + Sysvar>(&self, sysvar: &T) {
1081 let bank_forks = self.bank_forks.read().unwrap();
1082 let bank = bank_forks.working_bank();
1083 bank.set_sysvar_for_tests(sysvar);
1084 }
1085
1086 pub fn warp_to_slot(&mut self, warp_slot: Slot) -> Result<(), ProgramTestError> {
1088 let mut bank_forks = self.bank_forks.write().unwrap();
1089 let bank = bank_forks.working_bank();
1090
1091 bank.fill_bank_with_ticks_for_tests();
1094
1095 let working_slot = bank.slot();
1097 if warp_slot <= working_slot {
1098 return Err(ProgramTestError::InvalidWarpSlot);
1099 }
1100
1101 let pre_warp_slot = warp_slot - 1;
1105 let warp_bank = if pre_warp_slot == working_slot {
1106 bank.freeze();
1107 bank
1108 } else {
1109 bank_forks.insert(Bank::warp_from_parent(
1110 &bank,
1111 &Pubkey::default(),
1112 pre_warp_slot,
1113 ))
1114 };
1115 bank_forks.set_root(
1116 pre_warp_slot,
1117 &solana_runtime::accounts_background_service::AbsRequestSender::default(),
1118 Some(pre_warp_slot),
1119 );
1120
1121 bank_forks.insert(Bank::new_from_parent(
1123 &warp_bank,
1124 &Pubkey::default(),
1125 warp_slot,
1126 ));
1127
1128 let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
1131 w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
1136
1137 let bank = bank_forks.working_bank();
1138 self.last_blockhash = bank.last_blockhash();
1139 Ok(())
1140 }
1141
1142 pub async fn get_new_latest_blockhash(&mut self) -> io::Result<Hash> {
1144 let blockhash = self
1145 .banks_client
1146 .get_new_latest_blockhash(&self.last_blockhash)
1147 .await?;
1148 self.last_blockhash = blockhash;
1149 Ok(blockhash)
1150 }
1151}