From 771bf46e276cb199319fa7c84fb3936d06de3b90 Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Fri, 27 Nov 2020 10:17:32 +0530
Subject: [PATCH] Enable mock-keys to be created with the bench tool

This can help us create 'fake' key/value pairs for testing purposes and
memory usage analysis.

Signed-off-by: Sayan Nandan <nandansayan@outlook.com>
---
 tdb-bench/src/cli.yml | 11 +++++++++++
 tdb-bench/src/main.rs | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/tdb-bench/src/cli.yml b/tdb-bench/src/cli.yml
index 6af091c3..ae09dc82 100644
--- a/tdb-bench/src/cli.yml
+++ b/tdb-bench/src/cli.yml
@@ -63,3 +63,14 @@ args:
                   value_name: port
                   help: Sets the remote port to connect to
                   takes_value: true
+subcommands:
+        - testkey:
+                  about: This can be used to create 'mock' keys
+                  args:
+                          - count:
+                                    short: c
+                                    required: true
+                                    long: count
+                                    value_name: NUMBEROFKEYS
+                                    help: Sets the number of keys to create
+                                    takes_value: true
diff --git a/tdb-bench/src/main.rs b/tdb-bench/src/main.rs
index d509203f..f5de47a4 100644
--- a/tdb-bench/src/main.rs
+++ b/tdb-bench/src/main.rs
@@ -152,7 +152,6 @@ mod benchtool {
     pub fn runner() {
         let cfg_layout = load_yaml!("./cli.yml");
         let matches = App::from_yaml(cfg_layout).get_matches();
-        let json_out = matches.is_present("json");
         let mut host = match matches.value_of("host") {
             Some(h) => h.to_owned(),
             None => "127.0.0.1".to_owned(),
@@ -168,6 +167,40 @@ mod benchtool {
             },
             None => host.push_str("2003"),
         }
+        let rand = thread_rng();
+        if let Some(matches) = matches.subcommand_matches("testkey") {
+            let numkeys = matches.value_of("count").unwrap();
+            if let Ok(num) = numkeys.parse::<usize>() {
+                let mut np = Netpool::new(10, &host);
+                println!("Generating keys ...");
+                let keys: Vec<String> = (0..num)
+                    .into_iter()
+                    .map(|_| {
+                        let rand_string: String = rand.sample_iter(&Alphanumeric).take(8).collect();
+                        rand_string
+                    })
+                    .collect();
+                let values: Vec<String> = (0..num)
+                    .into_iter()
+                    .map(|_| {
+                        let rand_string: String = rand.sample_iter(&Alphanumeric).take(8).collect();
+                        rand_string
+                    })
+                    .collect();
+                let set_packs: Vec<Vec<u8>> = (0..num)
+                    .map(|idx| terrapipe::proc_query(format!("SET {} {}", keys[idx], values[idx])))
+                    .collect();
+                set_packs.into_iter().for_each(|packet| {
+                    np.execute(packet);
+                });
+                println!("Create mock keys!");
+                return;
+            } else {
+                eprintln!("ERROR: Invalid value for `count`");
+                std::process::exit(0x100);
+            }
+        }
+        let json_out = matches.is_present("json");
         let (max_connections, max_queries, packet_size) = match (
             matches
                 .value_of("connections")