diff --git a/Cargo.toml b/Cargo.toml index c49675ce7f..d67b97496a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -418,7 +418,7 @@ incremental = false [workspace.dependencies] async-trait = { default-features = false, version = "0.1.85" } base64 = { default-features = false, version = "0.23", features = ["std"] } -bytes = { default-features = false, version = "1.10", features = ["serde"] } +bytes = { default-features = false, version = "1.10.1", features = ["serde"] } cargo_metadata = { default-features = false, version = "0.23" } chrono = { default-features = false, version = "0.4.44" } crates_io_api = { default-features = false, version = "0.12" } diff --git a/src/gax-internal/src/grpc/grpc_rust.rs b/src/gax-internal/src/grpc/grpc_rust.rs index 454195cce6..e0e7ab9470 100644 --- a/src/gax-internal/src/grpc/grpc_rust.rs +++ b/src/gax-internal/src/grpc/grpc_rust.rs @@ -82,11 +82,6 @@ impl GrpcRustClient { Self::build(config, default_endpoint, Some(instrumentation)).await } - // TODO(#5991): Temporary helper for testing. Remove once `bidi_stream` is implemented. - pub fn invoker(&self) -> &Channel { - &self.inner.invoker - } - pub async fn execute( &self, _extensions: Extensions, @@ -423,7 +418,6 @@ fn make_tls_credentials() -> ClientBuilderResult> }) } -// TODO(#5991): Add integration tests for `GrpcRustClient::bidi_stream` and `GrpcRustClient::bidi_stream_with_status` (covering happy paths and request stream failures). #[cfg(test)] mod tests { use super::*; diff --git a/src/gax-internal/src/grpc/grpc_rust/bidi.rs b/src/gax-internal/src/grpc/grpc_rust/bidi.rs index 2c70880d71..8e6480ae4c 100644 --- a/src/gax-internal/src/grpc/grpc_rust/bidi.rs +++ b/src/gax-internal/src/grpc/grpc_rust/bidi.rs @@ -229,7 +229,6 @@ where } } -// TODO(#5991): Add tests for GrpcRustStreaming in an upcoming PR. #[cfg(test)] mod tests { use super::*; @@ -237,6 +236,7 @@ mod tests { use grpc::client::{RecvStream, ResponseStreamItem, SendOptions, SendStream}; use grpc::core::{RecvMessage, ResponseHeaders, SendMessage, Trailers}; use grpc::metadata::MetadataValue; + use grpc::{StatusCodeError, StatusError}; use pretty_assertions::assert_eq; use std::sync::{Arc, Mutex}; @@ -246,136 +246,204 @@ mod tests { value: String, } - // TODO(#5991): Add tests for failure paths. - #[tokio::test] - async fn bidi_call_yields_response_messages() -> anyhow::Result<()> { - // Arrange - const METHOD_NAME: &str = "/google.test.v1.Test/Bidi"; - const HEADER_KEY: &str = "x-response-header"; - const HEADER_VALUE: &str = "response-value"; - const REQUEST_VALUE: &str = "request"; - const RESPONSE_VALUE: &str = "response"; + #[derive(Clone, Default)] + struct MockSendStream { + observed_messages: Arc>>, + notify: Option>, + } - struct TestInvoker { - observed_headers: Arc>>, + impl MockSendStream { + fn capturing( observed_messages: Arc>>, notify: Arc, + ) -> Self { + Self { + observed_messages, + notify: Some(notify), + } } - impl Invoke for TestInvoker { - type SendStream = TestSendStream; - type RecvStream = TestRecvStream; - - async fn invoke( - &self, - headers: RequestHeaders, - _options: CallOptions, - ) -> (Self::SendStream, Self::RecvStream) { - *self.observed_headers.lock().expect("lock observed headers") = Some(headers); - ( - TestSendStream { - observed_messages: self.observed_messages.clone(), - notify: self.notify.clone(), - }, - TestRecvStream { - observed_messages: self.observed_messages.clone(), - notify: self.notify.clone(), - state: StreamState::default(), - }, - ) + fn noop() -> Self { + Self::default() + } + } + + impl SendStream for MockSendStream { + async fn send( + &mut self, + message: &dyn SendMessage, + _options: SendOptions, + ) -> Result<(), ()> { + let mut encoded = message.encode().map_err(|_| ())?; + let decoded = TestMessage::decode(&mut encoded).map_err(|_| ())?; + self.observed_messages + .lock() + .expect("lock observed messages") + .push(decoded); + if let Some(notify) = &self.notify { + notify.notify_one(); } + Ok(()) } + } - struct TestSendStream { + enum MockRecvAction { + WaitForMessage { observed_messages: Arc>>, notify: Arc, - } + }, + Headers(ResponseHeaders), + Message(TestMessage), + Trailers(Trailers), + } - impl SendStream for TestSendStream { - async fn send( - &mut self, - message: &dyn SendMessage, - _options: SendOptions, - ) -> Result<(), ()> { - let mut encoded = message.encode().map_err(|_| ())?; - let decoded = TestMessage::decode(&mut encoded).map_err(|_| ())?; - self.observed_messages - .lock() - .expect("lock observed messages") - .push(decoded); - self.notify.notify_one(); - Ok(()) + struct MockRecvStream { + actions: std::collections::VecDeque, + } + + impl MockRecvStream { + fn new(actions: impl IntoIterator) -> Self { + Self { + actions: actions.into_iter().collect(), } } - // TODO(#5991): Refactor common stream state test mocks across grpc_rust tests. - #[derive(Default)] - enum StreamState { - #[default] - Initial, - HeadersSent, - MessageSent, - Done, + fn immediate_trailers(trailers: Trailers) -> Self { + Self::new([MockRecvAction::Trailers(trailers)]) } - /// A mock [`RecvStream`] that simulates a gRPC response stream sequence: - /// - /// 1. Waits until at least one request message is sent by the client, then returns response headers. - /// 2. Returns a response message. - /// 3. Returns stream trailers followed by stream closure. - struct TestRecvStream { - observed_messages: Arc>>, - notify: Arc, - state: StreamState, + fn headers_and_trailers(headers: ResponseHeaders, trailers: Trailers) -> Self { + Self::new([ + MockRecvAction::Headers(headers), + MockRecvAction::Trailers(trailers), + ]) } + } - impl RecvStream for TestRecvStream { - async fn recv(&mut self, message: &mut dyn RecvMessage) -> ResponseStreamItem { - match self.state { - StreamState::Initial => { - self.state = StreamState::HeadersSent; - // Wait for the client request message to be sent before yielding initial headers. - while self - .observed_messages + impl RecvStream for MockRecvStream { + async fn recv(&mut self, message: &mut dyn RecvMessage) -> ResponseStreamItem { + while let Some(action) = self.actions.pop_front() { + match action { + MockRecvAction::WaitForMessage { + observed_messages, + notify, + } => { + while observed_messages .lock() - .expect("lock messages") + .expect("lock observed messages") .is_empty() { - self.notify.notified().await; + notify.notified().await; } - let mut metadata = grpc::metadata::MetadataMap::new(); - metadata.insert(HEADER_KEY, MetadataValue::from_static(HEADER_VALUE)); - ResponseStreamItem::Headers(ResponseHeaders::new().with_metadata(metadata)) } - StreamState::HeadersSent => { - self.state = StreamState::MessageSent; - // Emit a mock response message. - let response = TestMessage { - value: RESPONSE_VALUE.to_string(), - }; + MockRecvAction::Headers(headers) => { + return ResponseStreamItem::Headers(headers); + } + MockRecvAction::Message(response) => { let mut encoded = Bytes::from(response.encode_to_vec()); message .decode(&mut encoded) .expect("decode response message"); - ResponseStreamItem::Message + return ResponseStreamItem::Message; } - StreamState::MessageSent => { - self.state = StreamState::Done; - ResponseStreamItem::Trailers(Trailers::new(Ok(()))) + MockRecvAction::Trailers(trailers) => { + return ResponseStreamItem::Trailers(trailers); } - StreamState::Done => ResponseStreamItem::StreamClosed, } } + ResponseStreamItem::StreamClosed } + } + + // TODO(#5991): Refactor common stream test mocks into shared test_helpers across grpc_rust tests. + struct MockInvoker { + send_stream: Mutex>, + recv_stream: Mutex>, + observed_headers: Arc>>, + } + + impl MockInvoker { + fn new(send_stream: S, recv_stream: R) -> Self { + Self { + send_stream: Mutex::new(Some(send_stream)), + recv_stream: Mutex::new(Some(recv_stream)), + observed_headers: Arc::new(Mutex::new(None)), + } + } + + fn with_observed_headers( + send_stream: S, + recv_stream: R, + observed_headers: Arc>>, + ) -> Self { + Self { + send_stream: Mutex::new(Some(send_stream)), + recv_stream: Mutex::new(Some(recv_stream)), + observed_headers, + } + } + } + + impl Invoke for MockInvoker + where + S: SendStream + Send + 'static, + R: RecvStream + Send + 'static, + { + type SendStream = S; + type RecvStream = R; + + async fn invoke( + &self, + headers: RequestHeaders, + _options: CallOptions, + ) -> (Self::SendStream, Self::RecvStream) { + *self.observed_headers.lock().expect("lock observed headers") = Some(headers); + ( + self.send_stream + .lock() + .expect("lock send stream") + .take() + .expect("send stream should only be invoked once"), + self.recv_stream + .lock() + .expect("lock recv stream") + .take() + .expect("recv stream should only be invoked once"), + ) + } + } + + #[tokio::test] + async fn bidi_call_yields_response_messages() -> anyhow::Result<()> { + // Arrange + const METHOD_NAME: &str = "/google.test.v1.Test/Bidi"; + const HEADER_KEY: &str = "x-response-header"; + const HEADER_VALUE: &str = "response-value"; + const REQUEST_VALUE: &str = "request"; + const RESPONSE_VALUE: &str = "response"; let observed_headers = Arc::new(Mutex::new(None)); let observed_messages = Arc::new(Mutex::new(Vec::new())); let notify = Arc::new(tokio::sync::Notify::new()); - let invoker = TestInvoker { - observed_headers: observed_headers.clone(), - observed_messages: observed_messages.clone(), - notify, - }; + + let mut metadata = grpc::metadata::MetadataMap::new(); + metadata.insert(HEADER_KEY, MetadataValue::from_static(HEADER_VALUE)); + + let invoker = MockInvoker::with_observed_headers( + MockSendStream::capturing(observed_messages.clone(), notify.clone()), + MockRecvStream::new([ + MockRecvAction::WaitForMessage { + observed_messages: observed_messages.clone(), + notify, + }, + MockRecvAction::Headers(ResponseHeaders::new().with_metadata(metadata)), + MockRecvAction::Message(TestMessage { + value: RESPONSE_VALUE.to_string(), + }), + MockRecvAction::Trailers(Trailers::new(Ok(()))), + ]), + observed_headers.clone(), + ); let headers = RequestHeaders::new().with_method_name(METHOD_NAME); let request = TestMessage { value: REQUEST_VALUE.to_string(), @@ -420,4 +488,107 @@ mod tests { ); Ok(()) } + + #[tokio::test] + async fn bidi_call_yields_error_on_server_error_status() -> anyhow::Result<()> { + // Arrange + const METHOD_NAME: &str = "/google.test.v1.Test/Bidi"; + const ERROR_MESSAGE: &str = "stream aborted"; + + let err = StatusError::new(StatusCodeError::Aborted, ERROR_MESSAGE); + let invoker = MockInvoker::new( + MockSendStream::noop(), + MockRecvStream::headers_and_trailers(ResponseHeaders::new(), Trailers::new(Err(err))), + ); + let headers = RequestHeaders::new().with_method_name(METHOD_NAME); + + // Act + let response = + invoke_bidi::(&invoker, headers, tokio_stream::empty()) + .await?; + + // Assert + let mut stream = response.into_inner(); + let err = stream + .message() + .await + .expect_err("should return status error from trailers"); + assert_eq!(err.code(), tonic::Code::Aborted); + assert_eq!(err.message(), ERROR_MESSAGE); + assert_eq!(stream.message().await?, None); + + Ok(()) + } + + #[tokio::test] + async fn bidi_call_returns_error_on_immediate_trailers_only_status() -> anyhow::Result<()> { + // Arrange + const METHOD_NAME: &str = "/google.test.v1.Test/Bidi"; + const ERROR_MESSAGE: &str = "immediate failure"; + + let err = StatusError::new(StatusCodeError::Aborted, ERROR_MESSAGE); + let invoker = MockInvoker::new( + MockSendStream::noop(), + MockRecvStream::immediate_trailers(Trailers::new(Err(err))), + ); + let headers = RequestHeaders::new().with_method_name(METHOD_NAME); + + // Act + let err = + invoke_bidi::(&invoker, headers, tokio_stream::empty()) + .await + .expect_err("invoke_bidi should fail immediately on trailers-only error"); + + // Assert + assert_eq!(err.code(), tonic::Code::Aborted); + assert_eq!(err.message(), ERROR_MESSAGE); + + Ok(()) + } + + #[tokio::test] + async fn bidi_call_immediately_fails_if_initial_request_fails() -> anyhow::Result<()> { + // Arrange + const METHOD_NAME: &str = "/google.test.v1.Test/Bidi"; + const ERROR_MESSAGE_STREAM_CLOSED: &str = "grpc-rust request stream closed"; + + struct FailingSendStream; + impl SendStream for FailingSendStream { + async fn send( + &mut self, + _message: &dyn SendMessage, + _options: SendOptions, + ) -> Result<(), ()> { + Err(()) + } + } + + struct PendingRecvStream; + impl RecvStream for PendingRecvStream { + async fn recv(&mut self, _message: &mut dyn RecvMessage) -> ResponseStreamItem { + std::future::pending().await + } + } + + let invoker = MockInvoker::new(FailingSendStream, PendingRecvStream); + let headers = RequestHeaders::new().with_method_name(METHOD_NAME); + let request = TestMessage { + value: "msg".to_string(), + }; + + // Act + let err = invoke_bidi::( + &invoker, + headers, + tokio_stream::iter([request]), + ) + .await + .expect_err("invoke_bidi should fail immediately when initial outbound send fails"); + + // Assert + assert_eq!(err.code(), tonic::Code::Internal); + assert_eq!(err.message(), ERROR_MESSAGE_STREAM_CLOSED); + + Ok(()) + } } diff --git a/src/gax-internal/tests/grpc_rust_client_streaming.rs b/src/gax-internal/tests/grpc_rust_client_streaming.rs index 1f3452b090..63ab4ebf82 100644 --- a/src/gax-internal/tests/grpc_rust_client_streaming.rs +++ b/src/gax-internal/tests/grpc_rust_client_streaming.rs @@ -21,11 +21,10 @@ mod tests { // TODO(#5991): Consider refactoring some tests to run against both `grpc::Client` and `grpc::GrpcRustClient`. use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + use google_cloud_gax::options::RequestOptions; use google_cloud_gax_internal::grpc::GrpcRustClient; - use google_cloud_gax_internal::grpc::grpc_rust::bidi::{GrpcRustSend, ReceiveTask, RecvItem}; + use google_cloud_gax_internal::grpc::grpc_rust::GrpcRustStreaming; use google_cloud_gax_internal::options::ClientConfig; - use grpc::client::{CallOptions, Invoke, SendOptions, SendStream}; - use grpc::core::RequestHeaders; use grpc_server::google::test::v1::{EchoRequest, EchoResponse}; use grpc_server::start_echo_server; use pretty_assertions::assert_eq; @@ -34,177 +33,322 @@ mod tests { const MSG2: &str = "msg2"; #[tokio::test] - async fn test_grpc_rust_client_streaming_response_pump() -> anyhow::Result<()> { + async fn test_bidi_stream() -> anyhow::Result<()> { // Arrange - let (_client, mut send_stream, mut rx, _pump_task, _server_task) = - start_client_stream().await?; - - // Consume and assert initial headers emitted upon RPC start - let first_item = rx - .recv() - .await - .expect("response pump should yield at least one item")?; + let mut session = start_bidi_stream().await?; assert!( - matches!(first_item, Some(RecvItem::Headers(_))), - "expected initial headers" + !session.metadata.is_empty(), + "expected initial metadata headers from response" ); // Act - send_echo_request(&mut send_stream, MSG1).await?; + send_echo_request(&session.tx, MSG1).await?; // Assert - let res1 = recv_echo_response(&mut rx).await?; + let res1 = recv_echo_response(&mut session.stream).await?; assert_eq!(res1.message, MSG1); // Act - send_echo_request(&mut send_stream, MSG2).await?; + send_echo_request(&session.tx, MSG2).await?; // Assert - let res2 = recv_echo_response(&mut rx).await?; + let res2 = recv_echo_response(&mut session.stream).await?; assert_eq!(res2.message, MSG2); // Act - drop(send_stream); + drop(session.tx); // Assert - let end_res = rx - .recv() - .await - .expect("response pump should yield stream termination")?; - assert!( - end_res.is_none(), - "response pump should yield stream termination" - ); + let end_res = session.stream.message().await?; + assert_eq!(end_res, None, "stream should yield None upon completion"); Ok(()) } #[tokio::test] - async fn test_grpc_rust_client_streaming_stream_remains_usable_even_after_cancellation() - -> anyhow::Result<()> { + async fn test_bidi_stream_remains_usable_even_after_cancellation() -> anyhow::Result<()> { // Arrange - let (_client, mut send_stream, mut rx, _pump_task, _server_task) = - start_client_stream().await?; - consume_initial_headers(&mut rx).await?; + let mut session = start_bidi_stream().await?; // Act // Attempt to receive the next item with a short timeout before any message is - // sent, thereby cancelling rx.recv() - let recv_result = - tokio::time::timeout(std::time::Duration::from_millis(50), rx.recv()).await; + // sent, thereby cancelling recv_echo_response + let cancelled_read = tokio::time::timeout( + std::time::Duration::from_millis(50), + recv_echo_response(&mut session.stream), + ) + .await; // Assert assert!( - recv_result.is_err(), - "recv should time out when no message is sent" + cancelled_read.is_err(), + "cancelled_read should time out when no request is sent" ); // Act - // Send request message after read cancellation - send_echo_request(&mut send_stream, MSG1).await?; + send_echo_request(&session.tx, MSG1).await?; + + // Assert + let res = recv_echo_response(&mut session.stream).await?; + assert_eq!(res.message, MSG1); + + Ok(()) + } + + #[tokio::test] + async fn test_bidi_stream_drop_closes_channel() -> anyhow::Result<()> { + // Arrange + let mut session = start_bidi_stream().await?; + + // Act + send_echo_request(&session.tx, MSG1).await?; // Assert - // Receiver should still successfully yield the sent message - let res = recv_echo_response(&mut rx).await?; + let res = recv_echo_response(&mut session.stream).await?; assert_eq!(res.message, MSG1); + // Act + drop(session.stream); + + // Assert + tokio::time::timeout(std::time::Duration::from_secs(5), session.tx.closed()) + .await + .expect("dropping the stream should close the request channel"); + + Ok(()) + } + + #[tokio::test] + async fn test_bidi_stream_server_error_mid_stream() -> anyhow::Result<()> { + // Arrange + let mut session = start_bidi_stream().await?; + + // Act + send_echo_request(&session.tx, MSG1).await?; + + // Assert + let res1 = recv_echo_response(&mut session.stream).await?; + assert_eq!(res1.message, MSG1); + + // Act + // Sending an empty message causes our test echo server to return InvalidArgument and close the stream + send_echo_request(&session.tx, "").await?; + + // Assert + let err = session + .stream + .message() + .await + .expect_err("stream should return status error when server fails mid-stream"); + assert_eq!(err.code(), tonic::Code::InvalidArgument); + assert!( + err.message().contains("empty message"), + "expected 'empty message' in error message, got '{}'", + err.message() + ); + + // Assert + let subsequent = session.stream.message().await?; + assert_eq!( + subsequent, None, + "subsequent calls after stream termination should yield None" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_bidi_stream_initial_error() -> anyhow::Result<()> { + // Act + let res = start_bidi_stream_with_params("resource=error").await; + + // Assert + let err = res + .expect_err("bidi_stream should fail when server returns an initial error") + .downcast::() + .expect("expected a google_cloud_gax::error::Error"); + + let status = err.status().expect("expected status"); + assert_eq!(status.code, google_cloud_gax::error::rpc::Code::Aborted); + assert_eq!(status.message, "test with initial error"); + Ok(()) } #[tokio::test] - async fn test_grpc_rust_client_streaming_cancel_pump_task() -> anyhow::Result<()> { + async fn test_bidi_stream_with_status() -> anyhow::Result<()> { // Arrange - let (_client, mut send_stream, mut rx, pump_task, _server_task) = - start_client_stream().await?; - consume_initial_headers(&mut rx).await?; + let mut session = start_bidi_stream_with_status() + .await? + .expect("should succeed"); + assert!( + !session.metadata.is_empty(), + "expected initial metadata headers from response" + ); // Act - send_echo_request(&mut send_stream, MSG1).await?; + send_echo_request(&session.tx, MSG1).await?; // Assert - let res1 = recv_echo_response(&mut rx).await?; + let res1 = recv_echo_response(&mut session.stream).await?; assert_eq!(res1.message, MSG1); // Act - drop(pump_task); + send_echo_request(&session.tx, MSG2).await?; + + // Assert + let res2 = recv_echo_response(&mut session.stream).await?; + assert_eq!(res2.message, MSG2); + + // Act + drop(session.tx); + + // Assert + let end_res = session.stream.message().await?; + assert_eq!(end_res, None, "stream should yield None upon completion"); + + Ok(()) + } + + #[tokio::test] + async fn test_bidi_stream_with_status_initial_error() -> anyhow::Result<()> { + // Act + let result = start_bidi_stream_with_status_params("resource=error").await?; // Assert - let end_res = rx.recv().await; + let err = result.expect_err("should return Err(status) on initial error"); + assert_eq!(err.code(), tonic::Code::Aborted); assert!( - end_res.is_none(), - "channel should return None after pump_task is dropped" + err.message().contains("test with initial error"), + "expected 'test with initial error' in error message, got '{}'", + err.message() ); Ok(()) } - /// Starts an echo server and initializes a streaming RPC for testing. - async fn start_client_stream() -> anyhow::Result<( - GrpcRustClient, - impl SendStream, - tokio::sync::mpsc::Receiver>>>, - ReceiveTask, - tokio::task::JoinHandle<()>, - )> { + struct TestBidiSession { + /// For sending outbound request messages. + tx: tokio::sync::mpsc::Sender, + /// For reading inbound response messages. + stream: GrpcRustStreaming, + /// Initial server response metadata. + metadata: tonic::metadata::MetadataMap, + _client: GrpcRustClient, + /// Task handle for the echo server. + _server_task: tokio::task::JoinHandle<()>, + } + + impl std::fmt::Debug for TestBidiSession { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TestBidiSession") + .field("tx", &self.tx) + .field("metadata", &self.metadata) + .finish_non_exhaustive() + } + } + + /// Starts an echo server and initializes a bidirectional streaming RPC using `bidi_stream`. + async fn start_bidi_stream() -> anyhow::Result { + start_bidi_stream_with_params("").await + } + + /// Starts an echo server and initializes a bidirectional streaming RPC using `bidi_stream` with custom request parameters. + async fn start_bidi_stream_with_params( + request_params: &str, + ) -> anyhow::Result { let (endpoint, server_task) = start_echo_server().await?; let mut config = ClientConfig::default(); config.cred = Some(Anonymous::new().build()); let client = GrpcRustClient::new(config, &endpoint).await?; - // Start the RPC - let headers = RequestHeaders::new().with_method_name("/google.test.v1.EchoService/Chat"); - let (send_stream, recv_stream) = client - .invoker() - .invoke(headers, CallOptions::default()) - .await; + let (tx, rx) = tokio::sync::mpsc::channel::(10); + let request_stream = tokio_stream::wrappers::ReceiverStream::new(rx); - // Start the response pump - let (rx, pump_task) = ReceiveTask::start::(recv_stream); + let response = client + .bidi_stream::( + tonic::Extensions::new(), + http::uri::PathAndQuery::from_static("/google.test.v1.EchoService/Chat"), + request_stream, + RequestOptions::default(), + "test-only-api-client/1.0", + request_params, + ) + .await?; + + let (metadata, stream, _) = response.into_parts(); + Ok(TestBidiSession { + tx, + stream, + metadata, + _client: client, + _server_task: server_task, + }) + } - Ok((client, send_stream, rx, pump_task, server_task)) + /// Starts an echo server and initializes a bidirectional streaming RPC using `bidi_stream_with_status`. + async fn start_bidi_stream_with_status() -> anyhow::Result> { + start_bidi_stream_with_status_params("").await } - /// Sends the given message to the echo server - async fn send_echo_request(send_stream: &mut impl SendStream, msg: &str) -> anyhow::Result<()> { - send_stream - .send( - &GrpcRustSend(EchoRequest { - message: msg.to_string(), - ..Default::default() - }), - SendOptions::default(), + /// Starts an echo server and initializes a bidirectional streaming RPC using `bidi_stream_with_status` with custom request parameters. + async fn start_bidi_stream_with_status_params( + request_params: &str, + ) -> anyhow::Result> { + let (endpoint, server_task) = start_echo_server().await?; + let mut config = ClientConfig::default(); + config.cred = Some(Anonymous::new().build()); + let client = GrpcRustClient::new(config, &endpoint).await?; + + let (tx, rx) = tokio::sync::mpsc::channel::(10); + let request_stream = tokio_stream::wrappers::ReceiverStream::new(rx); + + let result = client + .bidi_stream_with_status::( + tonic::Extensions::new(), + http::uri::PathAndQuery::from_static("/google.test.v1.EchoService/Chat"), + request_stream, + RequestOptions::default(), + "test-only-api-client/1.0", + request_params, ) - .await - .map_err(|_| anyhow::anyhow!("failed to send message '{msg}'")) + .await?; + + match result { + Ok(response) => { + let (metadata, stream, _) = response.into_parts(); + Ok(Ok(TestBidiSession { + tx, + stream, + metadata, + _client: client, + _server_task: server_task, + })) + } + Err(status) => Ok(Err(status)), + } } - /// Consumes initial headers emitted by the response pump upon RPC start. - async fn consume_initial_headers( - rx: &mut tokio::sync::mpsc::Receiver>>>, + /// Sends an echo request with the given message string. + async fn send_echo_request( + tx: &tokio::sync::mpsc::Sender, + msg: &str, ) -> anyhow::Result<()> { - let item = rx - .recv() - .await - .expect("response pump should yield initial headers")?; - assert!( - matches!(item, Some(RecvItem::Headers(_))), - "expected initial headers, got {item:?}" - ); - Ok(()) + tx.send(EchoRequest { + message: msg.to_string(), + ..Default::default() + }) + .await + .map_err(|_| anyhow::anyhow!("failed to send message '{msg}'")) } - /// Receives a response message from the response pump. Fails if headers are received. + /// Receives the next echo response from the stream. async fn recv_echo_response( - rx: &mut tokio::sync::mpsc::Receiver>>>, + stream: &mut GrpcRustStreaming, ) -> anyhow::Result { - let res = rx - .recv() - .await - .expect("response pump should yield a response item")? - .expect("expected a response item"); - let RecvItem::Message(msg) = res else { - panic!("expected response message, got {res:?}"); - }; - Ok(msg) + stream + .message() + .await? + .ok_or_else(|| anyhow::anyhow!("expected response message, got end of stream")) } }