tonic_server_mock/
lib.rs

1use std::{convert::Infallible, sync::Arc};
2use tokio::{io::DuplexStream, sync::mpsc::Sender};
3use tonic::{
4    body::BoxBody,
5    codegen::{
6        http::{Request, Response},
7        Service,
8    },
9    server::NamedService,
10    transport::{Channel, Error as TransportError},
11};
12
13pub trait Svc:
14    Service<Request<BoxBody>, Response = Response<BoxBody>, Error = Infallible>
15    + NamedService
16    + Clone
17    + Send
18    + 'static
19where
20    Self::Future: Send + 'static,
21{
22}
23
24impl<S> Svc for S
25where
26    S: Service<Request<BoxBody>, Response = Response<BoxBody>, Error = Infallible>
27        + NamedService
28        + Clone
29        + Send
30        + 'static,
31    S::Future: Send + 'static,
32{
33}
34
35pub type ConnectionMockSender = Sender<Result<DuplexStream, TransportError>>;
36
37#[derive(Debug, Clone)]
38pub struct EndpointMock {
39    connection_sender: Arc<ConnectionMockSender>,
40}
41
42impl EndpointMock {
43    pub fn new(connection_sender: ConnectionMockSender) -> Self {
44        Self {
45            connection_sender: Arc::new(connection_sender),
46        }
47    }
48
49    pub async fn once(self) -> Channel {
50        self.connect().await
51    }
52
53    pub async fn connect(&self) -> Channel {
54        let connection_sender = Arc::clone(&self.connection_sender);
55
56        let client_connector =
57            ::tower::service_fn(move |/* mut */ uri: ::tonic::transport::Uri| {
58                tracing::info!("connection to {:?}", uri);
59                let connection_sender = Arc::clone(&connection_sender);
60                async move {
61                    let (client_io, server_io) = ::tokio::io::duplex(1024);
62                    connection_sender.send(Ok(server_io)).await.unwrap();
63                    Ok::<_, ::tonic::transport::Error>(::hyper_util::rt::TokioIo::new(client_io))
64                }
65            });
66
67        ::tonic::transport::Endpoint::try_from("http://[::1]:50051/pseudo-endpoint")
68            .unwrap()
69            .connect_with_connector(client_connector)
70            .await
71            .unwrap()
72    }
73}
74
75#[macro_export]
76macro_rules! mock_server_fn {
77    ($vis:vis $fn_name:ident; $($svc:ident),+; $logger:path) => {
78
79        $vis async fn $fn_name(
80            $($svc: impl $crate::Svc<Future: Send>),+
81        ) -> (impl ::futures::Future<Output = ()>, $crate::EndpointMock) {
82            use $logger::{info};
83
84            let (connection_sender, connections_receiver) = ::tokio::sync::mpsc::channel(32);
85            let incoming_connections: ::tokio_stream::wrappers::ReceiverStream<
86                Result<::tokio::io::DuplexStream, ::tonic::transport::Error>,
87            > = ::tokio_stream::wrappers::ReceiverStream::from(connections_receiver);
88
89            let router = ::tonic::transport::Server::builder()$(
90                .add_service($svc))+;
91
92            // Spawn server
93            let server_future = async move {
94                info!("start grpc server");
95                router
96                    .serve_with_incoming(incoming_connections)
97                    .await
98                    .unwrap();
99                info!("grpc server stopped");
100            };
101
102            let connector = $crate::EndpointMock::new(connection_sender);
103
104            (server_future, connector)
105        }
106    };
107
108    ($fn_name:ident; $($svc:ident),+; $logger:path) => {
109        mock_server_fn!(pub(crate) $fn_name; $($svc),+; $logger);
110    };
111
112    ($vis:vis $fn_name:ident; $($svc:ident),+) => {
113        mock_server_fn!($vis $fn_name; $($svc),+; ::tracing);
114    };
115
116    ($fn_name:ident; $($svc:ident),+) => {
117        mock_server_fn!(pub(crate) $fn_name; $($svc),+; ::tracing);
118    };
119}
tonic_server_mock/lib.rs

tonic_server_mock/
lib.rs
