mocktail 0.3.2

HTTP & gRPC server mocking for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! Mock gRPC service
use std::{convert::Infallible, sync::Arc};

use bytes::{Bytes, BytesMut};
use futures::{future::BoxFuture, StreamExt};
use http::{HeaderMap, HeaderValue};
use http_body::Frame;
use http_body_util::{BodyExt, StreamBody};
use hyper::{body::Incoming, service::Service};
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tracing::debug;

use crate::{request::Request, server::MockServerState, service::http::empty, Code};

use super::http::BoxBody;

/// Mock gRPC service.
#[derive(Debug, Clone)]
pub struct GrpcMockService {
    state: Arc<MockServerState>,
}

impl GrpcMockService {
    pub fn new(state: Arc<MockServerState>) -> Self {
        Self { state }
    }
}

impl Service<http::Request<Incoming>> for GrpcMockService {
    type Response = http::Response<BoxBody>;
    type Error = Infallible;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn call(&self, req: http::Request<Incoming>) -> Self::Future {
        let state = self.state.clone();
        let fut = async move {
            debug!(?req, "handling request");

            if req.method() != http::Method::POST {
                return Ok(http::Response::builder()
                    .status(http::StatusCode::METHOD_NOT_ALLOWED)
                    .header("Allow", "POST")
                    .body(empty())
                    .unwrap());
            }
            let content_type = req.headers().get("content-type");
            if !content_type.is_some_and(|v| {
                v.to_str()
                    .unwrap_or_default()
                    .starts_with("application/grpc")
            }) {
                return Ok(http::Response::builder()
                    .status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
                    .header("Accept-Post", "application/grpc")
                    .body(empty())
                    .unwrap());
            }

            let (parts, body) = req.into_parts();
            let mut stream = body.into_data_stream();

            // Create response stream
            let (response_tx, response_rx) =
                mpsc::channel::<Result<Frame<Bytes>, hyper::Error>>(32);
            let response_stream = ReceiverStream::new(response_rx);
            let response_body = BoxBody::new(StreamBody::new(response_stream));
            let response = http::Response::builder()
                .header("content-type", "application/grpc")
                .body(response_body)
                .unwrap();

            // Spawn task to handle request
            tokio::spawn(async move {
                let mut request = Request::from_parts(parts);
                let mut matched = false;
                let mut buf = BytesMut::new();

                while let Some(Ok(chunk)) = stream.next().await {
                    debug!(?chunk, "received chunk");
                    // Add chunk to body buffer
                    buf.extend(chunk);

                    // Match request to mock
                    request = request.with_body(buf.clone().freeze());
                    let mock = state.mocks().match_by_request(&request);
                    if let Some(mock) = mock {
                        matched = true;
                        debug!("mock found, sending response");
                        let mut response = mock.response;
                        // Send data frames
                        if !response.body().is_empty() {
                            while let Some(chunk) = response.body.next().await {
                                let _ = response_tx.send(Ok(Frame::data(chunk))).await;
                            }
                        }
                        // Send trailers frame
                        let mut trailers = HeaderMap::from(response.headers().clone());
                        trailers
                            .insert("grpc-status", response.status().as_grpc().to_header_value());
                        if let Some(message) = response.message() {
                            trailers
                                .insert("grpc-message", HeaderValue::from_str(message).unwrap());
                        }
                        let _ = response_tx.send(Ok(Frame::trailers(trailers))).await;
                        // Clear body buffer
                        buf.clear();
                    }
                }
                debug!("request stream closed");
                if !matched {
                    debug!(?request, "no mocks found, sending error");
                    let _ = response_tx
                        .send(Ok(Frame::trailers(mock_not_found_trailer())))
                        .await;
                }
            });

            Ok(response)
        };
        Box::pin(fut)
    }
}

fn mock_not_found_trailer() -> HeaderMap {
    let mut headers = HeaderMap::new();
    headers.insert("grpc-status", Code::NotFound.to_header_value());
    headers.insert("grpc-message", HeaderValue::from_static("mock not found"));
    headers
}