Use HeaderMap for headers in mock and response

lipanski · Oct 30, 2023 · bcdcb2a · bcdcb2a
1 parent b9007ea
commit bcdcb2a
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 27 deletions.
diff --git a/src/matcher.rs b/src/matcher.rs
@@ -1,4 +1,5 @@
 use assert_json_diff::{assert_json_matches_no_panic, CompareMode};
+use hyper::header::HeaderValue;
 use regex::Regex;
 use std::collections::HashMap;
 use std::convert::From;
@@ -108,7 +109,7 @@ impl fmt::Display for Matcher {
 }
 
 impl Matcher {
-    pub(crate) fn matches_values(&self, header_values: &[&str]) -> bool {
+    pub(crate) fn matches_values(&self, header_values: &[&HeaderValue]) -> bool {
         match self {
             Matcher::Missing => header_values.is_empty(),
             // AnyOf([…Missing…]) is handled here, but
@@ -122,7 +123,12 @@ impl Matcher {
                 matchers.iter().all(|m| m.matches_values(header_values))
             }
             _ => {
-                !header_values.is_empty() && header_values.iter().all(|val| self.matches_value(val))
+                !header_values.is_empty()
+                    && header_values.iter().all(|val| {
+                        val.to_str()
+                            .map(|val| self.matches_value(val))
+                            .unwrap_or(false)
+                    })
             }
         }
     }

diff --git a/src/mock.rs b/src/mock.rs
@@ -5,6 +5,8 @@ use crate::server::RemoteMock;
 use crate::server::State;
 use crate::Request;
 use crate::{Error, ErrorKind};
+use hyper::header::IntoHeaderName;
+use hyper::HeaderMap;
 use hyper::StatusCode;
 use rand::distributions::Alphanumeric;
 use rand::{thread_rng, Rng};
@@ -22,7 +24,7 @@ pub struct InnerMock {
     pub(crate) id: String,
     pub(crate) method: String,
     pub(crate) path: PathAndQueryMatcher,
-    pub(crate) headers: Vec<(String, Matcher)>,
+    pub(crate) headers: HeaderMap<Matcher>,
     pub(crate) body: Matcher,
     pub(crate) response: Response,
     pub(crate) hits: usize,
@@ -40,8 +42,8 @@ impl fmt::Display for InnerMock {
         formatted.push(' ');
         formatted.push_str(&self.path.to_string());
 
-        for &(ref key, ref value) in &self.headers {
-            formatted.push_str(key);
+        for (key, value) in &self.headers {
+            formatted.push_str(key.as_str());
             formatted.push_str(": ");
             formatted.push_str(&value.to_string());
             formatted.push_str("\r\n");
@@ -110,7 +112,7 @@ impl Mock {
                 .collect(),
             method: method.to_owned().to_uppercase(),
             path: PathAndQueryMatcher::Unified(path.into()),
-            headers: Vec::new(),
+            headers: HeaderMap::<Matcher>::default(),
             body: Matcher::Any,
             response: Response::default(),
             hits: 0,
@@ -200,10 +202,8 @@ impl Mock {
     ///   .match_header("authorization", "password");
     /// ```
     ///
-    pub fn match_header<M: Into<Matcher>>(mut self, field: &str, value: M) -> Self {
-        self.inner
-            .headers
-            .push((field.to_owned().to_lowercase(), value.into()));
+    pub fn match_header<T: IntoHeaderName, M: Into<Matcher>>(mut self, field: T, value: M) -> Self {
+        self.inner.headers.append(field, value.into());
 
         self
     }
@@ -283,11 +283,8 @@ impl Mock {
     /// s.mock("GET", "/").with_header("content-type", "application/json");
     /// ```
     ///
-    pub fn with_header(mut self, field: &str, value: &str) -> Self {
-        self.inner
-            .response
-            .headers
-            .push((field.to_owned(), value.to_owned()));
+    pub fn with_header<T: IntoHeaderName>(mut self, field: T, value: &str) -> Self {
+        self.inner.response.headers.append(field, value.to_owned());
 
         self
     }

diff --git a/src/request.rs b/src/request.rs
@@ -1,6 +1,8 @@
 use crate::{Error, ErrorKind};
 use hyper::body;
 use hyper::body::Buf;
+use hyper::header::AsHeaderName;
+use hyper::header::HeaderValue;
 use hyper::Body as HyperBody;
 use hyper::Request as HyperRequest;
 
@@ -41,17 +43,12 @@ impl Request {
     }
 
     /// Retrieves all the header values for the given header field name
-    pub fn header(&self, header_name: &str) -> Vec<&str> {
-        self.inner
-            .headers()
-            .get_all(header_name)
-            .iter()
-            .map(|item| item.to_str().unwrap())
-            .collect::<Vec<&str>>()
+    pub fn header<T: AsHeaderName>(&self, header_name: T) -> Vec<&HeaderValue> {
+        self.inner.headers().get_all(header_name).iter().collect()
     }
 
     /// Checks whether the provided header field exists
-    pub fn has_header(&self, header_name: &str) -> bool {
+    pub fn has_header<T: AsHeaderName>(&self, header_name: T) -> bool {
         self.inner.headers().contains_key(header_name)
     }
 

diff --git a/src/response.rs b/src/response.rs
@@ -1,6 +1,7 @@
 use crate::error::Error;
 use crate::Request;
 use futures::stream::Stream;
+use hyper::HeaderMap;
 use hyper::StatusCode;
 use std::fmt;
 use std::io;
@@ -12,7 +13,7 @@ use tokio::sync::mpsc;
 #[derive(Clone, Debug, PartialEq)]
 pub(crate) struct Response {
     pub status: StatusCode,
-    pub headers: Vec<(String, String)>,
+    pub headers: HeaderMap<String>,
     pub body: Body,
 }
 
@@ -55,9 +56,11 @@ impl PartialEq for Body {
 
 impl Default for Response {
     fn default() -> Self {
+        let mut headers = HeaderMap::with_capacity(1);
+        headers.insert("connection", "close".parse().unwrap());
         Self {
             status: StatusCode::OK,
-            headers: vec![("connection".into(), "close".into())],
+            headers,
             body: Body::Bytes(Vec::new()),
         }
     }

diff --git a/src/server.rs b/src/server.rs
@@ -44,7 +44,7 @@ impl RemoteMock {
         self.inner
             .headers
             .iter()
-            .all(|&(ref field, ref expected)| expected.matches_values(&request.header(field)))
+            .all(|(field, expected)| expected.matches_values(&request.header(field)))
     }
 
     fn body_matches(&self, request: &mut Request) -> bool {

diff --git a/tests/lib.rs b/tests/lib.rs
@@ -1,6 +1,7 @@
 #[macro_use]
 extern crate serde_json;
 
+use hyper::header::HeaderName;
 use mockito::{Matcher, Server};
 use rand::distributions::Alphanumeric;
 use rand::Rng;
@@ -699,7 +700,7 @@ fn test_mock_preserves_header_order() {
 
     // Add a large number of headers so getting the same order accidentally is unlikely.
     for i in 0..100 {
-        let field = format!("x-custom-header-{}", i);
+        let field: HeaderName = format!("x-custom-header-{}", i).try_into().unwrap();
         let value = "test";
         mock = mock.with_header(&field, value);
         expected_headers.push(format!("{}: {}", field, value));