use base64::{Engine as _, engine::general_purpose::STANDARD}; use futures_util::stream; 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, ) -> Result, Status> { Ok(Response::new(echo::EchoResponse { message: request.into_inner().message, })) } type ServerEchoStream = std::pin::Pin< Box> + Send>, >; async fn server_echo( &self, request: Request, ) -> Result, 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, ) -> Result, 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> + Send>, >; async fn server_echo( &self, _request: Request, ) -> Result, 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 fn descriptor_set_b64() -> String { STANDARD.encode(echo::FILE_DESCRIPTOR_SET) }