Merge pull request #224 from songokas/status-code-from-request

Allow returning a different status code based on a request

Author: lipanski

src/mock.rs

@@ -1,6 +1,6 @@
 use crate::diff;
 use crate::matcher::{Matcher, PathAndQueryMatcher, RequestMatcher};
-use crate::response::{Body, Header, Response};
+use crate::response::{Body, Header, Response, ResponseStatusCode};
 use crate::server::RemoteMock;
 use crate::server::State;
 use crate::Request;
@@ -346,10 +346,45 @@ impl Mock {
     ///
     #[track_caller]
     pub fn with_status(mut self, status: usize) -> Self {
-        self.inner.response.status = StatusCode::from_u16(status as u16)
-            .map_err(|_| Error::new_with_context(ErrorKind::InvalidStatusCode, status))
-            .unwrap();
+        self.inner.response.status = ResponseStatusCode::StatusCode(
+            StatusCode::from_u16(status as u16)
+                .map_err(|_| Error::new_with_context(ErrorKind::InvalidStatusCode, status))
+                .unwrap(),
+        );
+        self
+    }
 
+    ///
+    /// Sets the status code of the mock response dynamically while exposing the request object.
+    ///
+    /// You can use this method to provide custom status code for every incoming request.
+    ///
+    /// The function must be thread-safe. If it's a closure, it can't be borrowing its context.
+    /// Use `move` closures and `Arc` to share any data.
+    ///
+    /// ### Example
+    ///
+    /// ```
+    /// let mut s = mockito::Server::new();
+    ///
+    /// let _m = s.mock("GET", mockito::Matcher::Any).with_status_code_from_request(|request| {
+    ///     if request.path() == "/bob" {
+    ///         500
+    ///     } else if request.path() == "/alice" {
+    ///         400
+    ///     } else {
+    ///         404
+    ///     }
+    /// });
+    /// ```
+    ///
+    #[track_caller]
+    pub fn with_status_code_from_request(
+        mut self,
+        callback: impl Fn(&Request) -> usize + Send + Sync + 'static,
+    ) -> Self {
+        self.inner.response.status =
+            ResponseStatusCode::FnWithRequest(Arc::new(move |req| callback(req)));
         self
     }
 

src/response.rs

@@ -12,11 +12,44 @@ use tokio::sync::mpsc;
 
 #[derive(Clone, Debug, PartialEq)]
 pub(crate) struct Response {
-    pub status: StatusCode,
+    pub status: ResponseStatusCode,
     pub headers: HeaderMap<Header>,
     pub body: Body,
 }
 
+#[derive(Clone)]
+pub(crate) enum ResponseStatusCode {
+    StatusCode(StatusCode),
+    FnWithRequest(Arc<StatusCodeFnWithRequest>),
+}
+
+impl fmt::Debug for ResponseStatusCode {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            ResponseStatusCode::StatusCode(ref s) => s.fmt(f),
+            ResponseStatusCode::FnWithRequest(_) => f.write_str("<callback>"),
+        }
+    }
+}
+
+impl PartialEq for ResponseStatusCode {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (ResponseStatusCode::StatusCode(ref a), ResponseStatusCode::StatusCode(ref b)) => {
+                a == b
+            }
+            (
+                ResponseStatusCode::FnWithRequest(ref a),
+                ResponseStatusCode::FnWithRequest(ref b),
+            ) => std::ptr::eq(
+                a.as_ref() as *const StatusCodeFnWithRequest as *const u8,
+                b.as_ref() as *const StatusCodeFnWithRequest as *const u8,
+            ),
+            _ => false,
+        }
+    }
+}
+
 #[derive(Clone)]
 pub(crate) enum Header {
     String(String),
@@ -45,6 +78,7 @@ impl PartialEq for Header {
     }
 }
 
+type StatusCodeFnWithRequest = dyn Fn(&Request) -> usize + Send + Sync;
 type HeaderFnWithRequest = dyn Fn(&Request) -> String + Send + Sync;
 
 type BodyFnWithWriter = dyn Fn(&mut dyn io::Write) -> io::Result<()> + Send + Sync + 'static;
@@ -89,7 +123,7 @@ impl Default for Response {
         let mut headers = HeaderMap::with_capacity(1);
         headers.insert("connection", Header::String("close".to_string()));
         Self {
-            status: StatusCode::OK,
+            status: ResponseStatusCode::StatusCode(StatusCode::OK),
             headers,
             body: Body::Bytes(Bytes::new()),
         }

src/server.rs

@@ -1,6 +1,6 @@
 use crate::mock::InnerMock;
 use crate::request::Request;
-use crate::response::{Body as ResponseBody, ChunkedStream, Header};
+use crate::response::{Body as ResponseBody, ChunkedStream, Header, ResponseStatusCode};
 use crate::ServerGuard;
 use crate::{Error, ErrorKind, Matcher, Mock};
 use bytes::Bytes;
@@ -561,7 +561,15 @@ async fn handle_request(
 }
 
 fn respond_with_mock(request: Request, mock: &RemoteMock) -> Result<Response<Body>, Error> {
-    let status: StatusCode = mock.inner.response.status;
+    let status: StatusCode = match &mock.inner.response.status {
+        ResponseStatusCode::StatusCode(c) => *c,
+        ResponseStatusCode::FnWithRequest(status_code_fn) => {
+            let status = status_code_fn(&request);
+            StatusCode::from_u16(status as u16)
+                .map_err(|_| Error::new_with_context(ErrorKind::InvalidStatusCode, status))
+                .unwrap()
+        }
+    };
     let mut response = Response::builder().status(status);
 
     for (name, value) in mock.inner.response.headers.iter() {

tests/lib.rs

@@ -587,6 +587,20 @@ fn test_mock_with_custom_status() {
     assert_eq!("HTTP/1.1 499 <none>\r\n", status_line);
 }
 
+#[test]
+fn test_mock_with_status_code_from_request() {
+    let mut s = Server::new();
+    s.mock("GET", Matcher::Any)
+        .with_status_code_from_request(|request| if request.path() == "/world" { 500 } else { 404 })
+        .create();
+
+    let (status_line, _, _) = request(s.host_with_port(), "GET /world", "");
+    assert_eq!("HTTP/1.1 500 Internal Server Error\r\n", status_line);
+
+    let (status_line, _, _) = request(s.host_with_port(), "GET /", "");
+    assert_eq!("HTTP/1.1 404 Not Found\r\n", status_line);
+}
+
 #[test]
 fn test_mock_with_body() {
     let mut s = Server::new();