1 use super::*; 2 use bytes::Bytes; 3 use http_body::Body; 4 use pin_project::pin_project; 5 use std::{ 6 pin::Pin, 7 sync::{ 8 atomic::{AtomicUsize, Ordering::SeqCst}, 9 Arc, 10 }, 11 task::{ready, Context, Poll}, 12 }; 13 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; 14 use tonic::codec::CompressionEncoding; 15 use tonic::transport::{server::Connected, Channel}; 16 use tower_http::map_request_body::MapRequestBodyLayer; 17 18 macro_rules! parametrized_tests { 19 ($fn_name:ident, $($test_name:ident: $input:expr),+ $(,)?) => { 20 paste::paste! { 21 $( 22 #[tokio::test(flavor = "multi_thread")] 23 async fn [<$fn_name _ $test_name>]() { 24 let input = $input; 25 $fn_name(input).await; 26 } 27 )+ 28 } 29 } 30 } 31 32 pub(crate) use parametrized_tests; 33 34 /// A body that tracks how many bytes passes through it 35 #[pin_project] 36 pub struct CountBytesBody<B> { 37 #[pin] 38 pub inner: B, 39 pub counter: Arc<AtomicUsize>, 40 } 41 42 impl<B> Body for CountBytesBody<B> 43 where 44 B: Body<Data = Bytes>, 45 { 46 type Data = B::Data; 47 type Error = B::Error; 48 49 fn poll_data( 50 self: Pin<&mut Self>, 51 cx: &mut Context<'_>, 52 ) -> Poll<Option<Result<Self::Data, Self::Error>>> { 53 let this = self.project(); 54 let counter: Arc<AtomicUsize> = this.counter.clone(); 55 match ready!(this.inner.poll_data(cx)) { 56 Some(Ok(chunk)) => { 57 println!("response body chunk size = {}", chunk.len()); 58 counter.fetch_add(chunk.len(), SeqCst); 59 Poll::Ready(Some(Ok(chunk))) 60 } 61 x => Poll::Ready(x), 62 } 63 } 64 65 fn poll_trailers( 66 self: Pin<&mut Self>, 67 cx: &mut Context<'_>, 68 ) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> { 69 self.project().inner.poll_trailers(cx) 70 } 71 72 fn is_end_stream(&self) -> bool { 73 self.inner.is_end_stream() 74 } 75 76 fn size_hint(&self) -> http_body::SizeHint { 77 self.inner.size_hint() 78 } 79 } 80 81 #[allow(dead_code)] 82 pub fn measure_request_body_size_layer( 83 bytes_sent_counter: Arc<AtomicUsize>, 84 ) -> MapRequestBodyLayer<impl Fn(hyper::Body) -> hyper::Body + Clone> { 85 MapRequestBodyLayer::new(move |mut body: hyper::Body| { 86 let (mut tx, new_body) = hyper::Body::channel(); 87 88 let bytes_sent_counter = bytes_sent_counter.clone(); 89 tokio::spawn(async move { 90 while let Some(chunk) = body.data().await { 91 let chunk = chunk.unwrap(); 92 println!("request body chunk size = {}", chunk.len()); 93 bytes_sent_counter.fetch_add(chunk.len(), SeqCst); 94 tx.send_data(chunk).await.unwrap(); 95 } 96 97 if let Some(trailers) = body.trailers().await.unwrap() { 98 tx.send_trailers(trailers).await.unwrap(); 99 } 100 }); 101 102 new_body 103 }) 104 } 105 106 #[allow(dead_code)] 107 pub async fn mock_io_channel(client: tokio::io::DuplexStream) -> Channel { 108 let mut client = Some(client); 109 110 Endpoint::try_from("http://[::]:50051") 111 .unwrap() 112 .connect_with_connector(service_fn(move |_: Uri| { 113 let client = client.take().unwrap(); 114 async move { Ok::<_, std::io::Error>(client) } 115 })) 116 .await 117 .unwrap() 118 } 119 120 #[derive(Clone)] 121 pub struct AssertRightEncoding { 122 encoding: CompressionEncoding, 123 } 124 125 #[allow(dead_code)] 126 impl AssertRightEncoding { 127 pub fn new(encoding: CompressionEncoding) -> Self { 128 Self { encoding } 129 } 130 131 pub fn call<B: Body>(self, req: http::Request<B>) -> http::Request<B> { 132 let expected = match self.encoding { 133 CompressionEncoding::Gzip => "gzip", 134 CompressionEncoding::Zstd => "zstd", 135 _ => panic!("unexpected encoding {:?}", self.encoding), 136 }; 137 assert_eq!(req.headers().get("grpc-encoding").unwrap(), expected); 138 139 req 140 } 141 } 142