Skip to content

test(gax): add tests for bidi_stream and bidi_stream_with_status #6313

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
joshuatants wants to merge 1 commit into
base: grpc_rust
Choose a base branch
from
+358 −142
Draft

test(gax): add tests for bidi_stream and bidi_stream_with_status #6313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions src/gax-internal/src/grpc/grpc_rust.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Request, Response>(
&self,
_extensions: Extensions,
Expand Down Expand Up @@ -423,7 +418,6 @@ fn make_tls_credentials() -> ClientBuilderResult<Arc<RustlsChannelCredentials>>
})
}

// TODO(#5991): Add integration tests for `GrpcRustClient::bidi_stream` and `GrpcRustClient::bidi_stream_with_status` (covering happy paths and request stream failures).

@vsharonlynn vsharonlynn Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the unit test cases do not cover the case of request stream failure? Do we need to test this case?

#[cfg(test)]
mod tests {
use super::*;
Expand Down
137 changes: 103 additions & 34 deletions src/gax-internal/src/grpc/grpc_rust/bidi.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,14 +229,14 @@ where
}
}

// TODO(#5991): Add tests for GrpcRustStreaming in an upcoming PR.
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
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};

Expand All @@ -246,7 +246,38 @@ mod tests {
value: String,
}

// TODO(#5991): Add tests for failure paths.
struct TestSendStream {
observed_messages: Arc<Mutex<Vec<TestMessage>>>,
notify: Arc<tokio::sync::Notify>,
}

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(())
}
}

// TODO(#5991): Refactor common stream state test mocks across grpc_rust tests.
#[derive(Default)]
enum StreamState {
#[default]
Initial,
HeadersSent,
MessageSent,
Done,
}

#[tokio::test]
async fn bidi_call_yields_response_messages() -> anyhow::Result<()> {
// Arrange
Expand Down Expand Up @@ -286,38 +317,6 @@ mod tests {
}
}

struct TestSendStream {
observed_messages: Arc<Mutex<Vec<TestMessage>>>,
notify: Arc<tokio::sync::Notify>,
}

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(())
}
}

// TODO(#5991): Refactor common stream state test mocks across grpc_rust tests.
#[derive(Default)]
enum StreamState {
#[default]
Initial,
HeadersSent,
MessageSent,
Done,
}

/// 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.
Expand Down Expand Up @@ -420,4 +419,74 @@ mod tests {
);
Ok(())
}

#[tokio::test]
async fn bidi_call_yields_error_on_server_error_status() -> anyhow::Result<()> {

@vsharonlynn vsharonlynn Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Quick check of my understanding:
IIUC in this test case, the error is only observed when stream.message().await is called.

So this doesn't test the immediate error check in invoke_bidi on line 214.

Do you want to test the case where invoke_bidi returns an error immediately?

// Arrange
const METHOD_NAME: &str = "/google.test.v1.Test/Bidi";
const ERROR_MESSAGE: &str = "stream aborted";

struct TestErrorInvoker;

impl Invoke for TestErrorInvoker {
type SendStream = TestSendStream;
type RecvStream = TestErrorRecvStream;

async fn invoke(
&self,
_headers: RequestHeaders,
_options: CallOptions,
) -> (Self::SendStream, Self::RecvStream) {
(
TestSendStream {
observed_messages: Arc::new(Mutex::new(Vec::new())),
notify: Arc::new(tokio::sync::Notify::new()),
},
TestErrorRecvStream {
state: StreamState::default(),
},
)
}
}

struct TestErrorRecvStream {
state: StreamState,
}

impl RecvStream for TestErrorRecvStream {
async fn recv(&mut self, _message: &mut dyn RecvMessage) -> ResponseStreamItem {
match self.state {
StreamState::Initial => {
self.state = StreamState::HeadersSent;
ResponseStreamItem::Headers(ResponseHeaders::new())
}
_ => {
self.state = StreamState::Done;
let err = StatusError::new(StatusCodeError::Aborted, ERROR_MESSAGE);
ResponseStreamItem::Trailers(Trailers::new(Err(err)))
}
}
}
}

let invoker = TestErrorInvoker;
let headers = RequestHeaders::new().with_method_name(METHOD_NAME);

// Act
let response =
invoke_bidi::<TestMessage, TestMessage, _>(&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(())
}
}
Loading
Loading