cli: solana-tokens, validate inputs gracefully (#33926)
* cli: solana-tokens, refactor imports * cli: solana-tokens, validate inputs gracefully * change Allocation struct field types to simplify things * fix typos, apply some review suggestions * preserve backward compatibility for public APIs * apply latest review suggestions * address next batch of review comments --------- Co-authored-by: norwnd <norwnd>
This commit is contained in:
norwnd
GitHub
parent
808f67aead
commit
1c00d5d81a
|
|
@ -42,6 +42,7 @@ use {
|
|||
std::{
|
||||
cmp::{self},
|
||||
io,
|
||||
str::FromStr,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
|
|
@ -51,6 +52,7 @@ use {
|
|||
},
|
||||
};
|
||||
|
||||
/// Allocation is a helper (mostly for tests), prefer using TypedAllocation instead when possible.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Allocation {
|
||||
pub recipient: String,
|
||||
|
|
@ -58,6 +60,14 @@ pub struct Allocation {
|
|||
pub lockup_date: String,
|
||||
}
|
||||
|
||||
/// TypedAllocation is same as Allocation but contains typed fields.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
|
||||
pub struct TypedAllocation {
|
||||
pub recipient: Pubkey,
|
||||
pub amount: u64,
|
||||
pub lockup_date: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum FundingSource {
|
||||
FeePayer,
|
||||
|
|
@ -98,8 +108,20 @@ type StakeExtras = Vec<(Keypair, Option<DateTime<Utc>>)>;
|
|||
pub enum Error {
|
||||
#[error("I/O error")]
|
||||
IoError(#[from] io::Error),
|
||||
#[error("CSV file seems to be empty")]
|
||||
CsvIsEmptyError,
|
||||
#[error("CSV error")]
|
||||
CsvError(#[from] csv::Error),
|
||||
#[error("Bad input data for pubkey: {input}, error: {err}")]
|
||||
BadInputPubkeyError {
|
||||
input: String,
|
||||
err: pubkey::ParsePubkeyError,
|
||||
},
|
||||
#[error("Bad input data for lockup date: {input}, error: {err}")]
|
||||
BadInputLockupDate {
|
||||
input: String,
|
||||
err: chrono::ParseError,
|
||||
},
|
||||
#[error("PickleDb error")]
|
||||
PickleDbError(#[from] pickledb::error::Error),
|
||||
#[error("Transport error")]
|
||||
|
|
@ -118,15 +140,15 @@ pub enum Error {
|
|||
ExitSignal,
|
||||
}
|
||||
|
||||
fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
|
||||
fn merge_allocations(allocations: &[TypedAllocation]) -> Vec<TypedAllocation> {
|
||||
let mut allocation_map = IndexMap::new();
|
||||
for allocation in allocations {
|
||||
allocation_map
|
||||
.entry(&allocation.recipient)
|
||||
.or_insert(Allocation {
|
||||
recipient: allocation.recipient.clone(),
|
||||
.or_insert(TypedAllocation {
|
||||
recipient: allocation.recipient,
|
||||
amount: 0,
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
})
|
||||
.amount += allocation.amount;
|
||||
}
|
||||
|
|
@ -134,13 +156,13 @@ fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
|
|||
}
|
||||
|
||||
/// Return true if the recipient and lockups are the same
|
||||
fn has_same_recipient(allocation: &Allocation, transaction_info: &TransactionInfo) -> bool {
|
||||
allocation.recipient == transaction_info.recipient.to_string()
|
||||
&& allocation.lockup_date.parse().ok() == transaction_info.lockup_date
|
||||
fn has_same_recipient(allocation: &TypedAllocation, transaction_info: &TransactionInfo) -> bool {
|
||||
allocation.recipient == transaction_info.recipient
|
||||
&& allocation.lockup_date == transaction_info.lockup_date
|
||||
}
|
||||
|
||||
fn apply_previous_transactions(
|
||||
allocations: &mut Vec<Allocation>,
|
||||
allocations: &mut Vec<TypedAllocation>,
|
||||
transaction_infos: &[TransactionInfo],
|
||||
) {
|
||||
for transaction_info in transaction_infos {
|
||||
|
|
@ -179,7 +201,7 @@ fn transfer<S: Signer>(
|
|||
}
|
||||
|
||||
fn distribution_instructions(
|
||||
allocation: &Allocation,
|
||||
allocation: &TypedAllocation,
|
||||
new_stake_account_address: &Pubkey,
|
||||
args: &DistributeTokensArgs,
|
||||
lockup_date: Option<DateTime<Utc>>,
|
||||
|
|
@ -193,7 +215,7 @@ fn distribution_instructions(
|
|||
// No stake args; a simple token transfer.
|
||||
None => {
|
||||
let from = args.sender_keypair.pubkey();
|
||||
let to = allocation.recipient.parse().unwrap();
|
||||
let to = allocation.recipient;
|
||||
let lamports = allocation.amount;
|
||||
let instruction = system_instruction::transfer(&from, &to, lamports);
|
||||
vec![instruction]
|
||||
|
|
@ -203,7 +225,7 @@ fn distribution_instructions(
|
|||
Some(stake_args) => {
|
||||
let unlocked_sol = stake_args.unlocked_sol;
|
||||
let sender_pubkey = args.sender_keypair.pubkey();
|
||||
let recipient = allocation.recipient.parse().unwrap();
|
||||
let recipient = allocation.recipient;
|
||||
|
||||
let mut instructions = match &stake_args.sender_stake_args {
|
||||
// No source stake account, so create a recipient stake account directly.
|
||||
|
|
@ -304,7 +326,7 @@ fn distribution_instructions(
|
|||
fn build_messages(
|
||||
client: &RpcClient,
|
||||
db: &mut PickleDb,
|
||||
allocations: &[Allocation],
|
||||
allocations: &[TypedAllocation],
|
||||
args: &DistributeTokensArgs,
|
||||
exit: Arc<AtomicBool>,
|
||||
messages: &mut Vec<Message>,
|
||||
|
|
@ -318,7 +340,7 @@ fn build_messages(
|
|||
let associated_token_addresses = allocation_chunk
|
||||
.iter()
|
||||
.map(|x| {
|
||||
let wallet_address = x.recipient.parse().unwrap();
|
||||
let wallet_address = x.recipient;
|
||||
get_associated_token_address(&wallet_address, &spl_token_args.mint)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
|
@ -333,11 +355,7 @@ fn build_messages(
|
|||
return Err(Error::ExitSignal);
|
||||
}
|
||||
let new_stake_account_keypair = Keypair::new();
|
||||
let lockup_date = if allocation.lockup_date.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap())
|
||||
};
|
||||
let lockup_date = allocation.lockup_date;
|
||||
|
||||
let do_create_associated_token_account = if let Some(spl_token_args) = &args.spl_token_args
|
||||
{
|
||||
|
|
@ -382,7 +400,7 @@ fn build_messages(
|
|||
fn send_messages(
|
||||
client: &RpcClient,
|
||||
db: &mut PickleDb,
|
||||
allocations: &[Allocation],
|
||||
allocations: &[TypedAllocation],
|
||||
args: &DistributeTokensArgs,
|
||||
exit: Arc<AtomicBool>,
|
||||
messages: Vec<Message>,
|
||||
|
|
@ -404,7 +422,7 @@ fn send_messages(
|
|||
signers.push(&*sender_stake_args.stake_authority);
|
||||
signers.push(&*sender_stake_args.withdraw_authority);
|
||||
signers.push(&new_stake_account_keypair);
|
||||
if !allocation.lockup_date.is_empty() {
|
||||
if allocation.lockup_date.is_some() {
|
||||
if let Some(lockup_authority) = &sender_stake_args.lockup_authority {
|
||||
signers.push(&**lockup_authority);
|
||||
} else {
|
||||
|
|
@ -435,7 +453,7 @@ fn send_messages(
|
|||
args.stake_args.as_ref().map(|_| &new_stake_account_address);
|
||||
db::set_transaction_info(
|
||||
db,
|
||||
&allocation.recipient.parse().unwrap(),
|
||||
&allocation.recipient,
|
||||
allocation.amount,
|
||||
&transaction,
|
||||
new_stake_account_address_option,
|
||||
|
|
@ -455,7 +473,7 @@ fn send_messages(
|
|||
fn distribute_allocations(
|
||||
client: &RpcClient,
|
||||
db: &mut PickleDb,
|
||||
allocations: &[Allocation],
|
||||
allocations: &[TypedAllocation],
|
||||
args: &DistributeTokensArgs,
|
||||
exit: Arc<AtomicBool>,
|
||||
) -> Result<(), Error> {
|
||||
|
|
@ -490,63 +508,91 @@ fn distribute_allocations(
|
|||
fn read_allocations(
|
||||
input_csv: &str,
|
||||
transfer_amount: Option<u64>,
|
||||
require_lockup_heading: bool,
|
||||
with_lockup: bool,
|
||||
raw_amount: bool,
|
||||
) -> io::Result<Vec<Allocation>> {
|
||||
) -> Result<Vec<TypedAllocation>, Error> {
|
||||
let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
|
||||
let allocations = if let Some(amount) = transfer_amount {
|
||||
let recipients: Vec<String> = rdr
|
||||
.deserialize()
|
||||
.map(|recipient| recipient.unwrap())
|
||||
.collect();
|
||||
recipients
|
||||
.into_iter()
|
||||
.map(|recipient| Allocation {
|
||||
recipient,
|
||||
amount,
|
||||
lockup_date: "".to_string(),
|
||||
rdr.deserialize()
|
||||
.map(|recipient| {
|
||||
let recipient: String = recipient?;
|
||||
let recipient =
|
||||
Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
|
||||
input: recipient,
|
||||
err,
|
||||
})?;
|
||||
Ok(TypedAllocation {
|
||||
recipient,
|
||||
amount,
|
||||
lockup_date: None,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
} else if require_lockup_heading {
|
||||
let recipients: Vec<(String, f64, String)> = rdr
|
||||
.deserialize()
|
||||
.map(|recipient| recipient.unwrap())
|
||||
.collect();
|
||||
recipients
|
||||
.into_iter()
|
||||
.map(|(recipient, amount, lockup_date)| Allocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(amount),
|
||||
lockup_date,
|
||||
.collect::<Result<Vec<TypedAllocation>, Error>>()?
|
||||
} else if with_lockup {
|
||||
// We only support SOL token in "require lockup" mode.
|
||||
rdr.deserialize()
|
||||
.map(|recipient| {
|
||||
let (recipient, amount, lockup_date): (String, f64, String) = recipient?;
|
||||
let recipient =
|
||||
Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
|
||||
input: recipient,
|
||||
err,
|
||||
})?;
|
||||
let lockup_date = if !lockup_date.is_empty() {
|
||||
let lockup_date = lockup_date.parse::<DateTime<Utc>>().map_err(|err| {
|
||||
Error::BadInputLockupDate {
|
||||
input: lockup_date,
|
||||
err,
|
||||
}
|
||||
})?;
|
||||
Some(lockup_date)
|
||||
} else {
|
||||
// empty lockup date means no lockup, it's okay to have only some lockups specified
|
||||
None
|
||||
};
|
||||
Ok(TypedAllocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(amount),
|
||||
lockup_date,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
.collect::<Result<Vec<TypedAllocation>, Error>>()?
|
||||
} else if raw_amount {
|
||||
let recipients: Vec<(String, u64)> = rdr
|
||||
.deserialize()
|
||||
.map(|recipient| recipient.unwrap())
|
||||
.collect();
|
||||
recipients
|
||||
.into_iter()
|
||||
.map(|(recipient, amount)| Allocation {
|
||||
recipient,
|
||||
amount,
|
||||
lockup_date: "".to_string(),
|
||||
rdr.deserialize()
|
||||
.map(|recipient| {
|
||||
let (recipient, amount): (String, u64) = recipient?;
|
||||
let recipient =
|
||||
Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
|
||||
input: recipient,
|
||||
err,
|
||||
})?;
|
||||
Ok(TypedAllocation {
|
||||
recipient,
|
||||
amount,
|
||||
lockup_date: None,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
.collect::<Result<Vec<TypedAllocation>, Error>>()?
|
||||
} else {
|
||||
let recipients: Vec<(String, f64)> = rdr
|
||||
.deserialize()
|
||||
.map(|recipient| recipient.unwrap())
|
||||
.collect();
|
||||
recipients
|
||||
.into_iter()
|
||||
.map(|(recipient, amount)| Allocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(amount),
|
||||
lockup_date: "".to_string(),
|
||||
rdr.deserialize()
|
||||
.map(|recipient| {
|
||||
let (recipient, amount): (String, f64) = recipient?;
|
||||
let recipient =
|
||||
Pubkey::from_str(&recipient).map_err(|err| Error::BadInputPubkeyError {
|
||||
input: recipient,
|
||||
err,
|
||||
})?;
|
||||
Ok(TypedAllocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(amount),
|
||||
lockup_date: None,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
.collect::<Result<Vec<TypedAllocation>, Error>>()?
|
||||
};
|
||||
if allocations.is_empty() {
|
||||
return Err(Error::CsvIsEmptyError);
|
||||
}
|
||||
Ok(allocations)
|
||||
}
|
||||
|
||||
|
|
@ -566,11 +612,11 @@ pub fn process_allocations(
|
|||
args: &DistributeTokensArgs,
|
||||
exit: Arc<AtomicBool>,
|
||||
) -> Result<Option<usize>, Error> {
|
||||
let require_lockup_heading = args.stake_args.is_some();
|
||||
let mut allocations: Vec<Allocation> = read_allocations(
|
||||
let with_lockup = args.stake_args.is_some();
|
||||
let mut allocations: Vec<TypedAllocation> = read_allocations(
|
||||
&args.input_csv,
|
||||
args.transfer_amount,
|
||||
require_lockup_heading,
|
||||
with_lockup,
|
||||
args.spl_token_args.is_some(),
|
||||
)?;
|
||||
|
||||
|
|
@ -773,7 +819,7 @@ pub fn get_fee_estimate_for_messages(
|
|||
|
||||
fn check_payer_balances(
|
||||
messages: &[Message],
|
||||
allocations: &[Allocation],
|
||||
allocations: &[TypedAllocation],
|
||||
client: &RpcClient,
|
||||
args: &DistributeTokensArgs,
|
||||
) -> Result<(), Error> {
|
||||
|
|
@ -857,7 +903,7 @@ pub fn process_balances(
|
|||
args: &BalancesArgs,
|
||||
exit: Arc<AtomicBool>,
|
||||
) -> Result<(), Error> {
|
||||
let allocations: Vec<Allocation> =
|
||||
let allocations: Vec<TypedAllocation> =
|
||||
read_allocations(&args.input_csv, None, false, args.spl_token_args.is_some())?;
|
||||
let allocations = merge_allocations(&allocations);
|
||||
|
||||
|
|
@ -885,7 +931,7 @@ pub fn process_balances(
|
|||
if let Some(spl_token_args) = &args.spl_token_args {
|
||||
print_token_balances(client, allocation, spl_token_args)?;
|
||||
} else {
|
||||
let address: Pubkey = allocation.recipient.parse().unwrap();
|
||||
let address: Pubkey = allocation.recipient;
|
||||
let expected = lamports_to_sol(allocation.amount);
|
||||
let actual = lamports_to_sol(client.get_balance(&address).unwrap());
|
||||
println!(
|
||||
|
|
@ -909,9 +955,13 @@ pub fn process_transaction_log(args: &TransactionLogArgs) -> Result<(), Error> {
|
|||
|
||||
use {
|
||||
crate::db::check_output_file,
|
||||
solana_sdk::{pubkey::Pubkey, signature::Keypair},
|
||||
solana_sdk::{
|
||||
pubkey::{self, Pubkey},
|
||||
signature::Keypair,
|
||||
},
|
||||
tempfile::{tempdir, NamedTempFile},
|
||||
};
|
||||
|
||||
pub fn test_process_distribute_tokens_with_client(
|
||||
client: &RpcClient,
|
||||
sender_keypair: Keypair,
|
||||
|
|
@ -939,7 +989,7 @@ pub fn test_process_distribute_tokens_with_client(
|
|||
} else {
|
||||
sol_to_lamports(1000.0)
|
||||
};
|
||||
let alice_pubkey = solana_sdk::pubkey::new_rand();
|
||||
let alice_pubkey = pubkey::new_rand();
|
||||
let allocations_file = NamedTempFile::new().unwrap();
|
||||
let input_csv = allocations_file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file);
|
||||
|
|
@ -1039,7 +1089,7 @@ pub fn test_process_create_stake_with_client(client: &RpcClient, sender_keypair:
|
|||
.unwrap();
|
||||
|
||||
let expected_amount = sol_to_lamports(1000.0);
|
||||
let alice_pubkey = solana_sdk::pubkey::new_rand();
|
||||
let alice_pubkey = pubkey::new_rand();
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
|
|
@ -1161,7 +1211,7 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
|
|||
.unwrap();
|
||||
|
||||
let expected_amount = sol_to_lamports(1000.0);
|
||||
let alice_pubkey = solana_sdk::pubkey::new_rand();
|
||||
let alice_pubkey = pubkey::new_rand();
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
|
|
@ -1328,16 +1378,27 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn test_read_allocations() {
|
||||
let alice_pubkey = solana_sdk::pubkey::new_rand();
|
||||
let allocation = Allocation {
|
||||
recipient: alice_pubkey.to_string(),
|
||||
let alice_pubkey = pubkey::new_rand();
|
||||
let allocation = TypedAllocation {
|
||||
recipient: alice_pubkey,
|
||||
amount: 42,
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
};
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
wtr.serialize(&allocation).unwrap();
|
||||
wtr.serialize((
|
||||
"recipient".to_string(),
|
||||
"amount".to_string(),
|
||||
"require_lockup".to_string(),
|
||||
))
|
||||
.unwrap();
|
||||
wtr.serialize((
|
||||
allocation.recipient.to_string(),
|
||||
allocation.amount,
|
||||
allocation.lockup_date,
|
||||
))
|
||||
.unwrap();
|
||||
wtr.flush().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
|
|
@ -1345,10 +1406,10 @@ mod tests {
|
|||
vec![allocation]
|
||||
);
|
||||
|
||||
let allocation_sol = Allocation {
|
||||
recipient: alice_pubkey.to_string(),
|
||||
let allocation_sol = TypedAllocation {
|
||||
recipient: alice_pubkey,
|
||||
amount: sol_to_lamports(42.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
|
|
@ -1367,8 +1428,8 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn test_read_allocations_no_lockup() {
|
||||
let pubkey0 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey1 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey0 = pubkey::new_rand();
|
||||
let pubkey1 = pubkey::new_rand();
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
|
|
@ -1379,15 +1440,15 @@ mod tests {
|
|||
wtr.flush().unwrap();
|
||||
|
||||
let expected_allocations = vec![
|
||||
Allocation {
|
||||
recipient: pubkey0.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: pubkey0,
|
||||
amount: sol_to_lamports(42.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
Allocation {
|
||||
recipient: pubkey1.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: pubkey1,
|
||||
amount: sol_to_lamports(43.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
];
|
||||
assert_eq!(
|
||||
|
|
@ -1397,42 +1458,210 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic]
|
||||
fn test_read_allocations_malformed() {
|
||||
let pubkey0 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey1 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey0 = pubkey::new_rand();
|
||||
let pubkey1 = pubkey::new_rand();
|
||||
|
||||
// Empty file.
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(&file);
|
||||
wtr.flush().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
let got = read_allocations(&input_csv, None, false, false);
|
||||
assert!(matches!(got, Err(Error::CsvIsEmptyError)));
|
||||
|
||||
// Missing 2nd column.
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(&file);
|
||||
wtr.serialize("recipient".to_string()).unwrap();
|
||||
wtr.serialize(pubkey0.to_string()).unwrap();
|
||||
wtr.serialize(pubkey1.to_string()).unwrap();
|
||||
wtr.flush().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got = read_allocations(&input_csv, None, false, false);
|
||||
assert!(matches!(got, Err(Error::CsvError(..))));
|
||||
|
||||
// Missing 3rd column.
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(&file);
|
||||
wtr.serialize(("recipient".to_string(), "amount".to_string()))
|
||||
.unwrap();
|
||||
wtr.serialize((&pubkey0.to_string(), 42.0)).unwrap();
|
||||
wtr.serialize((&pubkey1.to_string(), 43.0)).unwrap();
|
||||
wtr.serialize((pubkey0.to_string(), "42.0".to_string()))
|
||||
.unwrap();
|
||||
wtr.serialize((pubkey1.to_string(), "43.0".to_string()))
|
||||
.unwrap();
|
||||
wtr.flush().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got = read_allocations(&input_csv, None, true, false);
|
||||
assert!(matches!(got, Err(Error::CsvError(..))));
|
||||
|
||||
let expected_allocations = vec![
|
||||
Allocation {
|
||||
recipient: pubkey0.to_string(),
|
||||
amount: sol_to_lamports(42.0),
|
||||
lockup_date: "".to_string(),
|
||||
},
|
||||
Allocation {
|
||||
recipient: pubkey1.to_string(),
|
||||
amount: sol_to_lamports(43.0),
|
||||
lockup_date: "".to_string(),
|
||||
},
|
||||
];
|
||||
assert_eq!(
|
||||
read_allocations(&input_csv, None, true, false).unwrap(),
|
||||
expected_allocations
|
||||
let generate_csv_file = |header: (String, String, String),
|
||||
data: Vec<(String, String, String)>,
|
||||
file: &NamedTempFile| {
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
wtr.serialize(header).unwrap();
|
||||
wtr.serialize(&data[0]).unwrap();
|
||||
wtr.serialize(&data[1]).unwrap();
|
||||
wtr.flush().unwrap();
|
||||
};
|
||||
|
||||
let default_header = (
|
||||
"recipient".to_string(),
|
||||
"amount".to_string(),
|
||||
"require_lockup".to_string(),
|
||||
);
|
||||
|
||||
// Bad pubkey (default).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(pubkey0.to_string(), "42.0".to_string(), "".to_string()),
|
||||
("bad pubkey".to_string(), "43.0".to_string(), "".to_string()),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got_err = read_allocations(&input_csv, None, false, false).unwrap_err();
|
||||
assert!(
|
||||
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
|
||||
);
|
||||
// Bad pubkey (with transfer amount).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(pubkey0.to_string(), "42.0".to_string(), "".to_string()),
|
||||
("bad pubkey".to_string(), "43.0".to_string(), "".to_string()),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got_err = read_allocations(&input_csv, Some(123), false, false).unwrap_err();
|
||||
assert!(
|
||||
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
|
||||
);
|
||||
// Bad pubkey (with require lockup).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(
|
||||
pubkey0.to_string(),
|
||||
"42.0".to_string(),
|
||||
"2021-02-07T00:00:00Z".to_string(),
|
||||
),
|
||||
(
|
||||
"bad pubkey".to_string(),
|
||||
"43.0".to_string(),
|
||||
"2021-02-07T00:00:00Z".to_string(),
|
||||
),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got_err = read_allocations(&input_csv, None, true, false).unwrap_err();
|
||||
assert!(
|
||||
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
|
||||
);
|
||||
// Bad pubkey (with raw amount).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(pubkey0.to_string(), "42".to_string(), "".to_string()),
|
||||
("bad pubkey".to_string(), "43".to_string(), "".to_string()),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got_err = read_allocations(&input_csv, None, false, true).unwrap_err();
|
||||
assert!(
|
||||
matches!(got_err, Error::BadInputPubkeyError { input, .. } if input == *"bad pubkey")
|
||||
);
|
||||
|
||||
// Bad value in 2nd column (default).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(
|
||||
pubkey0.to_string(),
|
||||
"bad amount".to_string(),
|
||||
"".to_string(),
|
||||
),
|
||||
(
|
||||
pubkey1.to_string(),
|
||||
"43.0".to_string().to_string(),
|
||||
"".to_string(),
|
||||
),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got = read_allocations(&input_csv, None, false, false);
|
||||
assert!(matches!(got, Err(Error::CsvError(..))));
|
||||
// Bad value in 2nd column (with require lockup).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(
|
||||
pubkey0.to_string(),
|
||||
"bad amount".to_string(),
|
||||
"".to_string(),
|
||||
),
|
||||
(pubkey1.to_string(), "43.0".to_string(), "".to_string()),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got = read_allocations(&input_csv, None, true, false);
|
||||
assert!(matches!(got, Err(Error::CsvError(..))));
|
||||
// Bad value in 2nd column (with raw amount).
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(pubkey0.to_string(), "42".to_string(), "".to_string()),
|
||||
(pubkey1.to_string(), "43.0".to_string(), "".to_string()), // bad raw amount
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got = read_allocations(&input_csv, None, false, true);
|
||||
assert!(matches!(got, Err(Error::CsvError(..))));
|
||||
|
||||
// Bad value in 3rd column.
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
generate_csv_file(
|
||||
default_header.clone(),
|
||||
vec![
|
||||
(
|
||||
pubkey0.to_string(),
|
||||
"42.0".to_string(),
|
||||
"2021-01-07T00:00:00Z".to_string(),
|
||||
),
|
||||
(
|
||||
pubkey1.to_string(),
|
||||
"43.0".to_string(),
|
||||
"bad lockup date".to_string(),
|
||||
),
|
||||
],
|
||||
&file,
|
||||
);
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let got_err = read_allocations(&input_csv, None, true, false).unwrap_err();
|
||||
assert!(
|
||||
matches!(got_err, Error::BadInputLockupDate { input, .. } if input == *"bad lockup date")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_allocations_transfer_amount() {
|
||||
let pubkey0 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey1 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey2 = solana_sdk::pubkey::new_rand();
|
||||
let pubkey0 = pubkey::new_rand();
|
||||
let pubkey1 = pubkey::new_rand();
|
||||
let pubkey2 = pubkey::new_rand();
|
||||
let file = NamedTempFile::new().unwrap();
|
||||
let input_csv = file.path().to_str().unwrap().to_string();
|
||||
let mut wtr = csv::WriterBuilder::new().from_writer(file);
|
||||
|
|
@ -1445,20 +1674,20 @@ mod tests {
|
|||
let amount = sol_to_lamports(1.5);
|
||||
|
||||
let expected_allocations = vec![
|
||||
Allocation {
|
||||
recipient: pubkey0.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: pubkey0,
|
||||
amount,
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
Allocation {
|
||||
recipient: pubkey1.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: pubkey1,
|
||||
amount,
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
Allocation {
|
||||
recipient: pubkey2.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: pubkey2,
|
||||
amount,
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
];
|
||||
assert_eq!(
|
||||
|
|
@ -1469,18 +1698,18 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn test_apply_previous_transactions() {
|
||||
let alice = solana_sdk::pubkey::new_rand();
|
||||
let bob = solana_sdk::pubkey::new_rand();
|
||||
let alice = pubkey::new_rand();
|
||||
let bob = pubkey::new_rand();
|
||||
let mut allocations = vec![
|
||||
Allocation {
|
||||
recipient: alice.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: alice,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
Allocation {
|
||||
recipient: bob.to_string(),
|
||||
TypedAllocation {
|
||||
recipient: bob,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
},
|
||||
];
|
||||
let transaction_infos = vec![TransactionInfo {
|
||||
|
|
@ -1493,24 +1722,24 @@ mod tests {
|
|||
|
||||
// Ensure that we applied the transaction to the allocation with
|
||||
// a matching recipient address (to bob, not alice).
|
||||
assert_eq!(allocations[0].recipient, alice.to_string());
|
||||
assert_eq!(allocations[0].recipient, alice);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_has_same_recipient() {
|
||||
let alice_pubkey = solana_sdk::pubkey::new_rand();
|
||||
let bob_pubkey = solana_sdk::pubkey::new_rand();
|
||||
let alice_pubkey = pubkey::new_rand();
|
||||
let bob_pubkey = pubkey::new_rand();
|
||||
let lockup0 = "2021-01-07T00:00:00Z".to_string();
|
||||
let lockup1 = "9999-12-31T23:59:59Z".to_string();
|
||||
let alice_alloc = Allocation {
|
||||
recipient: alice_pubkey.to_string(),
|
||||
let alice_alloc = TypedAllocation {
|
||||
recipient: alice_pubkey,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
};
|
||||
let alice_alloc_lockup0 = Allocation {
|
||||
recipient: alice_pubkey.to_string(),
|
||||
let alice_alloc_lockup0 = TypedAllocation {
|
||||
recipient: alice_pubkey,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: lockup0.clone(),
|
||||
lockup_date: lockup0.parse().ok(),
|
||||
};
|
||||
let alice_info = TransactionInfo {
|
||||
recipient: alice_pubkey,
|
||||
|
|
@ -1550,13 +1779,13 @@ mod tests {
|
|||
#[test]
|
||||
fn test_set_split_stake_lockup() {
|
||||
let lockup_date_str = "2021-01-07T00:00:00Z";
|
||||
let allocation = Allocation {
|
||||
recipient: Pubkey::default().to_string(),
|
||||
let allocation = TypedAllocation {
|
||||
recipient: Pubkey::default(),
|
||||
amount: sol_to_lamports(1.002_282_880),
|
||||
lockup_date: lockup_date_str.to_string(),
|
||||
lockup_date: lockup_date_str.parse().ok(),
|
||||
};
|
||||
let stake_account_address = solana_sdk::pubkey::new_rand();
|
||||
let new_stake_account_address = solana_sdk::pubkey::new_rand();
|
||||
let stake_account_address = pubkey::new_rand();
|
||||
let new_stake_account_address = pubkey::new_rand();
|
||||
let lockup_authority = Keypair::new();
|
||||
let lockup_authority_address = lockup_authority.pubkey();
|
||||
let sender_stake_args = SenderStakeArgs {
|
||||
|
|
@ -1613,12 +1842,12 @@ mod tests {
|
|||
sender_keypair_file: &str,
|
||||
fee_payer: &str,
|
||||
stake_args: Option<StakeArgs>,
|
||||
) -> (Vec<Allocation>, DistributeTokensArgs) {
|
||||
let recipient = solana_sdk::pubkey::new_rand();
|
||||
let allocations = vec![Allocation {
|
||||
recipient: recipient.to_string(),
|
||||
) -> (Vec<TypedAllocation>, DistributeTokensArgs) {
|
||||
let recipient = pubkey::new_rand();
|
||||
let allocations = vec![TypedAllocation {
|
||||
recipient,
|
||||
amount: allocation_amount,
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
}];
|
||||
let args = DistributeTokensArgs {
|
||||
sender_keypair: read_keypair_file(sender_keypair_file).unwrap().into(),
|
||||
|
|
@ -1890,10 +2119,10 @@ mod tests {
|
|||
|
||||
// Underfunded stake-account
|
||||
let expensive_allocation_amount = 5000.0;
|
||||
let expensive_allocations = vec![Allocation {
|
||||
recipient: solana_sdk::pubkey::new_rand().to_string(),
|
||||
let expensive_allocations = vec![TypedAllocation {
|
||||
recipient: pubkey::new_rand(),
|
||||
amount: sol_to_lamports(expensive_allocation_amount),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
}];
|
||||
let err_result = check_payer_balances(
|
||||
&[one_signer_message(&client)],
|
||||
|
|
@ -2108,10 +2337,10 @@ mod tests {
|
|||
spl_token_args: None,
|
||||
transfer_amount: None,
|
||||
};
|
||||
let allocation = Allocation {
|
||||
recipient: recipient.to_string(),
|
||||
let allocation = TypedAllocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
};
|
||||
|
||||
let mut messages: Vec<Message> = vec![];
|
||||
|
|
@ -2230,10 +2459,10 @@ mod tests {
|
|||
spl_token_args: None,
|
||||
transfer_amount: None,
|
||||
};
|
||||
let allocation = Allocation {
|
||||
recipient: recipient.to_string(),
|
||||
let allocation = TypedAllocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
};
|
||||
let message = transaction.message.clone();
|
||||
|
||||
|
|
@ -2329,10 +2558,10 @@ mod tests {
|
|||
.to_string();
|
||||
let mut db = db::open_db(&db_file, false).unwrap();
|
||||
let recipient = Pubkey::new_unique();
|
||||
let allocation = Allocation {
|
||||
recipient: recipient.to_string(),
|
||||
let allocation = TypedAllocation {
|
||||
recipient,
|
||||
amount: sol_to_lamports(1.0),
|
||||
lockup_date: "".to_string(),
|
||||
lockup_date: None,
|
||||
};
|
||||
// This is just dummy data; Args will not affect messages
|
||||
let args = DistributeTokensArgs {
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use {
|
||||
crate::{
|
||||
args::{DistributeTokensArgs, SplTokenArgs},
|
||||
commands::{get_fee_estimate_for_messages, Allocation, Error, FundingSource},
|
||||
commands::{get_fee_estimate_for_messages, Error, FundingSource, TypedAllocation},
|
||||
},
|
||||
console::style,
|
||||
solana_account_decoder::parse_token::{real_number_string, real_number_string_trimmed},
|
||||
|
|
@ -37,7 +37,7 @@ pub fn update_decimals(client: &RpcClient, args: &mut Option<SplTokenArgs>) -> R
|
|||
}
|
||||
|
||||
pub(crate) fn build_spl_token_instructions(
|
||||
allocation: &Allocation,
|
||||
allocation: &TypedAllocation,
|
||||
args: &DistributeTokensArgs,
|
||||
do_create_associated_token_account: bool,
|
||||
) -> Vec<Instruction> {
|
||||
|
|
@ -45,7 +45,7 @@ pub(crate) fn build_spl_token_instructions(
|
|||
.spl_token_args
|
||||
.as_ref()
|
||||
.expect("spl_token_args must be some");
|
||||
let wallet_address = allocation.recipient.parse().unwrap();
|
||||
let wallet_address = allocation.recipient;
|
||||
let associated_token_address =
|
||||
get_associated_token_address(&wallet_address, &spl_token_args.mint);
|
||||
let mut instructions = vec![];
|
||||
|
|
@ -75,7 +75,7 @@ pub(crate) fn build_spl_token_instructions(
|
|||
|
||||
pub(crate) fn check_spl_token_balances(
|
||||
messages: &[Message],
|
||||
allocations: &[Allocation],
|
||||
allocations: &[TypedAllocation],
|
||||
client: &RpcClient,
|
||||
args: &DistributeTokensArgs,
|
||||
created_accounts: u64,
|
||||
|
|
@ -112,10 +112,10 @@ pub(crate) fn check_spl_token_balances(
|
|||
|
||||
pub(crate) fn print_token_balances(
|
||||
client: &RpcClient,
|
||||
allocation: &Allocation,
|
||||
allocation: &TypedAllocation,
|
||||
spl_token_args: &SplTokenArgs,
|
||||
) -> Result<(), Error> {
|
||||
let address = allocation.recipient.parse().unwrap();
|
||||
let address = allocation.recipient;
|
||||
let expected = allocation.amount;
|
||||
let associated_token_address = get_associated_token_address(&address, &spl_token_args.mint);
|
||||
let recipient_account = client
|
||||
|
|
|
|||
Loading…
Reference in New Issue