Interpolate env variables

README.md

@@ -89,6 +89,10 @@ plan:
     request:
       url: http://localhost:9000/api/users.json
 
+  - name: Interpolate environment variables
+    request:
+      url: http://localhost:9000/api/{{ EDITOR }}
+
   - name: Support for POST method
     request:
       url: /api/users
@@ -166,6 +170,7 @@ This is the list of all features supported by the current version of `drill`:
 - **Delay:** introduce controlled delay between requests. Example: [assigns.yml](./example/assigns.yml)
 - **Dynamic urls:** execute requests with dynamic interpolations in the url, like `/api/users/{{ item }}`
 - **Dynamic headers:** execute requests with dynamic headers. Example: [headers.yml](./example/headers.yml)
+- **Interpolate environment variables:** set environment variables, like `/api/users/{{ EDITOR }}`
 - **Request dependencies:** create dependencies between requests with `assign` and url interpolations.
 - **Split files:** organize your benchmarks in multiple files and include them.
 - **CSV support:** read CSV files and build N requests fill dynamic interpolations with CSV data.

example/env.yml

@@ -0,0 +1,9 @@
+---
+
+base: 'http://localhost:9000'
+iterations: '{{ ITERATIONS }}'
+
+plan:
+  - name: Fetch users
+    request:
+      url: /api/{{ EDITOR }}.json

src/config.rs

@@ -1,5 +1,7 @@
 use yaml_rust::{Yaml, YamlLoader};
 
+use crate::benchmark::Context;
+use crate::interpolator;
 use crate::reader;
 
 const NCONCURRENCY: i64 = 1;
@@ -25,11 +27,14 @@ impl Config {
     let config_docs = YamlLoader::load_from_str(config_file.as_str()).unwrap();
     let config_doc = &config_docs[0];
 
-    let concurrency = read_i64_configuration(config_doc, "concurrency", NCONCURRENCY);
+    let context: Context = Context::new();
+    let interpolator = interpolator::Interpolator::new(&context);
+
+    let concurrency = read_i64_configuration(config_doc, &interpolator, "concurrency", NCONCURRENCY);
     let threads = std::cmp::min(num_cpus::get(), concurrency as usize);
-    let iterations = read_i64_configuration(config_doc, "iterations", NITERATIONS);
-    let rampup = read_i64_configuration(config_doc, "rampup", NRAMPUP);
-    let base = config_doc["base"].as_str().unwrap().to_owned();
+    let iterations = read_i64_configuration(config_doc, &interpolator, "iterations", NITERATIONS);
+    let rampup = read_i64_configuration(config_doc, &interpolator, "rampup", NRAMPUP);
+    let base = read_str_configuration(config_doc, &interpolator, "base", "");
 
     Config {
       base,
@@ -45,8 +50,35 @@ impl Config {
   }
 }
 
-fn read_i64_configuration(config_doc: &Yaml, name: &str, default: i64) -> i64 {
-  match config_doc[name].as_i64() {
+fn read_str_configuration(config_doc: &Yaml, interpolator: &interpolator::Interpolator, name: &str, default: &str) -> String {
+  match config_doc[name].as_str() {
+    Some(value) => {
+      if value.contains('{') {
+        interpolator.resolve(&value, true).to_owned()
+      } else {
+        value.to_owned()
+      }
+    }
+    None => {
+      if config_doc[name].as_str().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)
+  } else if let Some(key) = config_doc[name].as_str() {
+    interpolator.resolve(&key, false).parse::<i64>().ok()
+  } else {
+    None
+  };
+
+  match value {
     Some(value) => {
       if value < 0 {
         println!("Invalid negative {} value!", name);

src/interpolator.rs

@@ -35,6 +35,10 @@ impl<'a> Interpolator<'a> {
           return item;
         }
 
+        if let Some(item) = self.resolve_environment_interpolation(capture) {
+          return item;
+        }
+
         if strict {
           panic!("Unknown '{}' variable!", &capture);
         }
@@ -46,6 +50,13 @@ impl<'a> Interpolator<'a> {
       .to_string()
   }
 
+  fn resolve_environment_interpolation(&self, value: &str) -> Option<String> {
+    match std::env::vars().find(|tuple| tuple.0 == value) {
+      Some(tuple) => Some(tuple.1),
+      _ => None,
+    }
+  }
+
   fn resolve_context_interpolation(&self, cap_path: Vec<&str>) -> Option<String> {
     let (cap_root, cap_tail) = cap_path.split_at(1);
 
@@ -130,4 +141,16 @@ mod tests {
 
     assert_eq!(interpolated, "http://example.com/postalcode/{{ 5digitzip }}/view/{{ 5digitzip }}");
   }
+
+  #[test]
+  fn interpolates_environment_variables() {
+    std::env::set_var("FOO", "BAR");
+
+    let context: Context = Context::new();
+    let interpolator = Interpolator::new(&context);
+    let url = String::from("http://example.com/postalcode/{{ FOO }}");
+    let interpolated = interpolator.resolve(&url, true);
+
+    assert_eq!(interpolated, "http://example.com/postalcode/BAR");
+  }
 }