diff --git a/SYNTAX.md b/SYNTAX.md index bb48617..069bb87 100644 --- a/SYNTAX.md +++ b/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 diff --git a/example/headers.yml b/example/headers.yml index 56786af..35d9f5e 100644 --- a/example/headers.yml +++ b/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' diff --git a/example/server/server.js b/example/server/server.js index e58c5d3..0784391 100644 --- a/example/server/server.js +++ b/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}) diff --git a/src/actions/exec.rs b/src/actions/exec.rs index 2c32e03..b6e71c0 100644 --- a/src/actions/exec.rs +++ b/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, + pub index: Option, pub assign: Option, } @@ -22,7 +24,7 @@ impl Exec { item["exec"].as_hash().is_some() } - pub fn new(item: &Yaml, _with_item: Option) -> Exec { + pub fn new(item: &Yaml, with_item: Option, index: Option) -> Exec { let name = extract(item, "name"); let command = extract(&item["exec"], "command"); let assign = extract_optional(item, "assign"); @@ -30,6 +32,8 @@ impl Exec { 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)); diff --git a/src/actions/mod.rs b/src/actions/mod.rs index 46c9d9e..deed451 100644 --- a/src/actions/mod.rs +++ b/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") + } +} diff --git a/src/actions/request.rs b/src/actions/request.rs index 108f19b..d193ef8 100644 --- a/src/actions/request.rs +++ b/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 = 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 = serde_json::from_value(cookies.clone()).unwrap(); let cookie = cookies.iter().map(|(key, value)| format!("{}={}", key, value)).collect::>().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(); diff --git a/src/config.rs b/src/config.rs index edfd2f2..9dd16d8 100644 --- a/src/config.rs +++ b/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, } 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) -> LinkedHashMap { + 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) -> Vec { + 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) diff --git a/src/expandable/include.rs b/src/expandable/include.rs index 2579b49..c38cb0c 100644 --- a/src/expandable/include.rs +++ b/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) { diff --git a/src/expandable/mod.rs b/src/expandable/mod.rs index 4ce5a47..470b480 100644 --- a/src/expandable/mod.rs +++ b/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; diff --git a/src/expandable/multi_exec.rs b/src/expandable/multi_exec.rs new file mode 100644 index 0000000..29c2dde --- /dev/null +++ b/src/expandable/multi_exec.rs @@ -0,0 +1,93 @@ +use rand::seq::SliceRandom; +use rand::thread_rng; +use yaml_rust::Yaml; + +use super::pick; +use crate::actions::Exec; +use crate::benchmark::Benchmark; +use crate::interpolator::INTERPOLATION_REGEX; + +pub fn is_that_you(item: &Yaml) -> bool { + item["exec"].as_hash().is_some() && item["with_items"].as_vec().is_some() +} + +pub fn expand(item: &Yaml, benchmark: &mut Benchmark) { + if let Some(with_items) = item["with_items"].as_vec() { + let mut with_items_list = with_items.clone(); + + if let Some(shuffle) = item["shuffle"].as_bool() { + if shuffle { + let mut rng = thread_rng(); + with_items_list.shuffle(&mut rng); + } + } + + let pick = pick(item, &with_items_list); + for (index, with_item) in with_items_list.iter().take(pick).enumerate() { + let index = index as u32; + + let value: &str = with_item.as_str().unwrap_or(""); + + if INTERPOLATION_REGEX.is_match(value) { + panic!("Interpolations not supported in 'with_items' children!"); + } + + benchmark.push(Box::new(Exec::new(item, Some(with_item.clone()), Some(index)))); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn expand_multi() { + let text = "---\nname: foobar\nexec:\n command: echo {{ item }}\nwith_items:\n - 1\n - 2\n - 3"; + let docs = yaml_rust::YamlLoader::load_from_str(text).unwrap(); + let doc = &docs[0]; + let mut benchmark: Benchmark = Benchmark::new(); + + expand(doc, &mut benchmark); + + assert_eq!(is_that_you(doc), true); + assert_eq!(benchmark.len(), 3); + } + + #[test] + fn expand_multi_should_limit_execs_using_the_pick_option() { + let text = "---\nname: foobar\nexec:\n command: echo {{ item }}\npick: 2\nwith_items:\n - 1\n - 2\n - 3"; + let docs = yaml_rust::YamlLoader::load_from_str(text).unwrap(); + let doc = &docs[0]; + let mut benchmark: Benchmark = Benchmark::new(); + + expand(doc, &mut benchmark); + + assert_eq!(is_that_you(doc), true); + assert_eq!(benchmark.len(), 2); + } + + #[test] + fn expand_multi_should_work_with_pick_and_shuffle() { + let text = "---\nname: foobar\nexec:\n command: echo {{ item }}\npick: 1\nshuffle: true\nwith_items:\n - 1\n - 2\n - 3"; + let docs = yaml_rust::YamlLoader::load_from_str(text).unwrap(); + let doc = &docs[0]; + let mut benchmark: Benchmark = Benchmark::new(); + + expand(doc, &mut benchmark); + + assert_eq!(is_that_you(doc), true); + assert_eq!(benchmark.len(), 1); + } + + #[test] + #[should_panic] + fn runtime_expand() { + let text = "---\nname: foobar\nexec:\n command: echo {{ item }}\nwith_items:\n - 1\n - 2\n - foo{{ memory }}"; + let docs = yaml_rust::YamlLoader::load_from_str(text).unwrap(); + let doc = &docs[0]; + let mut benchmark: Benchmark = Benchmark::new(); + + expand(doc, &mut benchmark); + } +}