Add --allocations-csv option (solana-labs#14)

* Add allocations-csv option * Add tests or GTFO * Apply review feedback * apply feedback * Add read_allocations function
garious · May 5, 2020 · 0c36708 · 0c36708
1 parent 4e20253
commit 0c36708
Show file tree

Hide file tree

Showing 5 changed files with 155 additions and 43 deletions.
diff --git a/README.md b/README.md
@@ -17,7 +17,7 @@ primary_address,bid_amount_dollars
 ```
 
 ```bash
-solana-tokens distribute-tokens --from <KEYPAIR> --dollars-per-sol <NUMBER> --bids-csv <BIDS_CSV> <TRANSACTION_LOG> --fee-payer <KEYPAIR>
+solana-tokens distribute-tokens --from <KEYPAIR> --dollars-per-sol <NUMBER> --input_csv <BIDS_CSV> <TRANSACTION_LOG> --fee-payer <KEYPAIR>
 ```
 
 Example transaction log before:
@@ -31,7 +31,7 @@ Send tokens to the recipients in `<BIDS_CSV>` if the distribution is
 not already recordered in the transaction log.
 
 ```bash
-solana-tokens distribute-tokens --from <KEYPAIR> --dollars-per-sol <NUMBER> --bids-csv <BIDS_CSV> <TRANSACTION_LOG> --fee-payer <KEYPAIR>
+solana-tokens distribute-tokens --from <KEYPAIR> --dollars-per-sol <NUMBER> --input_csv <BIDS_CSV> <TRANSACTION_LOG> --fee-payer <KEYPAIR>
 ```
 
 Example output:
@@ -60,7 +60,7 @@ List the differences between a list of expected distributions and the record of
 transactions have already been sent.
 
 ```bash
-solana-tokens distribute-tokens --dollars-per-sol <NUMBER> --dry-run --bids-csv <BIDS_CSV> <TRANSACTION_LOG>
+solana-tokens distribute-tokens --dollars-per-sol <NUMBER> --dry-run --input_csv <BIDS_CSV> <TRANSACTION_LOG>
 ```
 
 Example bids.csv:

diff --git a/src/arg_parser.rs b/src/arg_parser.rs
@@ -42,20 +42,25 @@ where
                         .help("Transactions database file"),
                 )
                 .arg(
-                    Arg::with_name("bids_csv")
-                        .long("bids-csv")
+                    Arg::with_name("from_bids")
+                        .long("from-bids")
+                        .help("Input CSV contains bids in dollars, not allocations in SOL"),
+                )
+                .arg(
+                    Arg::with_name("input_csv")
+                        .long("input_csv")
                         .required(true)
                         .takes_value(true)
                         .value_name("FILE")
-                        .help("Bids CSV file"),
+                        .help("Input CSV file"),
                 )
                 .arg(
                     Arg::with_name("dollars_per_sol")
                         .long("dollars-per-sol")
                         .required(true)
                         .takes_value(true)
                         .value_name("NUMBER")
-                        .help("Dollars per SOL"),
+                        .help("Dollars per SOL, if input CSV contains bids")
                 )
                 .arg(
                     Arg::with_name("dry_run")
@@ -149,8 +154,8 @@ where
             SubCommand::with_name("balances")
                 .about("Balance of each account")
                 .arg(
-                    Arg::with_name("bids_csv")
-                        .long("bids-csv")
+                    Arg::with_name("input_csv")
+                        .long("input_csv")
                         .required(true)
                         .takes_value(true)
                         .value_name("FILE")
@@ -170,9 +175,10 @@ where
 
 fn parse_distribute_tokens_args(matches: &ArgMatches<'_>) -> DistributeTokensArgs<String> {
     DistributeTokensArgs {
-        bids_csv: value_t_or_exit!(matches, "bids_csv", String),
+        input_csv: value_t_or_exit!(matches, "input_csv", String),
+        from_bids: matches.is_present("from_bids"),
         transactions_db: value_t_or_exit!(matches, "transactions_db", String),
-        dollars_per_sol: value_t_or_exit!(matches, "dollars_per_sol", f64),
+        dollars_per_sol: value_t!(matches, "dollars_per_sol", f64).ok(),
         dry_run: matches.is_present("dry_run"),
         sender_keypair: value_t!(matches, "sender_keypair", String).ok(),
         fee_payer: value_t!(matches, "fee_payer", String).ok(),
@@ -194,7 +200,7 @@ fn parse_distribute_stake_args(matches: &ArgMatches<'_>) -> DistributeStakeArgs<
 
 fn parse_balances_args(matches: &ArgMatches<'_>) -> BalancesArgs {
     BalancesArgs {
-        bids_csv: value_t_or_exit!(matches, "bids_csv", String),
+        input_csv: value_t_or_exit!(matches, "input_csv", String),
         dollars_per_sol: value_t_or_exit!(matches, "dollars_per_sol", f64),
     }
 }

diff --git a/src/args.rs b/src/args.rs
@@ -5,9 +5,10 @@ use solana_sdk::{pubkey::Pubkey, signature::Signer};
 use std::error::Error;
 
 pub struct DistributeTokensArgs<K> {
-    pub bids_csv: String,
+    pub input_csv: String,
+    pub from_bids: bool,
     pub transactions_db: String,
-    pub dollars_per_sol: f64,
+    pub dollars_per_sol: Option<f64>,
     pub dry_run: bool,
     pub sender_keypair: Option<K>,
     pub fee_payer: Option<K>,
@@ -25,7 +26,7 @@ pub struct DistributeStakeArgs<P, K> {
 }
 
 pub struct BalancesArgs {
-    pub bids_csv: String,
+    pub input_csv: String,
     pub dollars_per_sol: f64,
 }
 
@@ -49,7 +50,8 @@ pub fn resolve_command(
             let mut wallet_manager = maybe_wallet_manager()?;
             let matches = ArgMatches::default();
             let resolved_args = DistributeTokensArgs {
-                bids_csv: args.bids_csv,
+                input_csv: args.input_csv,
+                from_bids: args.from_bids,
                 transactions_db: args.transactions_db,
                 dollars_per_sol: args.dollars_per_sol,
                 dry_run: args.dry_run,

diff --git a/src/tokens.rs b/src/tokens.rs
@@ -272,26 +272,42 @@ fn set_transaction_info(
     Ok(())
 }
 
+fn read_allocations(
+    args: &DistributeTokensArgs<Box<dyn Signer>>
+) -> Vec<Allocation> {
+    let rdr = ReaderBuilder::new()
+        .trim(Trim::All)
+        .from_path(&args.input_csv);
+    if args.from_bids {
+        let bids: Vec<Bid> = rdr.unwrap().deserialize().map(|bid| bid.unwrap()).collect();
+        bids
+            .into_iter()
+            .map(|bid| create_allocation(&bid, args.dollars_per_sol.unwrap()))
+            .collect()
+    } else {
+        rdr.unwrap().deserialize().map(|entry| entry.unwrap()).collect()
+    }
+}
+
 pub fn process_distribute_tokens<T: Client>(
     client: &ThinClient<T>,
     args: &DistributeTokensArgs<Box<dyn Signer>>,
 ) -> Result<(), Error> {
-    let mut rdr = ReaderBuilder::new()
-        .trim(Trim::All)
-        .from_path(&args.bids_csv)?;
-    let bids: Vec<Bid> = rdr.deserialize().map(|bid| bid.unwrap()).collect();
-    let mut allocations: Vec<Allocation> = bids
-        .into_iter()
-        .map(|bid| create_allocation(&bid, args.dollars_per_sol))
-        .collect();
+    let mut allocations: Vec<Allocation> = read_allocations(&args);
 
     let starting_total_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
     println!(
-        "{} ◎{} (${})",
+        "{} ◎{}",
         style(format!("{}", "Total in allocations_csv:")).bold(),
         starting_total_tokens,
-        starting_total_tokens * args.dollars_per_sol,
     );
+    if let Some(dollars_per_sol) = args.dollars_per_sol {
+        println!(
+            "{} ${}",
+            style(format!("{}", "Total in allocations_csv:")).bold(),
+            starting_total_tokens * dollars_per_sol,
+        );
+    }
 
     let mut db = open_db(&args.transactions_db)?;
     let transaction_infos = read_transaction_infos(&db);
@@ -335,23 +351,41 @@ pub fn process_distribute_tokens<T: Client>(
     let distributed_tokens: f64 = transaction_infos.iter().map(|x| x.amount).sum();
     let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
     println!(
-        "{} ◎{} (${})",
+        "{} ◎{}",
         style(format!("{}", "Distributed:")).bold(),
         distributed_tokens,
-        distributed_tokens * args.dollars_per_sol,
     );
+    if let Some(dollars_per_sol) = args.dollars_per_sol {
+        println!(
+            "{} ${}",
+            style(format!("{}", "Distributed:")).bold(),
+            distributed_tokens * dollars_per_sol,
+        );
+    }
     println!(
-        "{} ◎{} (${})",
+        "{} ◎{}",
         style(format!("{}", "Undistributed:")).bold(),
         undistributed_tokens,
-        undistributed_tokens * args.dollars_per_sol,
     );
+    if let Some(dollars_per_sol) = args.dollars_per_sol {
+        println!(
+            "{} ${}",
+            style(format!("{}", "Undistributed:")).bold(),
+            undistributed_tokens * dollars_per_sol,
+        );
+    }
     println!(
-        "{} ◎{} (${})\n",
+        "{} ◎{}",
         style(format!("{}", "Total:")).bold(),
         distributed_tokens + undistributed_tokens,
-        (distributed_tokens + undistributed_tokens) * args.dollars_per_sol,
     );
+    if let Some(dollars_per_sol) = args.dollars_per_sol {
+        println!(
+            "{} ${}",
+            style(format!("{}", "Total:")).bold(),
+            (distributed_tokens + undistributed_tokens) * dollars_per_sol,
+        );
+    }
 
     distribute_tokens(client, &mut db, &allocations, args)?;
 
@@ -391,7 +425,7 @@ pub fn process_balances<T: Client>(
 ) -> Result<(), csv::Error> {
     let mut rdr = ReaderBuilder::new()
         .trim(Trim::All)
-        .from_path(&args.bids_csv)?;
+        .from_path(&args.input_csv)?;
     let bids: Vec<Bid> = rdr.deserialize().map(|bid| bid.unwrap()).collect();
     let allocations: Vec<Allocation> = bids
         .into_iter()
@@ -426,7 +460,7 @@ pub fn process_balances<T: Client>(
 
 use solana_sdk::{pubkey::Pubkey, signature::Keypair};
 use tempfile::{tempdir, NamedTempFile};
-pub fn test_process_distribute_with_client<C: Client>(client: C, sender_keypair: Keypair) {
+pub fn test_process_distribute_bids_with_client<C: Client>(client: C, sender_keypair: Keypair) {
     let thin_client = ThinClient(client);
     let fee_payer = Keypair::new();
     thin_client
@@ -439,7 +473,7 @@ pub fn test_process_distribute_with_client<C: Client>(client: C, sender_keypair:
         accepted_amount_dollars: 1000.0,
     };
     let bids_file = NamedTempFile::new().unwrap();
-    let bids_csv = bids_file.path().to_str().unwrap().to_string();
+    let input_csv = bids_file.path().to_str().unwrap().to_string();
     let mut wtr = csv::WriterBuilder::new().from_writer(bids_file);
     wtr.serialize(&bid).unwrap();
     wtr.flush().unwrap();
@@ -456,15 +490,16 @@ pub fn test_process_distribute_with_client<C: Client>(client: C, sender_keypair:
         sender_keypair: Some(Box::new(sender_keypair)),
         fee_payer: Some(Box::new(fee_payer)),
         dry_run: false,
-        bids_csv,
+        from_bids: true,
+        input_csv,
         transactions_db: transactions_db.clone(),
-        dollars_per_sol: 0.22,
+        dollars_per_sol: Some(0.22),
     };
     process_distribute_tokens(&thin_client, &args).unwrap();
     let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap());
     assert_eq!(transaction_infos.len(), 1);
     assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string());
-    let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol;
+    let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol.unwrap();
     assert_eq!(transaction_infos[0].amount, expected_amount);
 
     assert_eq!(
@@ -477,7 +512,68 @@ pub fn test_process_distribute_with_client<C: Client>(client: C, sender_keypair:
     let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap());
     assert_eq!(transaction_infos.len(), 1);
     assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string());
-    let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol;
+    let expected_amount = bid.accepted_amount_dollars / args.dollars_per_sol.unwrap();
+    assert_eq!(transaction_infos[0].amount, expected_amount);
+
+    assert_eq!(
+        thin_client.get_balance(&alice_pubkey).unwrap(),
+        sol_to_lamports(expected_amount),
+    );
+}
+
+pub fn test_process_distribute_allocations_with_client<C: Client>(client: C, sender_keypair: Keypair) {
+    let thin_client = ThinClient(client);
+    let fee_payer = Keypair::new();
+    thin_client
+        .transfer(sol_to_lamports(1.0), &sender_keypair, &fee_payer.pubkey())
+        .unwrap();
+
+    let alice_pubkey = Pubkey::new_rand();
+    let allocation = Allocation {
+        recipient: alice_pubkey.to_string(),
+        amount: 1000.0,
+    };
+    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);
+    wtr.serialize(&allocation).unwrap();
+    wtr.flush().unwrap();
+
+    let dir = tempdir().unwrap();
+    let transactions_db = dir
+        .path()
+        .join("transactions.csv")
+        .to_str()
+        .unwrap()
+        .to_string();
+
+    let args: DistributeTokensArgs<Box<dyn Signer>> = DistributeTokensArgs {
+        sender_keypair: Some(Box::new(sender_keypair)),
+        fee_payer: Some(Box::new(fee_payer)),
+        dry_run: false,
+        input_csv,
+        from_bids: false,
+        transactions_db: transactions_db.clone(),
+        dollars_per_sol: None,
+    };
+    process_distribute_tokens(&thin_client, &args).unwrap();
+    let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap());
+    assert_eq!(transaction_infos.len(), 1);
+    assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string());
+    let expected_amount = allocation.amount;
+    assert_eq!(transaction_infos[0].amount, expected_amount);
+
+    assert_eq!(
+        thin_client.get_balance(&alice_pubkey).unwrap(),
+        sol_to_lamports(expected_amount),
+    );
+
+    // Now, run it again, and check there's no double-spend.
+    process_distribute_tokens(&thin_client, &args).unwrap();
+    let transaction_infos = read_transaction_infos(&open_db(&transactions_db).unwrap());
+    assert_eq!(transaction_infos.len(), 1);
+    assert_eq!(transaction_infos[0].recipient, alice_pubkey.to_string());
+    let expected_amount = allocation.amount;
     assert_eq!(transaction_infos[0].amount, expected_amount);
 
     assert_eq!(
@@ -589,11 +685,19 @@ mod tests {
     use solana_sdk::genesis_config::create_genesis_config;
 
     #[test]
-    fn test_process_distribute() {
+    fn test_process_distribute_bids() {
+        let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0));
+        let bank = Bank::new(&genesis_config);
+        let bank_client = BankClient::new(bank);
+        test_process_distribute_bids_with_client(bank_client, sender_keypair);
+    }
+
+    #[test]
+    fn test_process_distribute_allocations() {
         let (genesis_config, sender_keypair) = create_genesis_config(sol_to_lamports(9_000_000.0));
         let bank = Bank::new(&genesis_config);
         let bank_client = BankClient::new(bank);
-        test_process_distribute_with_client(bank_client, sender_keypair);
+        test_process_distribute_allocations_with_client(bank_client, sender_keypair);
     }
 
     #[test]

diff --git a/tests/tokens.rs b/tests/tokens.rs
@@ -1,7 +1,7 @@
 use solana_client::rpc_client::RpcClient;
 use solana_core::validator::{TestValidator, TestValidatorOptions};
 use solana_sdk::native_token::sol_to_lamports;
-use solana_tokens::tokens::test_process_distribute_with_client;
+use solana_tokens::tokens::test_process_distribute_bids_with_client;
 use std::fs::remove_dir_all;
 
 #[test]
@@ -12,7 +12,7 @@ fn test_process_distribute_with_rpc_client() {
         ..TestValidatorOptions::default()
     });
     let rpc_client = RpcClient::new_socket(validator.leader_data.rpc);
-    test_process_distribute_with_client(rpc_client, validator.alice);
+    test_process_distribute_bids_with_client(rpc_client, validator.alice);
 
     validator.server.close().unwrap();
     remove_dir_all(validator.ledger_path).unwrap();