Expose server mocks and enable mutability (#5)

Signed-off-by: declark1 <[email protected]>

Author: declark1

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mocktail-test"
-version = "0.1.2-alpha"
+version = "0.1.3-alpha"
 edition = "2021"
 authors = ["Dan Clark", "Gaurav Kumbhat"]
 

mocktail-test/Cargo.toml

@@ -32,6 +32,7 @@ mod tests {
         server.start().await?;
 
         let client = reqwest::Client::new();
+
         let response = client
             .post(server.url("/hello"))
             .json(&HelloRequest { name: "Dan".into() })
@@ -41,7 +42,6 @@ mod tests {
         let body = response.json::<HelloResponse>().await?;
         dbg!(&body);
 
-        let client = reqwest::Client::new();
         let response = client
             .post(server.url("/hello"))
             .json(&HelloRequest {
@@ -51,6 +51,35 @@ mod tests {
             .await?;
         assert!(response.status() == StatusCode::NOT_FOUND);
 
+        // Clear the mocks on the server
+        server.mocks().clear();
+
+        assert!(server.mocks().is_empty());
+
+        // Add a new mock to the server
+        server.mocks().insert(
+            MockPath::post("/hello"),
+            Mock::new(
+                MockRequest::json(HelloRequest {
+                    name: "There".into(),
+                }),
+                MockResponse::json(HelloResponse {
+                    message: "Hello There!".into(),
+                }),
+            ),
+        );
+
+        let response = client
+            .post(server.url("/hello"))
+            .json(&HelloRequest {
+                name: "There".into(),
+            })
+            .send()
+            .await?;
+        assert!(response.status().is_success());
+        let body = response.json::<HelloResponse>().await?;
+        dbg!(&body);
+
         Ok(())
     }
 }

mocktail-test/src/http_hello.rs

@@ -1,6 +1,6 @@
 [package]
 name = "mocktail"
-version = "0.1.2-alpha"
+version = "0.1.3-alpha"
 edition = "2021"
 authors = ["Dan Clark", "Gaurav Kumbhat"]
 description = "HTTP & gRPC server mocking for Rust"
@@ -17,6 +17,7 @@ categories = ["development-tools", "development-tools::testing"]
 doctest = false
 
 [dependencies]
+async-trait = "0.1.87"
 bytes = "1.9.0"
 futures = "0.3.31"
 h2 = { version = "0.4.8", features = ["stream"] }

mocktail/Cargo.toml

@@ -7,7 +7,7 @@ pub mod prelude {
 
     pub use crate::{
         mock::{Mock, MockBody, MockPath, MockRequest, MockResponse, MockSet},
-        server::{GrpcMockServer, HttpMockServer},
+        server::{GrpcMockServer, HttpMockServer, MockServer},
         utils::prost::MessageExt as _,
         Error,
     };

mocktail/src/lib.rs

@@ -1,18 +1,27 @@
 use std::collections::{hash_map, HashMap};
 
 use super::Mock;
-use crate::utils::HeaderMapExt;
 
 /// A set of mocks for a service.
 #[derive(Default, Debug, Clone)]
 pub struct MockSet(HashMap<MockPath, Vec<Mock>>);
 
 impl MockSet {
-    /// Creates a empty [`MockSet`].
+    /// Creates an empty mockset.
     pub fn new() -> Self {
         Self::default()
     }
 
+    /// Returns the number of entries in the mockset.
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Returns true if the mockset contains no entries.
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
     /// Inserts a [`Mock`].
     pub fn insert(&mut self, path: MockPath, mock: Mock) {
         match self.0.entry(path) {
@@ -25,6 +34,21 @@ impl MockSet {
         }
     }
 
+    /// Gets an entry for in-place manipulation.
+    pub fn entry(&mut self, path: MockPath) -> hash_map::Entry<'_, MockPath, Vec<Mock>> {
+        self.0.entry(path)
+    }
+
+    /// Removes an entry from the mockset.
+    pub fn remove(&mut self, path: &MockPath) -> Option<Vec<Mock>> {
+        self.0.remove(path)
+    }
+
+    /// Clears the mockset.
+    pub fn clear(&mut self) {
+        self.0.clear()
+    }
+
     /// Matches a [`Mock`] by path and predicate.
     pub fn find<P>(&self, path: &MockPath, predicate: P) -> Option<&Mock>
     where
@@ -39,22 +63,6 @@ impl MockSet {
     pub fn match_by_body(&self, path: &MockPath, body: &[u8]) -> Option<&Mock> {
         self.find(path, |mock| mock.request.body() == body)
     }
-
-    /// Matches a [`Mock`] by path, body, and headers.
-    pub fn match_by_body_and_headers(
-        &self,
-        path: &MockPath,
-        body: &[u8],
-        headers: &http::HeaderMap,
-    ) -> Option<&Mock> {
-        // `headers` must be a superset of `mock.request.headers`,
-        if let Some(mock) = self.match_by_body(path, body) {
-            if headers.is_superset(&mock.request.headers) {
-                return Some(mock);
-            }
-        }
-        None
-    }
 }
 
 impl FromIterator<(MockPath, Vec<Mock>)> for MockSet {

mocktail/src/mock/set.rs

@@ -1,4 +1,12 @@
-use crate::mock::MockSet;
+use std::{
+    net::SocketAddr,
+    sync::{RwLock, RwLockWriteGuard},
+};
+
+use async_trait::async_trait;
+use url::Url;
+
+use crate::{mock::MockSet, Error};
 
 mod grpc;
 pub use grpc::GrpcMockServer;
@@ -9,11 +17,29 @@ pub use http::HttpMockServer;
 /// Server state.
 #[derive(Debug)]
 pub struct ServerState {
-    pub mocks: MockSet,
+    pub mocks: RwLock<MockSet>,
 }
 
 impl ServerState {
     pub fn new(mocks: MockSet) -> Self {
-        Self { mocks }
+        Self {
+            mocks: RwLock::new(mocks),
+        }
     }
 }
+
+#[async_trait]
+pub trait MockServer {
+    async fn start(&self) -> Result<(), Error>;
+
+    /// Returns the server's service name.
+    fn name(&self) -> &str;
+
+    /// Returns the server's address.
+    fn addr(&self) -> SocketAddr;
+
+    /// Returns the url for a path.
+    fn url(&self, path: &str) -> Url;
+
+    fn mocks(&self) -> RwLockWriteGuard<'_, MockSet>;
+}

mocktail/src/server.rs

@@ -1,5 +1,11 @@
-use std::{convert::Infallible, net::SocketAddr, sync::Arc, time::Duration};
+use std::{
+    convert::Infallible,
+    net::SocketAddr,
+    sync::{Arc, RwLockWriteGuard},
+    time::Duration,
+};
 
+use async_trait::async_trait;
 use bytes::{Bytes, BytesMut};
 use futures::{future::BoxFuture, StreamExt};
 use http::{HeaderMap, HeaderValue, Request, Response};
@@ -12,7 +18,7 @@ use tokio_stream::wrappers::ReceiverStream;
 use tracing::{debug, error, info};
 use url::Url;
 
-use super::ServerState;
+use super::{MockServer, ServerState};
 use crate::{
     mock::{MockPath, MockSet, TonicBoxBody},
     utils::{find_available_port, has_content_type, tonic::CodeExt},
@@ -23,6 +29,7 @@ use crate::{
 pub struct GrpcMockServer {
     name: &'static str,
     addr: SocketAddr,
+    base_url: Url,
     state: Arc<ServerState>,
 }
 
@@ -31,14 +38,20 @@ impl GrpcMockServer {
     pub fn new(name: &'static str, mocks: MockSet) -> Result<Self, Error> {
         let port = find_available_port().unwrap();
         let addr: SocketAddr = format!("0.0.0.0:{}", port).parse().unwrap();
+        let base_url = Url::parse(&format!("http://{}", &addr)).unwrap();
+        let state = Arc::new(ServerState::new(mocks));
         Ok(Self {
             name,
             addr,
-            state: Arc::new(ServerState::new(mocks)),
+            base_url,
+            state,
         })
     }
+}
 
-    pub async fn start(&self) -> Result<(), Error> {
+#[async_trait]
+impl MockServer for GrpcMockServer {
+    async fn start(&self) -> Result<(), Error> {
         let service = GrpcMockSvc {
             state: self.state.clone(),
         };
@@ -63,28 +76,26 @@ impl GrpcMockServer {
             }
         });
 
-        // Cushion for server to become ready, there is probably a better approach :)
-        tokio::time::sleep(Duration::from_secs(1)).await;
+        // Give the server time to become ready
+        tokio::time::sleep(Duration::from_millis(10)).await;
 
         Ok(())
     }
 
-    /// Returns the server's service name.
-    pub fn name(&self) -> &str {
+    fn name(&self) -> &str {
         self.name
     }
 
-    /// Returns the server's address.
-    pub fn addr(&self) -> SocketAddr {
+    fn addr(&self) -> SocketAddr {
         self.addr
     }
 
-    pub fn base_url(&self) -> Url {
-        Url::parse(&format!("http://{}", self.addr())).unwrap()
+    fn url(&self, path: &str) -> Url {
+        self.base_url.join(path).unwrap()
     }
 
-    pub fn url(&self, path: &str) -> Url {
-        self.base_url().join(path).unwrap()
+    fn mocks(&self) -> RwLockWriteGuard<'_, MockSet> {
+        self.state.mocks.write().unwrap()
     }
 }
 
@@ -140,7 +151,13 @@ impl Service<Request<Incoming>> for GrpcMockSvc {
                     buf.extend(chunk);
                     // Attempt to match buffered data to mock
                     let body = buf.clone().freeze();
-                    if let Some(mock) = state.mocks.match_by_body(&path, &body) {
+                    let mock = state
+                        .mocks
+                        .read()
+                        .unwrap()
+                        .match_by_body(&path, &body)
+                        .cloned();
+                    if let Some(mock) = mock {
                         matched = true;
                         // A matching mock has been found, send response
                         debug!("mock found, sending response");
@@ -166,7 +183,6 @@ impl Service<Request<Incoming>> for GrpcMockSvc {
                 debug!("request stream closed");
                 if !matched {
                     debug!("no mocks found, sending error");
-                    dbg!(&state.mocks);
                     let mut trailers = HeaderMap::new();
                     trailers.insert("grpc-status", (tonic::Code::NotFound as i32).into());
                     trailers.insert("grpc-message", HeaderValue::from_static("mock not found"));