Rework TestValidator API to be more like ProgramTest

This commit is contained in:
Michael Vines
2020-12-11 20:03:15 -08:00
committed by mergify[bot]
parent 5237da4e01
commit 9f2d154588
6 changed files with 383 additions and 175 deletions

View File

@ -13,35 +13,165 @@ use {
hardened_unpack::MAX_GENESIS_ARCHIVE_UNPACKED_SIZE,
},
solana_sdk::{
account::Account,
clock::DEFAULT_MS_PER_SLOT,
commitment_config::CommitmentConfig,
fee_calculator::FeeRateGovernor,
fee_calculator::{FeeCalculator, FeeRateGovernor},
hash::Hash,
native_token::sol_to_lamports,
pubkey::Pubkey,
rent::Rent,
signature::{read_keypair_file, write_keypair_file, Keypair, Signer},
},
std::{
fs::remove_dir_all,
net::SocketAddr,
path::{Path, PathBuf},
sync::Arc,
thread::sleep,
time::Duration,
collections::HashMap, fs::remove_dir_all, net::SocketAddr, path::PathBuf, sync::Arc,
thread::sleep, time::Duration,
},
};
pub struct TestValidatorGenesisConfig {
pub fee_rate_governor: FeeRateGovernor,
pub mint_address: Pubkey,
pub rent: Rent,
#[derive(Clone)]
pub struct ProgramInfo {
pub program_id: Pubkey,
pub loader: Pubkey,
pub program_path: PathBuf,
}
#[derive(Default)]
pub struct TestValidatorStartConfig {
pub preserve_ledger: bool,
pub rpc_config: JsonRpcConfig,
pub rpc_ports: Option<(u16, u16)>, // (JsonRpc, JsonRpcPubSub), None == random ports
pub struct TestValidatorGenesis {
fee_rate_governor: FeeRateGovernor,
ledger_path: Option<PathBuf>,
rent: Rent,
rpc_config: JsonRpcConfig,
rpc_ports: Option<(u16, u16)>, // (JsonRpc, JsonRpcPubSub), None == random ports
accounts: HashMap<Pubkey, Account>,
programs: Vec<ProgramInfo>,
}
impl TestValidatorGenesis {
pub fn ledger_path<P: Into<PathBuf>>(&mut self, ledger_path: P) -> &mut Self {
self.ledger_path = Some(ledger_path.into());
self
}
pub fn fee_rate_governor(&mut self, fee_rate_governor: FeeRateGovernor) -> &mut Self {
self.fee_rate_governor = fee_rate_governor;
self
}
pub fn rent(&mut self, rent: Rent) -> &mut Self {
self.rent = rent;
self
}
pub fn rpc_config(&mut self, rpc_config: JsonRpcConfig) -> &mut Self {
self.rpc_config = rpc_config;
self
}
pub fn rpc_port(&mut self, rpc_port: u16) -> &mut Self {
self.rpc_ports = Some((rpc_port, rpc_port + 1));
self
}
/// Add an account to the test environment
pub fn add_account(&mut self, address: Pubkey, account: Account) -> &mut Self {
self.accounts.insert(address, account);
self
}
/// Add an account to the test environment with the account data in the provided `filename`
pub fn add_account_with_file_data(
&mut self,
address: Pubkey,
lamports: u64,
owner: Pubkey,
filename: &str,
) -> &mut Self {
self.add_account(
address,
Account {
lamports,
data: solana_program_test::read_file(
solana_program_test::find_file(filename).unwrap_or_else(|| {
panic!("Unable to locate {}", filename);
}),
),
owner,
executable: false,
rent_epoch: 0,
},
)
}
/// Add an account to the test environment with the account data in the provided as a base 64
/// string
pub fn add_account_with_base64_data(
&mut self,
address: Pubkey,
lamports: u64,
owner: Pubkey,
data_base64: &str,
) -> &mut Self {
self.add_account(
address,
Account {
lamports,
data: base64::decode(data_base64)
.unwrap_or_else(|err| panic!("Failed to base64 decode: {}", err)),
owner,
executable: false,
rent_epoch: 0,
},
)
}
/// Add a BPF program to the test environment.
///
/// `program_name` will also used to locate the BPF shared object in the current or fixtures
/// directory.
pub fn add_program(&mut self, program_name: &str, program_id: Pubkey) -> &mut Self {
let program_path = solana_program_test::find_file(&format!("{}.so", program_name))
.unwrap_or_else(|| panic!("Unable to locate program {}", program_name));
self.programs.push(ProgramInfo {
program_id,
loader: solana_sdk::bpf_loader::id(),
program_path,
});
self
}
/// Add a list of programs to the test environment.
///pub fn add_programs_with_path<'a>(&'a mut self, programs: &[ProgramInfo]) -> &'a mut Self {
pub fn add_programs_with_path(&mut self, programs: &[ProgramInfo]) -> &mut Self {
for program in programs {
self.programs.push(program.clone());
}
self
}
/// Start a test validator with the address of the mint account that will receive tokens
/// created at genesis.
///
pub fn start_with_mint_address(
&self,
mint_address: Pubkey,
) -> Result<TestValidator, Box<dyn std::error::Error>> {
TestValidator::start(mint_address, self)
}
/// Start a test validator
///
/// Returns a new `TestValidator` as well as the keypair for the mint account that will receive tokens
/// created at genesis.
///
/// This function panics on initialization failure.
pub fn start(&self) -> (TestValidator, Keypair) {
let mint_keypair = Keypair::new();
TestValidator::start(mint_keypair.pubkey(), self)
.map(|test_validator| (test_validator, mint_keypair))
.expect("Test validator failed to start")
}
}
pub struct TestValidator {
@ -56,106 +186,92 @@ pub struct TestValidator {
}
impl TestValidator {
/// The default test validator is intended to be generically suitable for unit testing.
///
/// It uses a unique temporary ledger that is deleted on `close` and randomly assigned ports.
/// All test tokens will be minted into `mint_address`
///
/// This function panics on initialization failure.
pub fn new(mint_address: Pubkey) -> Self {
let ledger_path = Self::initialize_ledger(
None,
TestValidatorGenesisConfig {
fee_rate_governor: FeeRateGovernor::default(),
mint_address,
rent: Rent::default(),
},
)
.unwrap();
Self::start(&ledger_path, TestValidatorStartConfig::default()).unwrap()
}
/// Create a `TestValidator` with no transaction fees and minimal rent.
/// Create and start a `TestValidator` with no transaction fees and minimal rent.
///
/// This function panics on initialization failure.
pub fn with_no_fees(mint_address: Pubkey) -> Self {
let ledger_path = Self::initialize_ledger(
None,
TestValidatorGenesisConfig {
fee_rate_governor: FeeRateGovernor::new(0, 0),
mint_address,
rent: Rent {
lamports_per_byte_year: 1,
exemption_threshold: 1.0,
..Rent::default()
},
},
)
.unwrap();
Self::start(&ledger_path, TestValidatorStartConfig::default()).unwrap()
TestValidatorGenesis::default()
.fee_rate_governor(FeeRateGovernor::new(0, 0))
.rent(Rent {
lamports_per_byte_year: 1,
exemption_threshold: 1.0,
..Rent::default()
})
.start_with_mint_address(mint_address)
.expect("validator start failed")
}
/// Create a `TestValidator` with custom transaction fees and minimal rent.
/// Create and start a `TestValidator` with custom transaction fees and minimal rent.
///
/// This function panics on initialization failure.
pub fn with_custom_fees(mint_address: Pubkey, target_lamports_per_signature: u64) -> Self {
let ledger_path = Self::initialize_ledger(
None,
TestValidatorGenesisConfig {
fee_rate_governor: FeeRateGovernor::new(target_lamports_per_signature, 0),
mint_address,
rent: Rent {
lamports_per_byte_year: 1,
exemption_threshold: 1.0,
..Rent::default()
},
},
)
.unwrap();
Self::start(&ledger_path, TestValidatorStartConfig::default()).unwrap()
TestValidatorGenesis::default()
.fee_rate_governor(FeeRateGovernor::new(target_lamports_per_signature, 0))
.rent(Rent {
lamports_per_byte_year: 1,
exemption_threshold: 1.0,
..Rent::default()
})
.start_with_mint_address(mint_address)
.expect("validator start failed")
}
/// Initialize the test validator's ledger directory
/// Initialize the ledger directory
///
/// If `ledger_path` is `None`, a temporary ledger will be created. Otherwise the ledger will
/// be initialized in the provided directory.
/// be initialized in the provided directory if it doesn't already exist.
///
/// Returns the path to the ledger directory.
pub fn initialize_ledger(
ledger_path: Option<&Path>,
config: TestValidatorGenesisConfig,
fn initialize_ledger(
mint_address: Pubkey,
config: &TestValidatorGenesis,
) -> Result<PathBuf, Box<dyn std::error::Error>> {
let TestValidatorGenesisConfig {
fee_rate_governor,
mint_address,
rent,
} = config;
let validator_identity_keypair = Keypair::new();
let validator_identity = Keypair::new();
let validator_vote_account = Keypair::new();
let validator_stake_account = Keypair::new();
let validator_identity_lamports = sol_to_lamports(500.);
let validator_stake_lamports = sol_to_lamports(1_000_000.);
let mint_lamports = sol_to_lamports(500_000_000.);
let initial_accounts = solana_program_test::programs::spl_programs(&rent);
let mut accounts = config.accounts.clone();
for (address, account) in solana_program_test::programs::spl_programs(&config.rent) {
accounts.entry(address).or_insert(account);
}
for program in &config.programs {
let data = solana_program_test::read_file(&program.program_path);
accounts.insert(
program.program_id,
Account {
lamports: Rent::default().minimum_balance(data.len()).min(1),
data,
owner: program.loader,
executable: true,
rent_epoch: 0,
},
);
}
let genesis_config = create_genesis_config_with_leader_ex(
mint_lamports,
&mint_address,
&validator_identity_keypair.pubkey(),
&validator_identity.pubkey(),
&validator_vote_account.pubkey(),
&validator_stake_account.pubkey(),
validator_stake_lamports,
validator_identity_lamports,
fee_rate_governor,
rent,
config.fee_rate_governor.clone(),
config.rent,
solana_sdk::genesis_config::ClusterType::Development,
initial_accounts,
accounts.into_iter().collect(),
);
let ledger_path = match ledger_path {
let ledger_path = match &config.ledger_path {
None => create_new_tmp_ledger!(&genesis_config).0,
Some(ledger_path) => {
if ledger_path.join("validator-keypair.json").exists() {
return Ok(ledger_path.to_path_buf());
}
let _ = create_new_ledger(
ledger_path,
&genesis_config,
@ -174,7 +290,7 @@ impl TestValidator {
};
write_keypair_file(
&validator_identity_keypair,
&validator_identity,
ledger_path.join("validator-keypair.json").to_str().unwrap(),
)?;
write_keypair_file(
@ -189,11 +305,14 @@ impl TestValidator {
}
/// Starts a TestValidator at the provided ledger directory
pub fn start(
ledger_path: &Path,
config: TestValidatorStartConfig,
fn start(
mint_address: Pubkey,
config: &TestValidatorGenesis,
) -> Result<Self, Box<dyn std::error::Error>> {
let validator_identity_keypair =
let preserve_ledger = config.ledger_path.is_some();
let ledger_path = TestValidator::initialize_ledger(mint_address, config)?;
let validator_identity =
read_keypair_file(ledger_path.join("validator-keypair.json").to_str().unwrap())?;
let validator_vote_account = read_keypair_file(
ledger_path
@ -202,7 +321,7 @@ impl TestValidator {
.unwrap(),
)?;
let mut node = Node::new_localhost_with_pubkey(&validator_identity_keypair.pubkey());
let mut node = Node::new_localhost_with_pubkey(&validator_identity.pubkey());
if let Some((rpc, rpc_pubsub)) = config.rpc_ports {
node.info.rpc = SocketAddr::new(node.info.gossip.ip(), rpc);
node.info.rpc_pubsub = SocketAddr::new(node.info.gossip.ip(), rpc_pubsub);
@ -216,7 +335,7 @@ impl TestValidator {
let validator_config = ValidatorConfig {
rpc_addrs: Some((node.info.rpc, node.info.rpc_pubsub)),
rpc_config: config.rpc_config,
rpc_config: config.rpc_config.clone(),
accounts_hash_interval_slots: 100,
account_paths: vec![ledger_path.join("accounts")],
poh_verify: false, // Skip PoH verification of ledger on startup for speed
@ -232,7 +351,7 @@ impl TestValidator {
let validator = Some(Validator::new(
node,
&Arc::new(validator_identity_keypair),
&Arc::new(validator_identity),
&ledger_path,
&validator_vote_account.pubkey(),
vec![Arc::new(validator_vote_account)],
@ -244,12 +363,12 @@ impl TestValidator {
// test validators concurrently...
discover_cluster(&gossip, 1).expect("TestValidator startup failed");
// This is a hack to delay until the single gossip commitment fees are non-zero for test
// consistency
// (fees from genesis are zero until the first block with a transaction in it is completed)
// This is a hack to delay until the fees are non-zero for test consistency
// (fees from genesis are zero until the first block with a transaction in it is completed
// due to a bug in the Bank)
{
let rpc_client =
RpcClient::new_with_commitment(rpc_url.clone(), CommitmentConfig::single_gossip());
RpcClient::new_with_commitment(rpc_url.clone(), CommitmentConfig::recent());
let fee_rate_governor = rpc_client
.get_fee_rate_governor()
.expect("get_fee_rate_governor")
@ -268,8 +387,8 @@ impl TestValidator {
}
Ok(TestValidator {
ledger_path: ledger_path.to_path_buf(),
preserve_ledger: false,
ledger_path,
preserve_ledger,
rpc_pubsub_url,
rpc_url,
gossip,
@ -279,30 +398,42 @@ impl TestValidator {
})
}
/// Return the test validator's TPU address
/// Return the validator's TPU address
pub fn tpu(&self) -> &SocketAddr {
&self.tpu
}
/// Return the test validator's Gossip address
/// Return the validator's Gossip address
pub fn gossip(&self) -> &SocketAddr {
&self.gossip
}
/// Return the test validator's JSON RPC URL
/// Return the validator's JSON RPC URL
pub fn rpc_url(&self) -> String {
self.rpc_url.clone()
}
/// Return the test validator's JSON RPC PubSub URL
/// Return the validator's JSON RPC PubSub URL
pub fn rpc_pubsub_url(&self) -> String {
self.rpc_pubsub_url.clone()
}
/// Return the vote account address of the validator
/// Return the validator's vote account address
pub fn vote_account_address(&self) -> Pubkey {
self.vote_account_address
}
/// Return an RpcClient for the validator. As a convenience, also return a recent blockhash and
/// associated fee calculator
pub fn rpc_client(&self) -> (RpcClient, Hash, FeeCalculator) {
let rpc_client =
RpcClient::new_with_commitment(self.rpc_url.clone(), CommitmentConfig::recent());
let (recent_blockhash, fee_calculator) = rpc_client
.get_recent_blockhash()
.expect("get_recent_blockhash");
(rpc_client, recent_blockhash, fee_calculator)
}
}
impl Drop for TestValidator {