Added ability to set default headers and to copy headers

SYNTAX.md

@@ -29,6 +29,8 @@ plan:
 - `iterations`: Number of loops is going to do (Optional, default: 1)
 - `concurrency`: Number of concurrent iterations. (Optional, default: max)
 - `rampup`: Amount of time it will take to start all iterations. (Optional)
+- `default_headers`: The list of headers you want all requests to share. (Optional)
+- `copy_headers`: The list of headers that you want to copy between requests if it appears in a response. (Optional)
 - `plan`: List of items to do in your benchmark. (Required)
 
 #### Plan items

example/headers.yml

@@ -2,23 +2,39 @@
 # This is an example of how to send custom headers.
 
 ---
-base: 'http://localhost:3000'
+base: 'http://localhost:9000'
 iterations: 1
+default_headers:
+  Authorization: Basic aHR0cHdhdGNoOmY=
+copy_headers:
+  - X-Test
 
 plan:
   - name: Custom headers
     request:
       url: /
       headers:
-        Authorization: Basic aHR0cHdhdGNoOmY=
         X-Foo: Bar
 
   - name: Dynamic Custom headers
     request:
       url: /
       headers:
-        Authorization: Basic aHR0cHdhdGNoOmY=
         X-Foo: Bar {{ item }}
     with_items:
       - 70
       - 73
+
+  - name: Fetch first test header
+    request:
+      url: /header
+
+  - name: Fetch second test header
+    request:
+      url: /header
+    assign: returned_headers
+
+  - name: Check the headers
+    assert:
+      key: returned_headers.headers.x-test
+      value: '2'

example/server/server.js

@@ -97,6 +97,16 @@ app.get('/', function(req, res){
   res.json({ status: ':D' })
 });
 
+app.get('/header', function(req, res){
+  header = req.get('X-Test');
+  if (header) {
+    res.header('X-Test', parseInt(header) + 1);
+  } else {
+    res.header('X-Test', '1');
+  }
+  res.send();
+});
+
 app.delete('/', function(req, res){
   req.session.counter = 1;
   res.json({counter: req.session.counter})

src/actions/exec.rs

@@ -5,7 +5,7 @@ use std::process::Command;
 use yaml_rust::Yaml;
 
 use crate::actions::Runnable;
-use crate::actions::{extract, extract_optional};
+use crate::actions::{extract, extract_optional, yaml_to_json};
 use crate::benchmark::{Context, Pool, Reports};
 use crate::config::Config;
 use crate::interpolator;
@@ -14,6 +14,8 @@ use crate::interpolator;
 pub struct Exec {
   name: String,
   command: String,
+  pub with_item: Option<Yaml>,
+  pub index: Option<u32>,
   pub assign: Option<String>,
 }
 
@@ -22,14 +24,16 @@ impl Exec {
     item["exec"].as_hash().is_some()
   }
 
-  pub fn new(item: &Yaml, _with_item: Option<Yaml>) -> Exec {
+  pub fn new(item: &Yaml, with_item: Option<Yaml>, index: Option<u32>) -> Exec {
     let name = extract(item, "name");
     let command = extract(&item["exec"], "command");
     let assign = extract_optional(item, "assign");
 
     Exec {
       name,
       command,
+      with_item,
+      index,
       assign,
     }
   }
@@ -38,8 +42,15 @@ impl Exec {
 #[async_trait]
 impl Runnable for Exec {
   async fn execute(&self, context: &mut Context, _reports: &mut Reports, _pool: &Pool, config: &Config) {
+    if self.with_item.is_some() {
+      context.insert("item".to_string(), yaml_to_json(self.with_item.clone().unwrap()));
+    }
     if !config.quiet {
-      println!("{:width$} {}", self.name.green(), self.command.cyan().bold(), width = 25);
+      if self.with_item.is_some() {
+        println!("{:width$} ({}) {}", self.name.green(), self.with_item.clone().unwrap().as_str().unwrap(), self.command.cyan().bold(), width = 25);
+      } else {
+        println!("{:width$} {}", self.name.green(), self.command.cyan().bold(), width = 25);
+      }
     }
 
     let final_command = interpolator::Interpolator::new(context).resolve(&self.command, !config.relaxed_interpolations);
@@ -50,6 +61,13 @@ impl Runnable for Exec {
 
     let output: String = String::from_utf8_lossy(&execution.stdout).into();
     let output = output.trim_end().to_string();
+    if !config.quiet {
+      if self.with_item.is_some() {
+        println!("{:width$} ({}) >>> {}", self.name.green(), self.with_item.clone().unwrap().as_str().unwrap(), output.cyan().bold(), width = 25);
+      } else {
+        println!("{:width$} >>> {}", self.name.green(), output.cyan().bold(), width = 25);
+      }
+    }
 
     if let Some(ref key) = self.assign {
       context.insert(key.to_owned(), json!(output));

src/actions/mod.rs

@@ -1,4 +1,5 @@
 use async_trait::async_trait;
+use serde_json::{json, Map, Value};
 use yaml_rust::Yaml;
 
 mod assert;
@@ -63,3 +64,31 @@ pub fn extract<'a>(item: &'a Yaml, attr: &'a str) -> String {
     panic!("Unknown node `{}` => {:?}", attr, item[attr]);
   }
 }
+
+pub fn yaml_to_json(data: Yaml) -> Value {
+  if let Some(b) = data.as_bool() {
+    json!(b)
+  } else if let Some(i) = data.as_i64() {
+    json!(i)
+  } else if let Some(s) = data.as_str() {
+    json!(s)
+  } else if let Some(h) = data.as_hash() {
+    let mut map = Map::new();
+
+    for (key, value) in h.iter() {
+      map.entry(key.as_str().unwrap()).or_insert(yaml_to_json(value.clone()));
+    }
+
+    json!(map)
+  } else if let Some(v) = data.as_vec() {
+    let mut array = Vec::new();
+
+    for value in v.iter() {
+      array.push(yaml_to_json(value.clone()));
+    }
+
+    json!(array)
+  } else {
+    panic!("Unknown Yaml node")
+  }
+}

src/actions/request.rs

@@ -4,7 +4,7 @@ use std::time::{Duration, Instant};
 use async_trait::async_trait;
 use colored::Colorize;
 use reqwest::{
-  header::{self, HeaderMap, HeaderName, HeaderValue},
+  header::{self, HeaderName, HeaderValue},
   ClientBuilder, Method, Response,
 };
 use std::fmt::Write;
@@ -14,7 +14,7 @@ use yaml_rust::Yaml;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Map, Value};
 
-use crate::actions::{extract, extract_optional};
+use crate::actions::{extract, extract_optional, yaml_to_json};
 use crate::benchmark::{Context, Pool, Reports};
 use crate::config::Config;
 use crate::interpolator;
@@ -168,9 +168,16 @@ impl Request {
     };
 
     // Headers
-    let mut headers = HeaderMap::new();
+    let mut headers = config.default_headers.clone();
     headers.insert(header::USER_AGENT, HeaderValue::from_str(USER_AGENT).unwrap());
 
+    if let Some(h) = context.get("headers") {
+      let h: Map<String, Value> = serde_json::from_value(h.clone()).unwrap();
+      for (header, value) in h {
+        headers.insert(HeaderName::from_bytes(header.as_bytes()).unwrap(), HeaderValue::from_str(value.as_str().unwrap()).unwrap());
+      }
+    }
+
     if let Some(cookies) = context.get("cookies") {
       let cookies: Map<String, Value> = serde_json::from_value(cookies.clone()).unwrap();
       let cookie = cookies.iter().map(|(key, value)| format!("{}={}", key, value)).collect::<Vec<_>>().join(";");
@@ -222,34 +229,6 @@ impl Request {
   }
 }
 
-fn yaml_to_json(data: Yaml) -> Value {
-  if let Some(b) = data.as_bool() {
-    json!(b)
-  } else if let Some(i) = data.as_i64() {
-    json!(i)
-  } else if let Some(s) = data.as_str() {
-    json!(s)
-  } else if let Some(h) = data.as_hash() {
-    let mut map = Map::new();
-
-    for (key, value) in h.iter() {
-      map.entry(key.as_str().unwrap()).or_insert(yaml_to_json(value.clone()));
-    }
-
-    json!(map)
-  } else if let Some(v) = data.as_vec() {
-    let mut array = Vec::new();
-
-    for value in v.iter() {
-      array.push(yaml_to_json(value.clone()));
-    }
-
-    json!(array)
-  } else {
-    panic!("Unknown Yaml node")
-  }
-}
-
 #[async_trait]
 impl Runnable for Request {
   async fn execute(&self, context: &mut Context, reports: &mut Reports, pool: &Pool, config: &Config) {
@@ -293,6 +272,13 @@ impl Runnable for Request {
 
           context.insert("cookies".to_string(), json!(cookies));
         }
+        let mut headers = HashMap::new();
+        for header in &config.copy_headers {
+          if let Some(v) = response.headers().get(header) {
+            headers.insert(header, v.to_str().unwrap());
+          }
+        }
+        context.insert("headers".to_string(), json!(headers));
 
         let data = if let Some(ref key) = self.assign {
           let mut headers = Map::new();

src/config.rs

@@ -1,3 +1,6 @@
+use colored::Colorize;
+use linked_hash_map::LinkedHashMap;
+use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
 use yaml_rust::{Yaml, YamlLoader};
 
 use crate::benchmark::Context;
@@ -18,6 +21,8 @@ pub struct Config {
   pub nanosec: bool,
   pub timeout: u64,
   pub verbose: bool,
+  pub default_headers: HeaderMap,
+  pub copy_headers: Vec<String>,
 }
 
 impl Config {
@@ -34,6 +39,23 @@ impl Config {
     let concurrency = read_i64_configuration(config_doc, &interpolator, "concurrency", iterations);
     let rampup = read_i64_configuration(config_doc, &interpolator, "rampup", NRAMPUP);
     let base = read_str_configuration(config_doc, &interpolator, "base", "");
+    let mut default_headers = HeaderMap::new();
+    let hash = read_hashmap_configuration(config_doc, "default_headers", LinkedHashMap::new());
+    for (key, val) in hash.iter() {
+      if let Some(vs) = val.as_str() {
+        default_headers.insert(
+          HeaderName::from_bytes(key.as_str().unwrap().as_bytes()).unwrap(),
+          HeaderValue::from_str(vs).unwrap(),
+        );
+      } else {
+        panic!("{} Headers must be strings!!", "WARNING!".yellow().bold());
+      }
+    }
+    let mut copy_headers = Vec::new();
+    for v in read_list_configuration(config_doc, "copy_headers", Vec::new()).iter() {
+      copy_headers.push(v.as_str().unwrap().to_string());
+    };
+
 
     if concurrency > iterations {
       panic!("The concurrency can not be higher than the number of iterations")
@@ -50,6 +72,8 @@ impl Config {
       nanosec,
       timeout,
       verbose,
+      default_headers,
+      copy_headers,
     }
   }
 }
@@ -73,6 +97,36 @@ fn read_str_configuration(config_doc: &Yaml, interpolator: &interpolator::Interp
   }
 }
 
+fn read_hashmap_configuration(config_doc: &Yaml, name: &str, default: LinkedHashMap<Yaml, Yaml>) -> LinkedHashMap<Yaml, Yaml> {
+  match config_doc[name].as_hash() {
+    Some(value) => {
+      value.clone()
+    }
+    None => {
+      if config_doc[name].as_hash().is_some() {
+        println!("Invalid {} value!", name);
+      }
+
+      default.to_owned()
+    }
+  }
+}
+
+fn read_list_configuration(config_doc: &Yaml, name: &str, default: Vec<Yaml>) -> Vec<Yaml> {
+  match config_doc[name].as_vec() {
+    Some(value) => {
+      value.clone()
+    }
+    None => {
+      if config_doc[name].as_vec().is_some() {
+        println!("Invalid {} value!", name);
+      }
+
+      default.to_owned()
+    }
+  }
+}
+
 fn read_i64_configuration(config_doc: &Yaml, interpolator: &interpolator::Interpolator, name: &str, default: i64) -> i64 {
   let value = if let Some(value) = config_doc[name].as_i64() {
     Some(value)

src/expandable/include.rs

@@ -5,7 +5,7 @@ use crate::interpolator::INTERPOLATION_REGEX;
 
 use crate::actions;
 use crate::benchmark::Benchmark;
-use crate::expandable::{include, multi_csv_request, multi_file_request, multi_iter_request, multi_request};
+use crate::expandable::{include, multi_csv_request, multi_exec, multi_file_request, multi_iter_request, multi_request};
 use crate::tags::Tags;
 
 use crate::reader;
@@ -36,7 +36,9 @@ pub fn expand_from_filepath(parent_path: &str, benchmark: &mut Benchmark, access
       continue;
     }
 
-    if multi_request::is_that_you(item) {
+    if multi_exec::is_that_you(item) {
+      multi_exec::expand(item, benchmark);
+    } else if multi_request::is_that_you(item) {
       multi_request::expand(item, benchmark);
     } else if multi_iter_request::is_that_you(item) {
       multi_iter_request::expand(item, benchmark);
@@ -49,7 +51,7 @@ pub fn expand_from_filepath(parent_path: &str, benchmark: &mut Benchmark, access
     } else if actions::Delay::is_that_you(item) {
       benchmark.push(Box::new(actions::Delay::new(item, None)));
     } else if actions::Exec::is_that_you(item) {
-      benchmark.push(Box::new(actions::Exec::new(item, None)));
+      benchmark.push(Box::new(actions::Exec::new(item, None, None)));
     } else if actions::Assign::is_that_you(item) {
       benchmark.push(Box::new(actions::Assign::new(item, None)));
     } else if actions::Assert::is_that_you(item) {

src/expandable/mod.rs

@@ -1,6 +1,7 @@
 pub mod include;
 
 mod multi_csv_request;
+mod multi_exec;
 mod multi_file_request;
 mod multi_iter_request;
 mod multi_request;