From d278b71cb22b1df7f9d614eda1ffa414c19d80b4 Mon Sep 17 00:00:00 2001
From: Robert Kelly <robert.l.kelly3@gmail.com>
Date: Fri, 11 May 2018 12:55:05 -0400
Subject: [PATCH] added tests and utility method for key generation

---
 src/bin/client-demo.rs  | 10 ++--------
 src/bin/genesis-demo.rs | 10 ++--------
 src/signature.rs        | 44 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/src/bin/client-demo.rs b/src/bin/client-demo.rs
index bc895d5755..3b459a5ad2 100644
--- a/src/bin/client-demo.rs
+++ b/src/bin/client-demo.rs
@@ -96,16 +96,10 @@ fn main() {
     let rnd = GenKeys::new(demo.mint.keypair().public_key_bytes());
     let tokens_per_user = 1_000;
 
-    let users: Vec<_> = (0..demo.num_accounts)
-        .into_par_iter()
-        .map(|_| {
-            let pkcs8 = rnd.new_key();
-            (pkcs8, tokens_per_user)
-        })
-        .collect();
+    let users = rnd.gen_n_keys(demo.num_accounts, tokens_per_user);
 
     println!("Creating keypairs...");
-    let txs = users.len() / 2;
+    let txs = demo.num_accounts / 2;
     let keypairs: Vec<_> = users
         .into_par_iter()
         .map(|(pkcs8, _)| KeyPair::from_pkcs8(Input::from(&pkcs8)).unwrap())
diff --git a/src/bin/genesis-demo.rs b/src/bin/genesis-demo.rs
index dd9ef3b717..10ea4b9ef5 100644
--- a/src/bin/genesis-demo.rs
+++ b/src/bin/genesis-demo.rs
@@ -37,17 +37,11 @@ fn main() {
     });
 
     let rnd = GenKeys::new(demo.mint.keypair().public_key_bytes());
+    let num_accounts = demo.num_accounts;
     let tokens_per_user = 1_000;
 
-    let users: Vec<_> = (0..demo.num_accounts)
-        .into_iter()
-        .map(|_| {
-            let pkcs8 = rnd.new_key();
-            (pkcs8, tokens_per_user)
-        })
-        .collect();
+    let users = rnd.gen_n_keys(num_accounts, tokens_per_user);
 
-    let num_accounts = users.len();
     let last_id = demo.mint.last_id();
     let mint_keypair = demo.mint.keypair();
 
diff --git a/src/signature.rs b/src/signature.rs
index d7f2bb0cfb..f554fd0b8c 100644
--- a/src/signature.rs
+++ b/src/signature.rs
@@ -3,6 +3,7 @@
 use generic_array::typenum::{U32, U64};
 use generic_array::GenericArray;
 use rand::{ChaChaRng, Rng, SeedableRng};
+use rayon::prelude::*;
 use ring::error::Unspecified;
 use ring::rand::SecureRandom;
 use ring::signature::Ed25519KeyPair;
@@ -63,6 +64,17 @@ impl GenKeys {
     pub fn new_key(&self) -> Vec<u8> {
         KeyPair::generate_pkcs8(self).unwrap().to_vec()
     }
+
+    pub fn gen_n_keys(&self, n_keys: i64, tokens_per_user: i64) -> Vec<(Vec<u8>, i64)> {
+        let users: Vec<_> = (0..n_keys)
+            .into_par_iter()
+            .map(|_| {
+                let pkcs8 = self.new_key();
+                (pkcs8, tokens_per_user)
+            })
+            .collect();
+        users
+    }
 }
 
 impl SecureRandom for GenKeys {
@@ -72,3 +84,35 @@ impl SecureRandom for GenKeys {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashSet;
+    use std::iter::FromIterator;
+
+    #[test]
+    fn test_new_key_is_redundant() {
+        let seed: &[_] = &[1, 2, 3, 4];
+        let rnd = GenKeys::new(seed);
+        let rnd2 = GenKeys::new(seed);
+
+        for _ in 0..100 {
+            assert_eq!(rnd.new_key(), rnd2.new_key());
+        }
+    }
+
+    #[test]
+    fn test_gen_n_keys() {
+        let seed: &[_] = &[1, 2, 3, 4];
+        let rnd = GenKeys::new(seed);
+        let rnd2 = GenKeys::new(seed);
+
+        let users1 = rnd.gen_n_keys(50, 1);
+        let users2 = rnd2.gen_n_keys(50, 1);
+
+        let users1_set: HashSet<(Vec<u8>, i64)> = HashSet::from_iter(users1.iter().cloned());
+        let users2_set: HashSet<(Vec<u8>, i64)> = HashSet::from_iter(users2.iter().cloned());
+        assert_eq!(users1_set, users2_set);
+    }
+}