crank/crates/crank-adapter-grpc/src/test_support.rs

use base64::{Engine as _, engine::general_purpose::STANDARD};
use futures_util::stream;
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};
use tokio::net::TcpListener;
use tonic::{Request, Response, Status, transport::Server};

pub mod echo {
    tonic::include_proto!("echo");
    pub const FILE_DESCRIPTOR_SET: &[u8] = tonic::include_file_descriptor_set!("echo_descriptor");
}

#[derive(Default)]
pub struct EchoServiceImpl;

#[tonic::async_trait]
impl echo::echo_service_server::EchoService for EchoServiceImpl {
    async fn unary_echo(
        &self,
        request: Request<echo::EchoRequest>,
    ) -> Result<Response<echo::EchoResponse>, Status> {
        Ok(Response::new(echo::EchoResponse {
            message: request.into_inner().message,
        }))
    }

    type ServerEchoStream = std::pin::Pin<
        Box<dyn futures_util::Stream<Item = Result<echo::EchoResponse, Status>> + Send>,
    >;

    async fn server_echo(
        &self,
        request: Request<echo::EchoRequest>,
    ) -> Result<Response<Self::ServerEchoStream>, Status> {
        let message = request.into_inner().message;
        let events = vec![
            Ok(echo::EchoResponse {
                message: format!("{message}-one"),
            }),
            Ok(echo::EchoResponse {
                message: format!("{message}-two"),
            }),
            Ok(echo::EchoResponse {
                message: format!("{message}-three"),
            }),
        ];

        Ok(Response::new(Box::pin(stream::iter(events))))
    }
}

pub async fn spawn_unary_echo_server() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let incoming = tonic::transport::server::TcpIncoming::from(listener);

    tokio::spawn(async move {
        Server::builder()
            .add_service(echo::echo_service_server::EchoServiceServer::new(
                EchoServiceImpl,
            ))
            .serve_with_incoming(incoming)
            .await
            .unwrap();
    });

    format!("http://{}", address)
}

pub async fn spawn_metadata_echo_server() -> String {
    #[derive(Default)]
    struct MetadataEchoServiceImpl;

    #[tonic::async_trait]
    impl echo::echo_service_server::EchoService for MetadataEchoServiceImpl {
        async fn unary_echo(
            &self,
            request: Request<echo::EchoRequest>,
        ) -> Result<Response<echo::EchoResponse>, Status> {
            let message = request.get_ref().message.clone();
            let request_id = request
                .metadata()
                .get("x-request-id")
                .and_then(|value| value.to_str().ok())
                .unwrap_or_default();
            let correlation_id = request
                .metadata()
                .get("x-correlation-id")
                .and_then(|value| value.to_str().ok())
                .unwrap_or_default();

            Ok(Response::new(echo::EchoResponse {
                message: format!("{}|{}|{}", message, request_id, correlation_id),
            }))
        }

        type ServerEchoStream = std::pin::Pin<
            Box<dyn futures_util::Stream<Item = Result<echo::EchoResponse, Status>> + Send>,
        >;

        async fn server_echo(
            &self,
            _request: Request<echo::EchoRequest>,
        ) -> Result<Response<Self::ServerEchoStream>, Status> {
            Ok(Response::new(Box::pin(stream::empty())))
        }
    }

    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let incoming = tonic::transport::server::TcpIncoming::from(listener);

    tokio::spawn(async move {
        Server::builder()
            .add_service(echo::echo_service_server::EchoServiceServer::new(
                MetadataEchoServiceImpl,
            ))
            .serve_with_incoming(incoming)
            .await
            .unwrap();
    });

    format!("http://{}", address)
}

pub async fn spawn_counting_unary_echo_server(request_count: Arc<AtomicUsize>) -> String {
    struct CountingEchoServiceImpl {
        request_count: Arc<AtomicUsize>,
    }

    #[tonic::async_trait]
    impl echo::echo_service_server::EchoService for CountingEchoServiceImpl {
        async fn unary_echo(
            &self,
            request: Request<echo::EchoRequest>,
        ) -> Result<Response<echo::EchoResponse>, Status> {
            let message = request.into_inner().message;
            let count = self.request_count.fetch_add(1, Ordering::SeqCst) + 1;

            Ok(Response::new(echo::EchoResponse {
                message: format!("{message}-{count}"),
            }))
        }

        type ServerEchoStream = std::pin::Pin<
            Box<dyn futures_util::Stream<Item = Result<echo::EchoResponse, Status>> + Send>,
        >;

        async fn server_echo(
            &self,
            _request: Request<echo::EchoRequest>,
        ) -> Result<Response<Self::ServerEchoStream>, Status> {
            Ok(Response::new(Box::pin(stream::empty())))
        }
    }

    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let incoming = tonic::transport::server::TcpIncoming::from(listener);

    tokio::spawn(async move {
        Server::builder()
            .add_service(echo::echo_service_server::EchoServiceServer::new(
                CountingEchoServiceImpl { request_count },
            ))
            .serve_with_incoming(incoming)
            .await
            .unwrap();
    });

    format!("http://{}", address)
}

pub fn descriptor_set_b64() -> String {
    STANDARD.encode(echo::FILE_DESCRIPTOR_SET)
}