Add Mock::with_body_from_request

src/lib.rs

@@ -642,6 +642,7 @@ use lazy_static::lazy_static;
 pub use legacy::{mock, reset, server_address, server_url};
 pub use matcher::Matcher;
 pub use mock::Mock;
+pub use request::Request;
 pub use server::{Server, ServerGuard};
 use tokio::runtime::Runtime;
 

src/mock.rs

@@ -3,6 +3,7 @@ use crate::diff;
 use crate::matcher::{Matcher, PathAndQueryMatcher};
 use crate::response::{Body, Response};
 use crate::server::RemoteMock;
+use crate::Request;
 use crate::{Error, ErrorKind};
 use hyper::StatusCode;
 use rand::distributions::Alphanumeric;
@@ -323,9 +324,38 @@ impl Mock {
     ///
     pub fn with_body_from_fn(
         mut self,
-        cb: impl Fn(&mut dyn io::Write) -> io::Result<()> + Send + Sync + 'static,
+        callback: impl Fn(&mut dyn io::Write) -> io::Result<()> + Send + Sync + 'static,
     ) -> Self {
-        self.inner.response.body = Body::Fn(Arc::new(cb));
+        self.inner.response.body = Body::Fn(Arc::new(callback));
+        self
+    }
+
+    ///
+    /// Sets the body of the mock response dynamically while exposing the request object.
+    ///
+    /// You can use this method to provide a custom reponse body for every incoming request.
+    ///
+    /// ### Example
+    ///
+    /// ```
+    /// let mut s = mockito::Server::new();
+    ///
+    /// let _m = s.mock("GET", mockito::Matcher::Any).with_body_from_request(|request| {
+    ///     if request.path() == "/bob" {
+    ///         "hello bob".into()
+    ///     } else if request.path() == "/alice" {
+    ///         "hello alice".into()
+    ///     } else {
+    ///         "hello world".into()
+    ///     }
+    /// });
+    /// ```
+    ///
+    pub fn with_body_from_request(
+        mut self,
+        callback: impl Fn(&Request) -> Vec<u8> + Send + Sync + 'static,
+    ) -> Self {
+        self.inner.response.body = Body::FnWithRequest(Arc::new(callback));
         self
     }
 

src/request.rs

@@ -4,24 +4,34 @@ use hyper::body::Buf;
 use hyper::Body as HyperBody;
 use hyper::Request as HyperRequest;
 
+///
+/// Stores a HTTP request
+///
 #[derive(Debug)]
-pub(crate) struct Request {
+pub struct Request {
     inner: HyperRequest<HyperBody>,
     body: Option<Vec<u8>>,
 }
 
 impl Request {
-    pub fn new(request: HyperRequest<HyperBody>) -> Self {
+    pub(crate) fn new(request: HyperRequest<HyperBody>) -> Self {
         Request {
             inner: request,
             body: None,
         }
     }
 
+    /// The HTTP method
     pub fn method(&self) -> &str {
         self.inner.method().as_ref()
     }
 
+    /// The path excluding the query part
+    pub fn path(&self) -> &str {
+        self.inner.uri().path()
+    }
+
+    /// The path including the query part
     pub fn path_and_query(&self) -> &str {
         self.inner
             .uri()
@@ -30,20 +40,31 @@ impl Request {
             .unwrap_or("")
     }
 
-    pub fn header(&self, field: &str) -> Vec<&str> {
+    /// 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(field)
+            .get_all(header_name)
             .iter()
             .map(|item| item.to_str().unwrap())
             .collect::<Vec<&str>>()
     }
 
+    /// Checks whether the provided header field exists
     pub fn has_header(&self, header_name: &str) -> bool {
         self.inner.headers().contains_key(header_name)
     }
 
-    pub async fn read_body(&mut self) -> &Vec<u8> {
+    /// Returns the request body or an error, if the body hasn't been read
+    /// up to this moment.
+    pub fn body(&self) -> Result<&Vec<u8>, Error> {
+        self.body
+            .as_ref()
+            .ok_or(Error::new(ErrorKind::RequestBodyFailure))
+    }
+
+    /// Reads the body (if it hasn't been read already) and returns it
+    pub(crate) async fn read_body(&mut self) -> &Vec<u8> {
         if self.body.is_none() {
             let raw_body = self.inner.body_mut();
             let mut buf = body::aggregate(raw_body)

src/response.rs

@@ -1,3 +1,4 @@
+use crate::Request;
 use core::task::Poll;
 use futures::stream::Stream;
 use hyper::StatusCode;
@@ -14,18 +15,21 @@ pub(crate) struct Response {
 }
 
 type BodyFn = dyn Fn(&mut dyn io::Write) -> io::Result<()> + Send + Sync + 'static;
+type BodyFnWithRequest = dyn Fn(&Request) -> Vec<u8> + Send + Sync + 'static;
 
 #[derive(Clone)]
 pub(crate) enum Body {
     Bytes(Vec<u8>),
     Fn(Arc<BodyFn>),
+    FnWithRequest(Arc<BodyFnWithRequest>),
 }
 
 impl fmt::Debug for Body {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match *self {
             Body::Bytes(ref b) => b.fmt(f),
             Body::Fn(_) => f.write_str("<callback>"),
+            Body::FnWithRequest(_) => f.write_str("<callback>"),
         }
     }
 }
@@ -38,6 +42,10 @@ impl PartialEq for Body {
                 a.as_ref() as *const BodyFn as *const u8,
                 b.as_ref() as *const BodyFn as *const u8,
             ),
+            (Body::FnWithRequest(ref a), Body::FnWithRequest(ref b)) => std::ptr::eq(
+                a.as_ref() as *const BodyFnWithRequest as *const u8,
+                b.as_ref() as *const BodyFnWithRequest as *const u8,
+            ),
             _ => false,
         }
     }

src/server.rs

@@ -374,7 +374,10 @@ async fn handle_request(
     }
 }
 
-async fn respond_with_mock(request: Request, mock: &RemoteMock) -> Result<Response<Body>, Error> {
+async fn respond_with_mock(
+    mut request: Request,
+    mock: &RemoteMock,
+) -> Result<Response<Body>, Error> {
     let status: StatusCode = mock.inner.response.status;
     let mut response = Response::builder().status(status);
 
@@ -399,6 +402,14 @@ async fn respond_with_mock(request: Request, mock: &RemoteMock) -> Result<Respon
 
                 Body::wrap_stream(chunked)
             }
+            ResponseBody::FnWithRequest(body_fn) => {
+                // Make sure to read the request body so that `Request::body` can
+                // return it, if needed
+                request.read_body().await;
+
+                let bytes = body_fn(&request);
+                Body::from(bytes)
+            }
         }
     } else {
         Body::empty()

tests/lib.rs

@@ -635,6 +635,49 @@ fn test_mock_with_fn_body() {
     assert_eq!("hello", body);
 }
 
+#[test]
+fn test_mock_with_body_from_request() {
+    let mut s = Server::new();
+    let _m = s
+        .mock("GET", Matcher::Any)
+        .with_body_from_request(|request| {
+            if request.path() == "/world" {
+                "hello world".into()
+            } else {
+                "just hello".into()
+            }
+        })
+        .create();
+
+    let (_, _, body) = request(&s.host_with_port(), "GET /world", "");
+    assert_eq!("hello world", body);
+
+    let (_, _, body) = request(&s.host_with_port(), "GET /", "");
+    assert_eq!("just hello", body);
+}
+
+#[test]
+fn test_mock_with_body_from_request_body() {
+    let mut s = Server::new();
+    let _m = s
+        .mock("GET", "/")
+        .with_body_from_request(|request| {
+            let body = std::str::from_utf8(request.body().unwrap()).unwrap();
+            if body == "test" {
+                "test".into()
+            } else {
+                "not a test".into()
+            }
+        })
+        .create();
+
+    let (_, _, body) = request_with_body(&s.host_with_port(), "GET /", "", "test");
+    assert_eq!("test", body);
+
+    let (_, _, body) = request_with_body(&s.host_with_port(), "GET /", "", "something else");
+    assert_eq!("not a test", body);
+}
+
 #[test]
 fn test_mock_with_header() {
     let mut s = Server::new();