gaspersic/solana
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Use u64 behind the scenes for solana-tokens (bp #13815) (#13818)

Browse Source 
* Use u64 behind the scenes for solana-tokens (#13815)

* Use u64 behind the scenes for Allocations

* Fixup readme

* Clippy and remove errant comments

(cherry picked from commit 5e2d38227f)

# Conflicts:
#	tokens/src/commands.rs

* Fix conflicts

Co-authored-by: Tyera Eulberg <teulberg@gmail.com>
Co-authored-by: Tyera Eulberg <tyera@solana.com>
This commit is contained in:
mergify[bot]
2020-11-26 03:05:36 +00:00
committed by GitHub
parent a09ee672a6
commit 64c2e759ab
7 changed files with 249 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
tokens/README.md
View File

@ -135,15 +135,16 @@ The Associated Token Account will be created, and funded by the fee_payer, if it
does not already exist. does not already exist.
Send SPL tokens to the recipients in `<RECIPIENTS_CSV>`. Send SPL tokens to the recipients in `<RECIPIENTS_CSV>`.
*NOTE:* the CSV expects SPL-token amounts in raw format (no decimals)
Example recipients.csv: Example recipients.csv:
```text ```text
recipient,amount recipient,amount
CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,75.4 CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,75400
C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,10 C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,10000
7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,42.1 7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,42100
7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,20 7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,20000
``` ```
You can check the status of the recipients before beginning a distribution. You You can check the status of the recipients before beginning a distribution. You
@ -158,9 +159,9 @@ Example output:
```text ```text
Token: JDte736XZ1jGUtfAS32DLpBUWBR7WGSHy1hSZ36VRQ5V Token: JDte736XZ1jGUtfAS32DLpBUWBR7WGSHy1hSZ36VRQ5V
Recipient Expected Balance Actual Balance Difference Recipient Expected Balance Actual Balance Difference
CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT 75.40 0.00 -75.40 CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT 75.400 0.000 -75.400
C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s 10.000 Associated token account not yet created C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s 10.000 Associated token account not yet created
7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr 42.10 0.00 -42.10 7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr 42.100 0.000 -42.100
7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1 20.000 Associated token account not yet created 7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1 20.000 Associated token account not yet created
``` ```
@ -195,10 +196,10 @@ Example updated recipients.csv:
```text ```text
recipient,amount recipient,amount
CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,100 CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,100000
C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,100 C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,100000
7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,100 7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,100000
7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,100 7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,100000
``` ```
Using dry-run: Using dry-run:

5
tokens/src/arg_parser.rs
View File

@ -11,6 +11,7 @@ use solana_clap_utils::{
}; };
use solana_cli_config::CONFIG_FILE; use solana_cli_config::CONFIG_FILE;
use solana_remote_wallet::remote_wallet::maybe_wallet_manager; use solana_remote_wallet::remote_wallet::maybe_wallet_manager;
use solana_sdk::native_token::sol_to_lamports;
use std::error::Error; use std::error::Error;
use std::ffi::OsString; use std::ffi::OsString;
use std::process::exit; use std::process::exit;
@ -360,7 +361,7 @@ fn parse_distribute_tokens_args(
fee_payer, fee_payer,
stake_args: None, stake_args: None,
spl_token_args: None, spl_token_args: None,
transfer_amount: value_of(matches, "transfer_amount"), transfer_amount: value_of(matches, "transfer_amount").map(sol_to_lamports),
}) })
} }
@ -423,7 +424,7 @@ fn parse_distribute_stake_args(
let stake_args = StakeArgs { let stake_args = StakeArgs {
stake_account_address, stake_account_address,
unlocked_sol: value_t_or_exit!(matches, "unlocked_sol", f64), unlocked_sol: sol_to_lamports(value_t_or_exit!(matches, "unlocked_sol", f64)),
stake_authority, stake_authority,
withdraw_authority, withdraw_authority,
lockup_authority, lockup_authority,

4
tokens/src/args.rs
View File

@ -9,11 +9,11 @@ pub struct DistributeTokensArgs {
pub fee_payer: Box<dyn Signer>, pub fee_payer: Box<dyn Signer>,
pub stake_args: Option<StakeArgs>, pub stake_args: Option<StakeArgs>,
pub spl_token_args: Option<SplTokenArgs>, pub spl_token_args: Option<SplTokenArgs>,
pub transfer_amount: Option<f64>, pub transfer_amount: Option<u64>,
} }
pub struct StakeArgs { pub struct StakeArgs {
pub unlocked_sol: f64, pub unlocked_sol: u64,
pub stake_account_address: Pubkey, pub stake_account_address: Pubkey,
pub stake_authority: Box<dyn Signer>, pub stake_authority: Box<dyn Signer>,
pub withdraw_authority: Box<dyn Signer>, pub withdraw_authority: Box<dyn Signer>,

251
tokens/src/commands.rs
View File

@ -11,7 +11,9 @@ use indexmap::IndexMap;
use indicatif::{ProgressBar, ProgressStyle}; use indicatif::{ProgressBar, ProgressStyle};
use pickledb::PickleDb; use pickledb::PickleDb;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use solana_account_decoder::parse_token::{pubkey_from_spl_token_v2_0, spl_token_v2_0_pubkey}; use solana_account_decoder::parse_token::{
pubkey_from_spl_token_v2_0, spl_token_v2_0_pubkey, token_amount_to_ui_amount,
};
use solana_banks_client::{BanksClient, BanksClientExt}; use solana_banks_client::{BanksClient, BanksClientExt};
use solana_sdk::{ use solana_sdk::{
commitment_config::CommitmentLevel, commitment_config::CommitmentLevel,
@ -39,7 +41,7 @@ use tokio::time::sleep;
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct Allocation { pub struct Allocation {
pub recipient: String, pub recipient: String,
pub amount: f64, pub amount: u64,
pub lockup_date: String, pub lockup_date: String,
} }
@ -102,7 +104,7 @@ fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
.entry(&allocation.recipient) .entry(&allocation.recipient)
.or_insert(Allocation { .or_insert(Allocation {
recipient: allocation.recipient.clone(), recipient: allocation.recipient.clone(),
amount: 0.0, amount: 0,
lockup_date: "".to_string(), lockup_date: "".to_string(),
}) })
.amount += allocation.amount; .amount += allocation.amount;
@ -131,11 +133,11 @@ fn apply_previous_transactions(
break; break;
} else { } else {
amount -= allocation.amount; amount -= allocation.amount;
allocation.amount = 0.0; allocation.amount = 0;
} }
} }
} }
allocations.retain(|x| x.amount > f64::EPSILON); allocations.retain(|x| x.amount > 0);
} }
async fn transfer<S: Signer>( async fn transfer<S: Signer>(
@ -165,7 +167,7 @@ fn distribution_instructions(
if args.stake_args.is_none() && args.spl_token_args.is_none() { if args.stake_args.is_none() && args.spl_token_args.is_none() {
let from = args.sender_keypair.pubkey(); let from = args.sender_keypair.pubkey();
let to = allocation.recipient.parse().unwrap(); let to = allocation.recipient.parse().unwrap();
let lamports = sol_to_lamports(allocation.amount); let lamports = allocation.amount;
let instruction = system_instruction::transfer(&from, &to, lamports); let instruction = system_instruction::transfer(&from, &to, lamports);
return vec![instruction]; return vec![instruction];
} }
@ -183,7 +185,7 @@ fn distribution_instructions(
let mut instructions = stake_instruction::split( let mut instructions = stake_instruction::split(
&stake_args.stake_account_address, &stake_args.stake_account_address,
&stake_authority, &stake_authority,
sol_to_lamports(allocation.amount - unlocked_sol), allocation.amount - unlocked_sol,
&new_stake_account_address, &new_stake_account_address,
); );
@ -227,7 +229,7 @@ fn distribution_instructions(
instructions.push(system_instruction::transfer( instructions.push(system_instruction::transfer(
&sender_pubkey, &sender_pubkey,
&recipient, &recipient,
sol_to_lamports(unlocked_sol), unlocked_sol,
)); ));
instructions instructions
@ -251,7 +253,7 @@ async fn distribute_allocations(
Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap()) Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap())
}; };
let (decimals, do_create_associated_token_account) = let (display_amount, decimals, do_create_associated_token_account) =
if let Some(spl_token_args) = &args.spl_token_args { if let Some(spl_token_args) = &args.spl_token_args {
let wallet_address = allocation.recipient.parse().unwrap(); let wallet_address = allocation.recipient.parse().unwrap();
let associated_token_address = get_associated_token_address( let associated_token_address = get_associated_token_address(
@ -266,15 +268,16 @@ async fn distribute_allocations(
created_accounts += 1; created_accounts += 1;
} }
( (
token_amount_to_ui_amount(allocation.amount, spl_token_args.decimals).ui_amount,
spl_token_args.decimals as usize, spl_token_args.decimals as usize,
do_create_associated_token_account, do_create_associated_token_account,
) )
} else { } else {
(9, false) (lamports_to_sol(allocation.amount), 9, false)
}; };
println!( println!(
"{:<44} {:>24.2$}", "{:<44} {:>24.2$}",
allocation.recipient, allocation.amount, decimals allocation.recipient, display_amount, decimals
); );
let instructions = distribution_instructions( let instructions = distribution_instructions(
allocation, allocation,
@ -352,8 +355,9 @@ async fn distribute_allocations(
fn read_allocations( fn read_allocations(
input_csv: &str, input_csv: &str,
transfer_amount: Option<f64>, transfer_amount: Option<u64>,
require_lockup_heading: bool, require_lockup_heading: bool,
raw_amount: bool,
) -> io::Result<Vec<Allocation>> { ) -> io::Result<Vec<Allocation>> {
let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?; let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
let allocations = if let Some(amount) = transfer_amount { let allocations = if let Some(amount) = transfer_amount {
@ -370,7 +374,31 @@ fn read_allocations(
}) })
.collect() .collect()
} else if require_lockup_heading { } else if require_lockup_heading {
rdr.deserialize().map(|entry| entry.unwrap()).collect() 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()
} 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(),
})
.collect()
} else { } else {
let recipients: Vec<(String, f64)> = rdr let recipients: Vec<(String, f64)> = rdr
.deserialize() .deserialize()
@ -380,7 +408,7 @@ fn read_allocations(
.into_iter() .into_iter()
.map(|(recipient, amount)| Allocation { .map(|(recipient, amount)| Allocation {
recipient, recipient,
amount, amount: sol_to_lamports(amount),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}) })
.collect() .collect()
@ -405,10 +433,15 @@ pub async fn process_allocations(
&args.input_csv, &args.input_csv,
args.transfer_amount, args.transfer_amount,
require_lockup_heading, require_lockup_heading,
args.spl_token_args.is_some(),
)?; )?;
let is_sol = args.spl_token_args.is_none();
let starting_total_tokens = Token::from(allocations.iter().map(|x| x.amount).sum(), is_sol); let starting_total_tokens = allocations.iter().map(|x| x.amount).sum();
let starting_total_tokens = if let Some(spl_token_args) = &args.spl_token_args {
Token::spl_token(starting_total_tokens, spl_token_args.decimals)
} else {
Token::sol(starting_total_tokens)
};
println!( println!(
"{} {}", "{} {}",
style("Total in input_csv:").bold(), style("Total in input_csv:").bold(),
@ -428,8 +461,20 @@ pub async fn process_allocations(
return Ok(confirmations); return Ok(confirmations);
} }
let distributed_tokens = Token::from(transaction_infos.iter().map(|x| x.amount).sum(), is_sol); let distributed_tokens = transaction_infos.iter().map(|x| x.amount).sum();
let undistributed_tokens = Token::from(allocations.iter().map(|x| x.amount).sum(), is_sol); let undistributed_tokens = allocations.iter().map(|x| x.amount).sum();
let (distributed_tokens, undistributed_tokens) =
if let Some(spl_token_args) = &args.spl_token_args {
(
Token::spl_token(distributed_tokens, spl_token_args.decimals),
Token::spl_token(undistributed_tokens, spl_token_args.decimals),
)
} else {
(
Token::sol(distributed_tokens),
Token::sol(undistributed_tokens),
)
};
println!("{} {}", style("Distributed:").bold(), distributed_tokens,); println!("{} {}", style("Distributed:").bold(), distributed_tokens,);
println!( println!(
"{} {}", "{} {}",
@ -546,7 +591,7 @@ async fn check_payer_balances(
client: &mut BanksClient, client: &mut BanksClient,
args: &DistributeTokensArgs, args: &DistributeTokensArgs,
) -> Result<(), Error> { ) -> Result<(), Error> {
let mut undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum(); let mut undistributed_tokens: u64 = allocations.iter().map(|x| x.amount).sum();
let (fee_calculator, _blockhash, _last_valid_slot) = client.get_fees().await?; let (fee_calculator, _blockhash, _last_valid_slot) = client.get_fees().await?;
let fees = fee_calculator let fees = fee_calculator
@ -555,26 +600,22 @@ async fn check_payer_balances(
.unwrap(); .unwrap();
let (distribution_source, unlocked_sol_source) = if let Some(stake_args) = &args.stake_args { let (distribution_source, unlocked_sol_source) = if let Some(stake_args) = &args.stake_args {
let total_unlocked_sol = allocations.len() as f64 * stake_args.unlocked_sol; let total_unlocked_sol = allocations.len() as u64 * stake_args.unlocked_sol;
undistributed_tokens -= total_unlocked_sol; undistributed_tokens -= total_unlocked_sol;
( (
stake_args.stake_account_address, stake_args.stake_account_address,
Some(( Some((args.sender_keypair.pubkey(), total_unlocked_sol)),
args.sender_keypair.pubkey(),
sol_to_lamports(total_unlocked_sol),
)),
) )
} else { } else {
(args.sender_keypair.pubkey(), None) (args.sender_keypair.pubkey(), None)
}; };
let allocation_lamports = sol_to_lamports(undistributed_tokens);
if let Some((unlocked_sol_source, total_unlocked_sol)) = unlocked_sol_source { if let Some((unlocked_sol_source, total_unlocked_sol)) = unlocked_sol_source {
let staker_balance = client.get_balance(distribution_source).await?; let staker_balance = client.get_balance(distribution_source).await?;
if staker_balance < allocation_lamports { if staker_balance < undistributed_tokens {
return Err(Error::InsufficientFunds( return Err(Error::InsufficientFunds(
vec![FundingSource::StakeAccount].into(), vec![FundingSource::StakeAccount].into(),
lamports_to_sol(allocation_lamports), lamports_to_sol(undistributed_tokens),
)); ));
} }
if args.fee_payer.pubkey() == unlocked_sol_source { if args.fee_payer.pubkey() == unlocked_sol_source {
@ -603,10 +644,10 @@ async fn check_payer_balances(
} }
} else if args.fee_payer.pubkey() == distribution_source { } else if args.fee_payer.pubkey() == distribution_source {
let balance = client.get_balance(args.fee_payer.pubkey()).await?; let balance = client.get_balance(args.fee_payer.pubkey()).await?;
if balance < fees + allocation_lamports { if balance < fees + undistributed_tokens {
return Err(Error::InsufficientFunds( return Err(Error::InsufficientFunds(
vec![FundingSource::SystemAccount, FundingSource::FeePayer].into(), vec![FundingSource::SystemAccount, FundingSource::FeePayer].into(),
lamports_to_sol(fees + allocation_lamports), lamports_to_sol(fees + undistributed_tokens),
)); ));
} }
} else { } else {
@ -618,10 +659,10 @@ async fn check_payer_balances(
)); ));
} }
let sender_balance = client.get_balance(distribution_source).await?; let sender_balance = client.get_balance(distribution_source).await?;
if sender_balance < allocation_lamports { if sender_balance < undistributed_tokens {
return Err(Error::InsufficientFunds( return Err(Error::InsufficientFunds(
vec![FundingSource::SystemAccount].into(), vec![FundingSource::SystemAccount].into(),
lamports_to_sol(allocation_lamports), lamports_to_sol(undistributed_tokens),
)); ));
} }
} }
@ -629,7 +670,8 @@ async fn check_payer_balances(
} }
pub async fn process_balances(client: &mut BanksClient, args: &BalancesArgs) -> Result<(), Error> { pub async fn process_balances(client: &mut BanksClient, args: &BalancesArgs) -> Result<(), Error> {
let allocations: Vec<Allocation> = read_allocations(&args.input_csv, None, false)?; let allocations: Vec<Allocation> =
read_allocations(&args.input_csv, None, false, args.spl_token_args.is_some())?;
let allocations = merge_allocations(&allocations); let allocations = merge_allocations(&allocations);
let token = if let Some(spl_token_args) = &args.spl_token_args { let token = if let Some(spl_token_args) = &args.spl_token_args {
@ -653,7 +695,7 @@ pub async fn process_balances(client: &mut BanksClient, args: &BalancesArgs) ->
print_token_balances(client, allocation, spl_token_args).await?; print_token_balances(client, allocation, spl_token_args).await?;
} else { } else {
let address: Pubkey = allocation.recipient.parse().unwrap(); let address: Pubkey = allocation.recipient.parse().unwrap();
let expected = lamports_to_sol(sol_to_lamports(allocation.amount)); let expected = lamports_to_sol(allocation.amount);
let actual = lamports_to_sol(client.get_balance(address).await.unwrap()); let actual = lamports_to_sol(client.get_balance(address).await.unwrap());
println!( println!(
"{:<44} {:>24.9} {:>24.9} {:>24.9}", "{:<44} {:>24.9} {:>24.9} {:>24.9}",
@ -680,7 +722,7 @@ use tempfile::{tempdir, NamedTempFile};
pub async fn test_process_distribute_tokens_with_client( pub async fn test_process_distribute_tokens_with_client(
client: &mut BanksClient, client: &mut BanksClient,
sender_keypair: Keypair, sender_keypair: Keypair,
transfer_amount: Option<f64>, transfer_amount: Option<u64>,
) { ) {
let fee_payer = Keypair::new(); let fee_payer = Keypair::new();
let transaction = transfer( let transaction = transfer(
@ -703,20 +745,21 @@ pub async fn test_process_distribute_tokens_with_client(
sol_to_lamports(1.0), sol_to_lamports(1.0),
); );
let alice_pubkey = solana_sdk::pubkey::new_rand(); let expected_amount = if let Some(amount) = transfer_amount {
let allocation = Allocation {
recipient: alice_pubkey.to_string(),
amount: if let Some(amount) = transfer_amount {
amount amount
} else { } else {
1000.0 sol_to_lamports(1000.0)
},
lockup_date: "".to_string(),
}; };
let alice_pubkey = solana_sdk::pubkey::new_rand();
let allocations_file = NamedTempFile::new().unwrap(); let allocations_file = NamedTempFile::new().unwrap();
let input_csv = allocations_file.path().to_str().unwrap().to_string(); let input_csv = allocations_file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file); let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file);
wtr.serialize(&allocation).unwrap(); wtr.write_record(&["recipient", "amount"]).unwrap();
wtr.write_record(&[
alice_pubkey.to_string(),
lamports_to_sol(expected_amount).to_string(),
])
.unwrap();
wtr.flush().unwrap(); wtr.flush().unwrap();
let dir = tempdir().unwrap(); let dir = tempdir().unwrap();
@ -748,11 +791,7 @@ pub async fn test_process_distribute_tokens_with_client(
db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap()); db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
assert_eq!(transaction_infos.len(), 1); assert_eq!(transaction_infos.len(), 1);
assert_eq!(transaction_infos[0].recipient, alice_pubkey); assert_eq!(transaction_infos[0].recipient, alice_pubkey);
let expected_amount = sol_to_lamports(allocation.amount); assert_eq!(transaction_infos[0].amount, expected_amount);
assert_eq!(
sol_to_lamports(transaction_infos[0].amount),
expected_amount
);
assert_eq!( assert_eq!(
client.get_balance(alice_pubkey).await.unwrap(), client.get_balance(alice_pubkey).await.unwrap(),
@ -767,11 +806,7 @@ pub async fn test_process_distribute_tokens_with_client(
db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap()); db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
assert_eq!(transaction_infos.len(), 1); assert_eq!(transaction_infos.len(), 1);
assert_eq!(transaction_infos[0].recipient, alice_pubkey); assert_eq!(transaction_infos[0].recipient, alice_pubkey);
let expected_amount = sol_to_lamports(allocation.amount); assert_eq!(transaction_infos[0].amount, expected_amount);
assert_eq!(
sol_to_lamports(transaction_infos[0].amount),
expected_amount
);
assert_eq!( assert_eq!(
client.get_balance(alice_pubkey).await.unwrap(), client.get_balance(alice_pubkey).await.unwrap(),
@ -825,16 +860,19 @@ pub async fn test_process_distribute_stake_with_client(
.await .await
.unwrap(); .unwrap();
let expected_amount = sol_to_lamports(1000.0);
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = solana_sdk::pubkey::new_rand();
let allocation = Allocation {
recipient: alice_pubkey.to_string(),
amount: 1000.0,
lockup_date: "".to_string(),
};
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
let input_csv = file.path().to_str().unwrap().to_string(); let input_csv = file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(file); let mut wtr = csv::WriterBuilder::new().from_writer(file);
wtr.serialize(&allocation).unwrap(); wtr.write_record(&["recipient", "amount", "lockup_date"])
.unwrap();
wtr.write_record(&[
alice_pubkey.to_string(),
lamports_to_sol(expected_amount).to_string(),
"".to_string(),
])
.unwrap();
wtr.flush().unwrap(); wtr.flush().unwrap();
let dir = tempdir().unwrap(); let dir = tempdir().unwrap();
@ -853,7 +891,7 @@ pub async fn test_process_distribute_stake_with_client(
stake_authority: Box::new(stake_authority), stake_authority: Box::new(stake_authority),
withdraw_authority: Box::new(withdraw_authority), withdraw_authority: Box::new(withdraw_authority),
lockup_authority: None, lockup_authority: None,
unlocked_sol: 1.0, unlocked_sol: sol_to_lamports(1.0),
}; };
let args = DistributeTokensArgs { let args = DistributeTokensArgs {
fee_payer: Box::new(fee_payer), fee_payer: Box::new(fee_payer),
@ -873,11 +911,7 @@ pub async fn test_process_distribute_stake_with_client(
db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap()); db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
assert_eq!(transaction_infos.len(), 1); assert_eq!(transaction_infos.len(), 1);
assert_eq!(transaction_infos[0].recipient, alice_pubkey); assert_eq!(transaction_infos[0].recipient, alice_pubkey);
let expected_amount = sol_to_lamports(allocation.amount); assert_eq!(transaction_infos[0].amount, expected_amount);
assert_eq!(
sol_to_lamports(transaction_infos[0].amount),
expected_amount
);
assert_eq!( assert_eq!(
client.get_balance(alice_pubkey).await.unwrap(), client.get_balance(alice_pubkey).await.unwrap(),
@ -897,11 +931,7 @@ pub async fn test_process_distribute_stake_with_client(
db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap()); db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
assert_eq!(transaction_infos.len(), 1); assert_eq!(transaction_infos.len(), 1);
assert_eq!(transaction_infos[0].recipient, alice_pubkey); assert_eq!(transaction_infos[0].recipient, alice_pubkey);
let expected_amount = sol_to_lamports(allocation.amount); assert_eq!(transaction_infos[0].amount, expected_amount);
assert_eq!(
sol_to_lamports(transaction_infos[0].amount),
expected_amount
);
assert_eq!( assert_eq!(
client.get_balance(alice_pubkey).await.unwrap(), client.get_balance(alice_pubkey).await.unwrap(),
@ -952,7 +982,7 @@ pub(crate) mod tests {
test_process_distribute_tokens_with_client( test_process_distribute_tokens_with_client(
&mut banks_client, &mut banks_client,
sender_keypair, sender_keypair,
Some(1.5), Some(sol_to_lamports(1.5)),
) )
.await; .await;
}); });
@ -974,7 +1004,7 @@ pub(crate) mod tests {
let alice_pubkey = solana_sdk::pubkey::new_rand(); let alice_pubkey = solana_sdk::pubkey::new_rand();
let allocation = Allocation { let allocation = Allocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey.to_string(),
amount: 42.0, amount: 42,
lockup_date: "".to_string(), lockup_date: "".to_string(),
}; };
let file = NamedTempFile::new().unwrap(); let file = NamedTempFile::new().unwrap();
@ -984,12 +1014,27 @@ pub(crate) mod tests {
wtr.flush().unwrap(); wtr.flush().unwrap();
assert_eq!( assert_eq!(
read_allocations(&input_csv, None, false).unwrap(), read_allocations(&input_csv, None, false, true).unwrap(),
vec![allocation.clone()] vec![allocation]
);
let allocation_sol = Allocation {
recipient: alice_pubkey.to_string(),
amount: sol_to_lamports(42.0),
lockup_date: "".to_string(),
};
assert_eq!(
read_allocations(&input_csv, None, true, true).unwrap(),
vec![allocation_sol.clone()]
); );
assert_eq!( assert_eq!(
read_allocations(&input_csv, None, true).unwrap(), read_allocations(&input_csv, None, false, false).unwrap(),
vec![allocation] vec![allocation_sol.clone()]
);
assert_eq!(
read_allocations(&input_csv, None, true, false).unwrap(),
vec![allocation_sol]
); );
} }
@ -1009,17 +1054,17 @@ pub(crate) mod tests {
let expected_allocations = vec![ let expected_allocations = vec![
Allocation { Allocation {
recipient: pubkey0.to_string(), recipient: pubkey0.to_string(),
amount: 42.0, amount: sol_to_lamports(42.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}, },
Allocation { Allocation {
recipient: pubkey1.to_string(), recipient: pubkey1.to_string(),
amount: 43.0, amount: sol_to_lamports(43.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}, },
]; ];
assert_eq!( assert_eq!(
read_allocations(&input_csv, None, false).unwrap(), read_allocations(&input_csv, None, false, false).unwrap(),
expected_allocations expected_allocations
); );
} }
@ -1041,17 +1086,17 @@ pub(crate) mod tests {
let expected_allocations = vec![ let expected_allocations = vec![
Allocation { Allocation {
recipient: pubkey0.to_string(), recipient: pubkey0.to_string(),
amount: 42.0, amount: sol_to_lamports(42.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}, },
Allocation { Allocation {
recipient: pubkey1.to_string(), recipient: pubkey1.to_string(),
amount: 43.0, amount: sol_to_lamports(43.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}, },
]; ];
assert_eq!( assert_eq!(
read_allocations(&input_csv, None, true).unwrap(), read_allocations(&input_csv, None, true, false).unwrap(),
expected_allocations expected_allocations
); );
} }
@ -1070,7 +1115,7 @@ pub(crate) mod tests {
wtr.serialize(&pubkey2.to_string()).unwrap(); wtr.serialize(&pubkey2.to_string()).unwrap();
wtr.flush().unwrap(); wtr.flush().unwrap();
let amount = 1.5; let amount = sol_to_lamports(1.5);
let expected_allocations = vec![ let expected_allocations = vec![
Allocation { Allocation {
@ -1090,7 +1135,7 @@ pub(crate) mod tests {
}, },
]; ];
assert_eq!( assert_eq!(
read_allocations(&input_csv, Some(amount), false).unwrap(), read_allocations(&input_csv, Some(amount), false, false).unwrap(),
expected_allocations expected_allocations
); );
} }
@ -1102,18 +1147,18 @@ pub(crate) mod tests {
let mut allocations = vec![ let mut allocations = vec![
Allocation { Allocation {
recipient: alice.to_string(), recipient: alice.to_string(),
amount: 1.0, amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}, },
Allocation { Allocation {
recipient: bob.to_string(), recipient: bob.to_string(),
amount: 1.0, amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}, },
]; ];
let transaction_infos = vec![TransactionInfo { let transaction_infos = vec![TransactionInfo {
recipient: bob, recipient: bob,
amount: 1.0, amount: sol_to_lamports(1.0),
..TransactionInfo::default() ..TransactionInfo::default()
}]; }];
apply_previous_transactions(&mut allocations, &transaction_infos); apply_previous_transactions(&mut allocations, &transaction_infos);
@ -1132,12 +1177,12 @@ pub(crate) mod tests {
let lockup1 = "9999-12-31T23:59:59Z".to_string(); let lockup1 = "9999-12-31T23:59:59Z".to_string();
let alice_alloc = Allocation { let alice_alloc = Allocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey.to_string(),
amount: 1.0, amount: sol_to_lamports(1.0),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}; };
let alice_alloc_lockup0 = Allocation { let alice_alloc_lockup0 = Allocation {
recipient: alice_pubkey.to_string(), recipient: alice_pubkey.to_string(),
amount: 1.0, amount: sol_to_lamports(1.0),
lockup_date: lockup0.clone(), lockup_date: lockup0.clone(),
}; };
let alice_info = TransactionInfo { let alice_info = TransactionInfo {
@ -1180,7 +1225,7 @@ pub(crate) mod tests {
let lockup_date_str = "2021-01-07T00:00:00Z"; let lockup_date_str = "2021-01-07T00:00:00Z";
let allocation = Allocation { let allocation = Allocation {
recipient: Pubkey::default().to_string(), recipient: Pubkey::default().to_string(),
amount: 1.0, amount: sol_to_lamports(1.0),
lockup_date: lockup_date_str.to_string(), lockup_date: lockup_date_str.to_string(),
}; };
let stake_account_address = solana_sdk::pubkey::new_rand(); let stake_account_address = solana_sdk::pubkey::new_rand();
@ -1191,7 +1236,7 @@ pub(crate) mod tests {
stake_authority: Box::new(Keypair::new()), stake_authority: Box::new(Keypair::new()),
withdraw_authority: Box::new(Keypair::new()), withdraw_authority: Box::new(Keypair::new()),
lockup_authority: Some(Box::new(lockup_authority)), lockup_authority: Some(Box::new(lockup_authority)),
unlocked_sol: 1.0, unlocked_sol: sol_to_lamports(1.0),
}; };
let args = DistributeTokensArgs { let args = DistributeTokensArgs {
fee_payer: Box::new(Keypair::new()), fee_payer: Box::new(Keypair::new()),
@ -1231,7 +1276,7 @@ pub(crate) mod tests {
} }
fn initialize_check_payer_balances_inputs( fn initialize_check_payer_balances_inputs(
allocation_amount: f64, allocation_amount: u64,
sender_keypair_file: &str, sender_keypair_file: &str,
fee_payer: &str, fee_payer: &str,
stake_args: Option<StakeArgs>, stake_args: Option<StakeArgs>,
@ -1275,7 +1320,7 @@ pub(crate) mod tests {
// Fully funded payer // Fully funded payer
let (allocations, mut args) = initialize_check_payer_balances_inputs( let (allocations, mut args) = initialize_check_payer_balances_inputs(
allocation_amount, sol_to_lamports(allocation_amount),
&sender_keypair_file, &sender_keypair_file,
&sender_keypair_file, &sender_keypair_file,
None, None,
@ -1387,7 +1432,7 @@ pub(crate) mod tests {
// Fully funded payers // Fully funded payers
let (allocations, mut args) = initialize_check_payer_balances_inputs( let (allocations, mut args) = initialize_check_payer_balances_inputs(
allocation_amount, sol_to_lamports(allocation_amount),
&funded_payer_keypair_file, &funded_payer_keypair_file,
&sender_keypair_file, &sender_keypair_file,
None, None,
@ -1435,8 +1480,8 @@ pub(crate) mod tests {
} }
async fn initialize_stake_account( async fn initialize_stake_account(
stake_account_amount: f64, stake_account_amount: u64,
unlocked_sol: f64, unlocked_sol: u64,
sender_keypair: &Keypair, sender_keypair: &Keypair,
banks_client: &mut BanksClient, banks_client: &mut BanksClient,
) -> StakeArgs { ) -> StakeArgs {
@ -1455,7 +1500,7 @@ pub(crate) mod tests {
&stake_account_address, &stake_account_address,
&authorized, &authorized,
&lockup, &lockup,
sol_to_lamports(stake_account_amount), stake_account_amount,
); );
let message = Message::new(&instructions, Some(&sender_keypair.pubkey())); let message = Message::new(&instructions, Some(&sender_keypair.pubkey()));
let signers = [sender_keypair, &stake_account_keypair]; let signers = [sender_keypair, &stake_account_keypair];
@ -1493,8 +1538,8 @@ pub(crate) mod tests {
let allocation_amount = 1000.0; let allocation_amount = 1000.0;
let unlocked_sol = 1.0; let unlocked_sol = 1.0;
let stake_args = initialize_stake_account( let stake_args = initialize_stake_account(
allocation_amount, sol_to_lamports(allocation_amount),
unlocked_sol, sol_to_lamports(unlocked_sol),
&sender_keypair, &sender_keypair,
&mut banks_client, &mut banks_client,
) )
@ -1502,7 +1547,7 @@ pub(crate) mod tests {
// Fully funded payer & stake account // Fully funded payer & stake account
let (allocations, mut args) = initialize_check_payer_balances_inputs( let (allocations, mut args) = initialize_check_payer_balances_inputs(
allocation_amount, sol_to_lamports(allocation_amount),
&sender_keypair_file, &sender_keypair_file,
&sender_keypair_file, &sender_keypair_file,
Some(stake_args), Some(stake_args),
@ -1515,7 +1560,7 @@ pub(crate) mod tests {
let expensive_allocation_amount = 5000.0; let expensive_allocation_amount = 5000.0;
let expensive_allocations = vec![Allocation { let expensive_allocations = vec![Allocation {
recipient: solana_sdk::pubkey::new_rand().to_string(), recipient: solana_sdk::pubkey::new_rand().to_string(),
amount: expensive_allocation_amount, amount: sol_to_lamports(expensive_allocation_amount),
lockup_date: "".to_string(), lockup_date: "".to_string(),
}]; }];
let err_result = let err_result =
@ -1617,8 +1662,8 @@ pub(crate) mod tests {
let allocation_amount = 1000.0; let allocation_amount = 1000.0;
let unlocked_sol = 1.0; let unlocked_sol = 1.0;
let stake_args = initialize_stake_account( let stake_args = initialize_stake_account(
allocation_amount, sol_to_lamports(allocation_amount),
unlocked_sol, sol_to_lamports(unlocked_sol),
&sender_keypair, &sender_keypair,
&mut banks_client, &mut banks_client,
) )
@ -1642,7 +1687,7 @@ pub(crate) mod tests {
// Fully funded payers // Fully funded payers
let (allocations, mut args) = initialize_check_payer_balances_inputs( let (allocations, mut args) = initialize_check_payer_balances_inputs(
allocation_amount, sol_to_lamports(allocation_amount),
&funded_payer_keypair_file, &funded_payer_keypair_file,
&sender_keypair_file, &sender_keypair_file,
Some(stake_args), Some(stake_args),

8
tokens/src/db.rs
View File

@ -8,7 +8,7 @@ use std::{cmp::Ordering, fs, io, path::Path};
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct TransactionInfo { pub struct TransactionInfo {
pub recipient: Pubkey, pub recipient: Pubkey,
pub amount: f64, pub amount: u64,
pub new_stake_account_address: Option<Pubkey>, pub new_stake_account_address: Option<Pubkey>,
pub finalized_date: Option<DateTime<Utc>>, pub finalized_date: Option<DateTime<Utc>>,
pub transaction: Transaction, pub transaction: Transaction,
@ -19,7 +19,7 @@ pub struct TransactionInfo {
#[derive(Serialize, Deserialize, Debug, Default, PartialEq)] #[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
struct SignedTransactionInfo { struct SignedTransactionInfo {
recipient: String, recipient: String,
amount: f64, amount: u64,
#[serde(skip_serializing_if = "String::is_empty", default)] #[serde(skip_serializing_if = "String::is_empty", default)]
new_stake_account_address: String, new_stake_account_address: String,
finalized_date: Option<DateTime<Utc>>, finalized_date: Option<DateTime<Utc>>,
@ -34,7 +34,7 @@ impl Default for TransactionInfo {
}; };
Self { Self {
recipient: Pubkey::default(), recipient: Pubkey::default(),
amount: 0.0, amount: 0,
new_stake_account_address: None, new_stake_account_address: None,
finalized_date: None, finalized_date: None,
transaction, transaction,
@ -104,7 +104,7 @@ pub fn read_transaction_infos(db: &PickleDb) -> Vec<TransactionInfo> {
pub fn set_transaction_info( pub fn set_transaction_info(
db: &mut PickleDb, db: &mut PickleDb,
recipient: &Pubkey, recipient: &Pubkey,
amount: f64, amount: u64,
transaction: &Transaction, transaction: &Transaction,
new_stake_account_address: Option<&Pubkey>, new_stake_account_address: Option<&Pubkey>,
finalized: bool, finalized: bool,

87
tokens/src/spl_token.rs
View File

@ -85,7 +85,7 @@ pub fn build_spl_token_instructions(
&associated_token_address, &associated_token_address,
&spl_token_v2_0_pubkey(&args.sender_keypair.pubkey()), &spl_token_v2_0_pubkey(&args.sender_keypair.pubkey()),
&[], &[],
spl_token_amount(allocation.amount, spl_token_args.decimals), allocation.amount,
spl_token_args.decimals, spl_token_args.decimals,
) )
.unwrap(); .unwrap();
@ -104,8 +104,7 @@ pub async fn check_spl_token_balances(
.spl_token_args .spl_token_args
.as_ref() .as_ref()
.expect("spl_token_args must be some"); .expect("spl_token_args must be some");
let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum(); let allocation_amount: u64 = allocations.iter().map(|x| x.amount).sum();
let allocation_amount = spl_token_amount(undistributed_tokens, spl_token_args.decimals);
let (fee_calculator, _blockhash, _last_valid_slot) = client.get_fees().await?; let (fee_calculator, _blockhash, _last_valid_slot) = client.get_fees().await?;
let fees = fee_calculator let fees = fee_calculator
@ -152,18 +151,21 @@ pub async fn print_token_balances(
.get_account(pubkey_from_spl_token_v2_0(&associated_token_address)) .get_account(pubkey_from_spl_token_v2_0(&associated_token_address))
.await? .await?
.unwrap_or_default(); .unwrap_or_default();
let (actual, difference) = let (actual, difference) = if let Ok(recipient_token) =
if let Ok(recipient_token) = SplTokenAccount::unpack(&recipient_account.data) { SplTokenAccount::unpack(&recipient_account.data)
let actual = token_amount_to_ui_amount(recipient_token.amount, spl_token_args.decimals) {
.ui_amount; let actual_ui_amount =
token_amount_to_ui_amount(recipient_token.amount, spl_token_args.decimals).ui_amount;
let expected_ui_amount =
token_amount_to_ui_amount(expected, spl_token_args.decimals).ui_amount;
( (
style(format!( style(format!(
"{:>24.1$}", "{:>24.1$}",
actual, spl_token_args.decimals as usize actual_ui_amount, spl_token_args.decimals as usize
)), )),
format!( format!(
"{:>24.1$}", "{:>24.1$}",
actual - expected, actual_ui_amount - expected_ui_amount,
spl_token_args.decimals as usize spl_token_args.decimals as usize
), ),
) )
@ -175,7 +177,11 @@ pub async fn print_token_balances(
}; };
println!( println!(
"{:<44} {:>24.4$} {:>24} {:>24}", "{:<44} {:>24.4$} {:>24} {:>24}",
allocation.recipient, expected, actual, difference, spl_token_args.decimals as usize allocation.recipient,
token_amount_to_ui_amount(expected, spl_token_args.decimals).ui_amount,
actual,
difference,
spl_token_args.decimals as usize
); );
Ok(()) Ok(())
} }
@ -186,7 +192,6 @@ mod tests {
use crate::{ use crate::{
commands::{process_allocations, tests::tmp_file_path, Allocation}, commands::{process_allocations, tests::tmp_file_path, Allocation},
db::{self, check_output_file}, db::{self, check_output_file},
spl_token::spl_token_amount,
}; };
use solana_account_decoder::parse_token::{spl_token_id_v2_0, spl_token_v2_0_pubkey}; use solana_account_decoder::parse_token::{spl_token_id_v2_0, spl_token_v2_0_pubkey};
use solana_program_test::*; use solana_program_test::*;
@ -310,7 +315,7 @@ mod tests {
async fn test_process_distribute_spl_tokens_with_client( async fn test_process_distribute_spl_tokens_with_client(
banks_client: &mut BanksClient, banks_client: &mut BanksClient,
fee_payer: Keypair, fee_payer: Keypair,
transfer_amount: Option<f64>, transfer_amount: Option<u64>,
recent_blockhash: Hash, recent_blockhash: Hash,
) { ) {
// Initialize Token Mint // Initialize Token Mint
@ -357,23 +362,16 @@ mod tests {
let allocation_amount = if let Some(amount) = transfer_amount { let allocation_amount = if let Some(amount) = transfer_amount {
amount amount
} else { } else {
1000.0 100_000
}; };
let allocations_file = NamedTempFile::new().unwrap(); let allocations_file = NamedTempFile::new().unwrap();
let input_csv = allocations_file.path().to_str().unwrap().to_string(); let input_csv = allocations_file.path().to_str().unwrap().to_string();
let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file); let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file);
let allocation = Allocation { wtr.write_record(&["recipient", "amount"]).unwrap();
recipient: wallet_address_0.to_string(), wtr.write_record(&[wallet_address_0.to_string(), allocation_amount.to_string()])
amount: allocation_amount, .unwrap();
lockup_date: "".to_string(), wtr.write_record(&[wallet_address_1.to_string(), allocation_amount.to_string()])
}; .unwrap();
wtr.serialize(&allocation).unwrap();
let allocation = Allocation {
recipient: wallet_address_1.to_string(),
amount: allocation_amount,
lockup_date: "".to_string(),
};
wtr.serialize(&allocation).unwrap();
wtr.flush().unwrap(); wtr.flush().unwrap();
let dir = tempdir().unwrap(); let dir = tempdir().unwrap();
@ -421,15 +419,8 @@ mod tests {
assert!(transaction_infos assert!(transaction_infos
.iter() .iter()
.any(|info| info.recipient == pubkey_from_spl_token_v2_0(&wallet_address_1))); .any(|info| info.recipient == pubkey_from_spl_token_v2_0(&wallet_address_1)));
let expected_amount = spl_token_amount(allocation.amount, decimals); assert_eq!(transaction_infos[0].amount, allocation_amount);
assert_eq!( assert_eq!(transaction_infos[1].amount, allocation_amount);
spl_token_amount(transaction_infos[0].amount, decimals),
expected_amount
);
assert_eq!(
spl_token_amount(transaction_infos[1].amount, decimals),
expected_amount
);
let recipient_account_0 = banks_client let recipient_account_0 = banks_client
.get_account(pubkey_from_spl_token_v2_0(&associated_token_address_0)) .get_account(pubkey_from_spl_token_v2_0(&associated_token_address_0))
@ -440,7 +431,7 @@ mod tests {
SplTokenAccount::unpack(&recipient_account_0.data) SplTokenAccount::unpack(&recipient_account_0.data)
.unwrap() .unwrap()
.amount, .amount,
expected_amount, allocation_amount,
); );
let recipient_account_1 = banks_client let recipient_account_1 = banks_client
.get_account(pubkey_from_spl_token_v2_0(&associated_token_address_1)) .get_account(pubkey_from_spl_token_v2_0(&associated_token_address_1))
@ -451,7 +442,7 @@ mod tests {
SplTokenAccount::unpack(&recipient_account_1.data) SplTokenAccount::unpack(&recipient_account_1.data)
.unwrap() .unwrap()
.amount, .amount,
expected_amount, allocation_amount,
); );
check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap()); check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
@ -467,15 +458,8 @@ mod tests {
assert!(transaction_infos assert!(transaction_infos
.iter() .iter()
.any(|info| info.recipient == pubkey_from_spl_token_v2_0(&wallet_address_1))); .any(|info| info.recipient == pubkey_from_spl_token_v2_0(&wallet_address_1)));
let expected_amount = spl_token_amount(allocation.amount, decimals); assert_eq!(transaction_infos[0].amount, allocation_amount);
assert_eq!( assert_eq!(transaction_infos[1].amount, allocation_amount);
spl_token_amount(transaction_infos[0].amount, decimals),
expected_amount
);
assert_eq!(
spl_token_amount(transaction_infos[1].amount, decimals),
expected_amount
);
let recipient_account_0 = banks_client let recipient_account_0 = banks_client
.get_account(pubkey_from_spl_token_v2_0(&associated_token_address_0)) .get_account(pubkey_from_spl_token_v2_0(&associated_token_address_0))
@ -486,7 +470,7 @@ mod tests {
SplTokenAccount::unpack(&recipient_account_0.data) SplTokenAccount::unpack(&recipient_account_0.data)
.unwrap() .unwrap()
.amount, .amount,
expected_amount, allocation_amount,
); );
let recipient_account_1 = banks_client let recipient_account_1 = banks_client
.get_account(pubkey_from_spl_token_v2_0(&associated_token_address_1)) .get_account(pubkey_from_spl_token_v2_0(&associated_token_address_1))
@ -497,7 +481,7 @@ mod tests {
SplTokenAccount::unpack(&recipient_account_1.data) SplTokenAccount::unpack(&recipient_account_1.data)
.unwrap() .unwrap()
.amount, .amount,
expected_amount, allocation_amount,
); );
check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap()); check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
@ -521,7 +505,7 @@ mod tests {
test_process_distribute_spl_tokens_with_client( test_process_distribute_spl_tokens_with_client(
&mut banks_client, &mut banks_client,
payer, payer,
Some(105.5), Some(10550),
recent_blockhash, recent_blockhash,
) )
.await; .await;
@ -563,7 +547,7 @@ mod tests {
let unfunded_fee_payer = Keypair::new(); let unfunded_fee_payer = Keypair::new();
let allocation_amount = 42.0; let allocation_amount = 4200;
let allocations = vec![Allocation { let allocations = vec![Allocation {
recipient: Pubkey::new_unique().to_string(), recipient: Pubkey::new_unique().to_string(),
amount: allocation_amount, amount: allocation_amount,
@ -628,7 +612,10 @@ mod tests {
.unwrap_err(); .unwrap_err();
if let Error::InsufficientFunds(sources, amount) = err_result { if let Error::InsufficientFunds(sources, amount) = err_result {
assert_eq!(sources, vec![FundingSource::SplTokenAccount].into()); assert_eq!(sources, vec![FundingSource::SplTokenAccount].into());
assert!((amount - allocation_amount).abs() < f64::EPSILON); assert!(
(amount - token_amount_to_ui_amount(allocation_amount, decimals).ui_amount).abs()
< f64::EPSILON
);
} else { } else {
panic!("check_spl_token_balances should have errored"); panic!("check_spl_token_balances should have errored");
} }

37
tokens/src/token_display.rs
View File

@ -1,3 +1,5 @@
use solana_account_decoder::parse_token::token_amount_to_ui_amount;
use solana_sdk::native_token::lamports_to_sol;
use std::{ use std::{
fmt::{Debug, Display, Formatter, Result}, fmt::{Debug, Display, Formatter, Result},
ops::Add, ops::Add,
@ -12,25 +14,39 @@ pub enum TokenType {
} }
pub struct Token { pub struct Token {
amount: f64, amount: u64,
decimals: u8,
token_type: TokenType, token_type: TokenType,
} }
impl Token { impl Token {
fn write_with_symbol(&self, f: &mut Formatter) -> Result { fn write_with_symbol(&self, f: &mut Formatter) -> Result {
match &self.token_type { match &self.token_type {
TokenType::Sol => write!(f, "{}{}", SOL_SYMBOL, self.amount,), TokenType::Sol => {
TokenType::SplToken => write!(f, "{} tokens", self.amount,), let amount = lamports_to_sol(self.amount);
write!(f, "{}{}", SOL_SYMBOL, amount)
}
TokenType::SplToken => {
let amount = token_amount_to_ui_amount(self.amount, self.decimals).ui_amount;
write!(f, "{} tokens", amount)
}
} }
} }
pub fn from(amount: f64, is_sol: bool) -> Self { pub fn sol(amount: u64) -> Self {
let token_type = if is_sol { Self {
TokenType::Sol amount,
} else { decimals: 9,
TokenType::SplToken token_type: TokenType::Sol,
}; }
Self { amount, token_type } }
pub fn spl_token(amount: u64, decimals: u8) -> Self {
Self {
amount,
decimals,
token_type: TokenType::SplToken,
}
} }
} }
@ -53,6 +69,7 @@ impl Add for Token {
if self.token_type == other.token_type { if self.token_type == other.token_type {
Self { Self {
amount: self.amount + other.amount, amount: self.amount + other.amount,
decimals: self.decimals,
token_type: self.token_type, token_type: self.token_type,
} }
} else { } else {