From c18c0b3cf10750c5df3d782859a01440e5053b4e Mon Sep 17 00:00:00 2001
From: xnuter <retunsky@gmail.com>
Date: Thu, 17 Mar 2022 21:38:25 -0700
Subject: [PATCH] Ability to specify multiple options for a header value.

---
 src/configuration.rs      |  5 ++--
 src/http_bench_session.rs | 63 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/src/configuration.rs b/src/configuration.rs
index d07b547..25589c3 100644
--- a/src/configuration.rs
+++ b/src/configuration.rs
@@ -97,7 +97,8 @@ struct HttpOptions {
     /// Target, e.g. https://my-service.com:8443/8kb Can be multiple ones (with random choice balancing).
     #[clap()]
     target: Vec<String>,
-    /// Headers in \"Name:Value\" form. Can be provided multiple times.
+    /// Headers in "Name:Value1" form. E.g. `-H "Authentication:Bearer token" -H "Date:2022-03-17"`
+    /// It can contain multiple values, e.g. "Name:Value1:Value2:Value3". In this case a random one is chosen for each request.
     #[clap(short = 'H', long)]
     header: Vec<String>,
     /// Method. By default GET.
@@ -268,7 +269,7 @@ impl BenchmarkConfig {
                                                 .next()
                                                 .expect("Header name is missing")
                                                 .to_string(),
-                                            split.collect::<Vec<&str>>().join(":"),
+                                            split.map(String::from).collect::<Vec<String>>(),
                                         )
                                     })
                                     .collect(),
diff --git a/src/http_bench_session.rs b/src/http_bench_session.rs
index 980f5f1..8e08e55 100644
--- a/src/http_bench_session.rs
+++ b/src/http_bench_session.rs
@@ -45,7 +45,7 @@ pub struct HttpRequest {
     #[builder(default = "\"GET\".to_string()")]
     method: String,
     #[builder(default)]
-    headers: Vec<(String, String)>,
+    headers: Vec<(String, Vec<String>)>,
     #[builder(default)]
     body: Vec<u8>,
 }
@@ -178,7 +178,8 @@ impl HttpRequest {
             for (key, value) in self.headers.iter() {
                 request_builder = request_builder.header(
                     HeaderName::from_str(key).expect("Header name must be valid at this point"),
-                    HeaderValue::from_str(value).expect("Header value must be valid at this point"),
+                    HeaderValue::from_str(&value[thread_rng().gen_range(0..value.len())])
+                        .expect("Header value must be valid at this point"),
                 );
             }
         }
@@ -242,8 +243,8 @@ mod tests {
                 HttpRequestBuilder::default()
                     .url(vec![format!("{}/1", url)])
                     .headers(vec![
-                        ("x-header".to_string(), "value1".to_string()),
-                        ("x-another-header".to_string(), "value2".to_string()),
+                        ("x-header".to_string(), vec!["value1".to_string()]),
+                        ("x-another-header".to_string(), vec!["value2".to_string()]),
                     ])
                     .build()
                     .unwrap(),
@@ -281,8 +282,8 @@ mod tests {
                     .url(vec![format!("{}/1", url)])
                     .method("PUT".to_string())
                     .headers(vec![
-                        ("x-header".to_string(), "value1".to_string()),
-                        ("x-another-header".to_string(), "value2".to_string()),
+                        ("x-header".to_string(), vec!["value1".to_string()]),
+                        ("x-another-header".to_string(), vec!["value2".to_string()]),
                     ])
                     .body("abcd".as_bytes().to_vec())
                     .build()
@@ -321,8 +322,8 @@ mod tests {
                     .url(vec![format!("{}/1", url)])
                     .method("POST".to_string())
                     .headers(vec![
-                        ("x-header".to_string(), "value1".to_string()),
-                        ("x-another-header".to_string(), "value2".to_string()),
+                        ("x-header".to_string(), vec!["value1".to_string()]),
+                        ("x-another-header".to_string(), vec!["value2".to_string()]),
                     ])
                     .body("abcd".as_bytes().to_vec())
                     .build()
@@ -340,6 +341,52 @@ mod tests {
         assert_eq!("200 OK".to_string(), stats.status);
     }
 
+    #[tokio::test]
+    async fn test_success_multiheader_request() {
+        let m1 = mock("GET", "/1")
+            .match_header("x-header", "value11")
+            .with_status(200)
+            .with_header("content-type", "text/plain")
+            .with_body("abcd")
+            .expect_at_least(1)
+            .create();
+
+        let m2 = mock("GET", "/1")
+            .match_header("x-header", "value12")
+            .with_status(201)
+            .with_header("content-type", "text/plain")
+            .with_body("efg")
+            .expect_at_least(1)
+            .create();
+
+        let url = mockito::server_url().to_string();
+        println!("Url: {}", url);
+        let http_bench = HttpBenchAdapterBuilder::default()
+            .request(
+                HttpRequestBuilder::default()
+                    .url(vec![format!("{}/1", url)])
+                    .method("GET".to_string())
+                    .headers(vec![(
+                        "x-header".to_string(),
+                        vec!["value11".to_string(), "value12".to_string()],
+                    )])
+                    .build()
+                    .unwrap(),
+            )
+            .config(HttpClientConfigBuilder::default().build().unwrap())
+            .build()
+            .unwrap();
+
+        let client = http_bench.build_client().expect("Client is built");
+
+        for _ in 0..128 {
+            http_bench.send_request(&client).await;
+        }
+
+        m1.assert();
+        m2.assert();
+    }
+
     #[tokio::test]
     async fn test_failed_request() {
         let body = "world";