From 1c104a1a6c68272d388db673cf628b3ce9914ffe Mon Sep 17 00:00:00 2001 From: Bradley White <14679271+devbww@users.noreply.github.com> Date: Fri, 11 Feb 2022 19:30:20 -0500 Subject: [PATCH 1/2] fix(spanner): add Options save/restore to PartialResultSetSource Save the options in force during `PartialResultSetSource` construction, and then restore them over `PartialResultSetReader` calls. This means that `NextRow()`, `Finish()` and `TryCancel()` will have the right options installed during any RPCs they execute. Tighten up the test by installing a non-matching `OptionsSpan` before any `NextRow()` calls are made, and by expecting no `TryCancel()`s in cases that complete their enumeration. This is the Spanner equivalent of #8256. --- .../internal/partial_result_set_source.cc | 2 + .../internal/partial_result_set_source.h | 4 +- .../partial_result_set_source_test.cc | 454 ++++++++++++++---- google/cloud/spanner/results.h | 2 + 4 files changed, 375 insertions(+), 87 deletions(-) diff --git a/google/cloud/spanner/internal/partial_result_set_source.cc b/google/cloud/spanner/internal/partial_result_set_source.cc index 1cbfa47fd278c..fa31d57dd6a72 100644 --- a/google/cloud/spanner/internal/partial_result_set_source.cc +++ b/google/cloud/spanner/internal/partial_result_set_source.cc @@ -84,6 +84,7 @@ PartialResultSetSource::~PartialResultSetSource() { // If there is actual data in the streaming RPC Finish() can deadlock, so // before trying to read the final status we need to cancel the streaming // RPC. + internal::OptionsSpan span(options_); reader_->TryCancel(); // The user didn't iterate over all the data; finish the stream on their // behalf, but we have no way to communicate error status. @@ -95,6 +96,7 @@ PartialResultSetSource::~PartialResultSetSource() { } Status PartialResultSetSource::ReadFromStream() { + internal::OptionsSpan span(options_); auto result_set = reader_->Read(); if (!result_set) { // Read() returns false for end of stream, whether we read all the data or diff --git a/google/cloud/spanner/internal/partial_result_set_source.h b/google/cloud/spanner/internal/partial_result_set_source.h index ec3eda2f456fd..427d8cb41b1fe 100644 --- a/google/cloud/spanner/internal/partial_result_set_source.h +++ b/google/cloud/spanner/internal/partial_result_set_source.h @@ -19,6 +19,7 @@ #include "google/cloud/spanner/results.h" #include "google/cloud/spanner/value.h" #include "google/cloud/spanner/version.h" +#include "google/cloud/options.h" #include "google/cloud/status.h" #include "google/cloud/status_or.h" #include "absl/types/optional.h" @@ -60,10 +61,11 @@ class PartialResultSetSource : public ResultSourceInterface { private: explicit PartialResultSetSource( std::unique_ptr reader) - : reader_(std::move(reader)) {} + : options_(internal::CurrentOptions()), reader_(std::move(reader)) {} Status ReadFromStream(); + Options options_; std::unique_ptr reader_; absl::optional metadata_; absl::optional stats_; diff --git a/google/cloud/spanner/internal/partial_result_set_source_test.cc b/google/cloud/spanner/internal/partial_result_set_source_test.cc index 6cb6b5e769ce1..5b3bad9b9df20 100644 --- a/google/cloud/spanner/internal/partial_result_set_source_test.cc +++ b/google/cloud/spanner/internal/partial_result_set_source_test.cc @@ -16,6 +16,7 @@ #include "google/cloud/spanner/row.h" #include "google/cloud/spanner/testing/mock_partial_result_set_reader.h" #include "google/cloud/spanner/value.h" +#include "google/cloud/options.h" #include "google/cloud/testing_util/is_proto_equal.h" #include "google/cloud/testing_util/status_matchers.h" #include "absl/memory/memory.h" @@ -42,6 +43,20 @@ using ::google::protobuf::TextFormat; using ::testing::HasSubstr; using ::testing::Return; +struct StringOption { + using Type = std::string; +}; + +// Create the `PartialResultSetSource` within a particular `OptionsSpan` so +// that we might check that all `PartialResultSetReader` calls happen within +// a matching span. +StatusOr> CreatePartialResultSetSource( + std::unique_ptr reader, + std::string const& string_option) { + internal::OptionsSpan span(Options{}.set(string_option)); + return PartialResultSetSource::Create(std::move(reader)); +} + MATCHER_P(IsValidAndEquals, expected, "Verifies that a StatusOr contains the given Row") { return arg && *arg == expected; @@ -50,12 +65,18 @@ MATCHER_P(IsValidAndEquals, expected, /// @test Verify the behavior when the initial `Read()` fails. TEST(PartialResultSetSourceTest, InitialReadFailure) { auto grpc_reader = absl::make_unique(); - EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(absl::optional{})); + EXPECT_CALL(*grpc_reader, Read()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "InitialReadFailure"); + return absl::optional{}; + }); Status finish_status(StatusCode::kInvalidArgument, "invalid"); EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(finish_status)); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "InitialReadFailure"); EXPECT_THAT(reader, StatusIs(StatusCode::kInvalidArgument)); } @@ -79,12 +100,28 @@ TEST(PartialResultSetSourceTest, ReadSuccessThenFailure) { spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response)) - .WillOnce(Return(absl::optional{})); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ReadSuccessThenFailure"); + return response; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ReadSuccessThenFailure"); + return absl::optional{}; + }); Status finish_status(StatusCode::kCancelled, "cancelled"); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(finish_status)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([&finish_status] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ReadSuccessThenFailure"); + return finish_status; + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + // The first call to NextRow() yields a row but the second gives an error. - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ReadSuccessThenFailure"); EXPECT_STATUS_OK(reader.status()); EXPECT_THAT((*reader)->NextRow(), IsValidAndEquals(MakeTestRow({{"AnInt", spanner::Value(80)}}))); @@ -96,12 +133,25 @@ TEST(PartialResultSetSourceTest, ReadSuccessThenFailure) { TEST(PartialResultSetSourceTest, MissingMetadata) { auto grpc_reader = absl::make_unique(); spanner_proto::PartialResultSet response; - EXPECT_CALL(*grpc_reader, Read()).WillOnce(Return(response)); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); + EXPECT_CALL(*grpc_reader, Read()).WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingMetadata"); + return response; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingMetadata"); + return Status(); + }); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).Times(1); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingMetadata"); + }); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = + CreatePartialResultSetSource(std::move(grpc_reader), "MissingMetadata"); EXPECT_THAT(reader, StatusIs(StatusCode::kInternal, "response contained no metadata")); } @@ -116,11 +166,26 @@ TEST(PartialResultSetSourceTest, MissingRowTypeNoData) { spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response)) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingRowTypeNoData"); + return response; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingRowTypeNoData"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingRowTypeNoData"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "MissingRowTypeNoData"); ASSERT_STATUS_OK(reader); EXPECT_THAT((*reader)->NextRow(), IsValidAndEquals(spanner::Row{})); } @@ -136,12 +201,25 @@ TEST(PartialResultSetSourceTest, MissingRowTypeWithData) { values: { string_value: "10" })pb"; spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); - EXPECT_CALL(*grpc_reader, Read()).WillOnce(Return(response)); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); + EXPECT_CALL(*grpc_reader, Read()).WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingRowTypeWithData"); + return response; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingRowTypeWithData"); + return Status(); + }); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).Times(1); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MissingRowTypeWithData"); + }); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "MissingRowTypeWithData"); ASSERT_STATUS_OK(reader); StatusOr row = (*reader)->NextRow(); EXPECT_THAT(row, StatusIs(StatusCode::kInternal, @@ -190,11 +268,25 @@ TEST(PartialResultSetSourceTest, SingleResponse) { spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response)) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "SingleResponse"); + return response; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "SingleResponse"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), "SingleResponse"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = + CreatePartialResultSetSource(std::move(grpc_reader), "SingleResponse"); EXPECT_STATUS_OK(reader.status()); // Verify the returned metadata is correct. @@ -308,15 +400,46 @@ TEST(PartialResultSetSourceTest, MultipleResponses) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])) - .WillOnce(Return(response[2])) - .WillOnce(Return(response[3])) - .WillOnce(Return(response[4])) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return response[1]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return response[2]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return response[3]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return response[4]; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "MultipleResponses"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = + CreatePartialResultSetSource(std::move(grpc_reader), "MultipleResponses"); EXPECT_STATUS_OK(reader.status()); // Verify the returned rows are correct. @@ -365,13 +488,36 @@ TEST(PartialResultSetSourceTest, ResponseWithNoValues) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])) - .WillOnce(Return(response[2])) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ResponseWithNoValues"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ResponseWithNoValues"); + return response[1]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ResponseWithNoValues"); + return response[2]; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ResponseWithNoValues"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ResponseWithNoValues"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ResponseWithNoValues"); EXPECT_STATUS_OK(reader.status()); // Verify the returned row is correct. @@ -426,15 +572,46 @@ TEST(PartialResultSetSourceTest, ChunkedStringValueWellFormed) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])) - .WillOnce(Return(response[2])) - .WillOnce(Return(response[3])) - .WillOnce(Return(response[4])) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return response[1]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return response[2]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return response[3]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return response[4]; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedStringValueWellFormed"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ChunkedStringValueWellFormed"); EXPECT_STATUS_OK(reader.status()); // Verify the returned values are correct. @@ -476,13 +653,30 @@ TEST(PartialResultSetSourceTest, ChunkedValueSetNoValue) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoValue"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoValue"); + return response[1]; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoValue"); + return Status(); + }); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).Times(1); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoValue"); + }); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ChunkedValueSetNoValue"); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -519,13 +713,30 @@ TEST(PartialResultSetSourceTest, ChunkedValueSetNoFollowingValue) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoFollowingValue"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoFollowingValue"); + return response[1]; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoFollowingValue"); + return Status(); + }); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).Times(1); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetNoFollowingValue"); + }); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ChunkedValueSetNoFollowingValue"); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -562,12 +773,31 @@ TEST(PartialResultSetSourceTest, ChunkedValueSetAtEndOfStream) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetAtEndOfStream"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetAtEndOfStream"); + return response[1]; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetAtEndOfStream"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueSetAtEndOfStream"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ChunkedValueSetAtEndOfStream"); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -607,14 +837,35 @@ TEST(PartialResultSetSourceTest, ChunkedValueMergeFailure) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])) - .WillOnce(Return(response[2])); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueMergeFailure"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueMergeFailure"); + return response[1]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueMergeFailure"); + return response[2]; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueMergeFailure"); + return Status(); + }); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).Times(1); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ChunkedValueMergeFailure"); + }); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ChunkedValueMergeFailure"); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -678,15 +929,46 @@ TEST(PartialResultSetSourceTest, ErrorOnIncompleteRow) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce(Return(response[0])) - .WillOnce(Return(response[1])) - .WillOnce(Return(response[2])) - .WillOnce(Return(response[3])) - .WillOnce(Return(response[4])) - .WillOnce(Return(absl::optional{})); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(Status())); - - auto reader = PartialResultSetSource::Create(std::move(grpc_reader)); + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return response[0]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return response[1]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return response[2]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return response[3]; + }) + .WillOnce([&response] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return response[4]; + }) + .WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return absl::optional{}; + }); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { + EXPECT_EQ(internal::CurrentOptions().get(), + "ErrorOnIncompleteRow"); + return Status(); + }); + EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); + + internal::OptionsSpan overlay(Options{}.set("uh-oh")); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader), + "ErrorOnIncompleteRow"); EXPECT_STATUS_OK(reader.status()); // Verify the first two rows are correct. diff --git a/google/cloud/spanner/results.h b/google/cloud/spanner/results.h index 29f2cc80eb92e..5549789c18420 100644 --- a/google/cloud/spanner/results.h +++ b/google/cloud/spanner/results.h @@ -29,6 +29,7 @@ namespace google { namespace cloud { namespace spanner_internal { GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + class ResultSourceInterface { public: virtual ~ResultSourceInterface() = default; @@ -249,6 +250,7 @@ class ProfileDmlResult { private: std::unique_ptr source_; }; + GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace spanner } // namespace cloud From 3e475aab1df59db0360992887bafbd1a0c15de5c Mon Sep 17 00:00:00 2001 From: Bradley White <14679271+devbww@users.noreply.github.com> Date: Sun, 13 Feb 2022 03:31:51 -0500 Subject: [PATCH 2/2] Refactor the PartialResultSetReader mocks. --- .../partial_result_set_source_test.cc | 442 +++++------------- 1 file changed, 112 insertions(+), 330 deletions(-) diff --git a/google/cloud/spanner/internal/partial_result_set_source_test.cc b/google/cloud/spanner/internal/partial_result_set_source_test.cc index 5b3bad9b9df20..8fcfaea370ef8 100644 --- a/google/cloud/spanner/internal/partial_result_set_source_test.cc +++ b/google/cloud/spanner/internal/partial_result_set_source_test.cc @@ -41,22 +41,47 @@ using ::google::cloud::testing_util::IsProtoEqual; using ::google::cloud::testing_util::StatusIs; using ::google::protobuf::TextFormat; using ::testing::HasSubstr; -using ::testing::Return; +using ::testing::UnitTest; + +std::string CurrentTestName() { + return UnitTest::GetInstance()->current_test_info()->name(); +} struct StringOption { using Type = std::string; }; -// Create the `PartialResultSetSource` within a particular `OptionsSpan` so -// that we might check that all `PartialResultSetReader` calls happen within -// a matching span. +// Create the `PartialResultSetSource` within an `OptionsSpan` that has its +// `StringOption` set to the current test name, so that we might check that +// all `PartialResultSetReader` calls happen within a matching span. StatusOr> CreatePartialResultSetSource( - std::unique_ptr reader, - std::string const& string_option) { - internal::OptionsSpan span(Options{}.set(string_option)); + std::unique_ptr reader) { + internal::OptionsSpan span(Options{}.set(CurrentTestName())); return PartialResultSetSource::Create(std::move(reader)); } +// Returns a functor that expects the current `StringOption` to match the test +// name. +std::function VoidMock() { + return [] { + EXPECT_EQ(internal::CurrentOptions().get(), + CurrentTestName()); + }; +} + +// Returns a functor that will return the argument after expecting that the +// current `StringOption` matches the test name. +template +std::function ResultMock(T const& result) { + return [result]() { + EXPECT_EQ(internal::CurrentOptions().get(), + CurrentTestName()); + return result; + }; +} + +using ReadResult = absl::optional; + MATCHER_P(IsValidAndEquals, expected, "Verifies that a StatusOr contains the given Row") { return arg && *arg == expected; @@ -65,19 +90,15 @@ MATCHER_P(IsValidAndEquals, expected, /// @test Verify the behavior when the initial `Read()` fails. TEST(PartialResultSetSourceTest, InitialReadFailure) { auto grpc_reader = absl::make_unique(); - EXPECT_CALL(*grpc_reader, Read()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "InitialReadFailure"); - return absl::optional{}; - }); - Status finish_status(StatusCode::kInvalidArgument, "invalid"); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce(Return(finish_status)); + EXPECT_CALL(*grpc_reader, Read()) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()) + .WillOnce(ResultMock(Status(StatusCode::kInvalidArgument, "invalid"))); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "InitialReadFailure"); - EXPECT_THAT(reader, StatusIs(StatusCode::kInvalidArgument)); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); + EXPECT_THAT(reader, StatusIs(StatusCode::kInvalidArgument, "invalid")); } /** @@ -100,58 +121,33 @@ TEST(PartialResultSetSourceTest, ReadSuccessThenFailure) { spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ReadSuccessThenFailure"); - return response; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ReadSuccessThenFailure"); - return absl::optional{}; - }); - Status finish_status(StatusCode::kCancelled, "cancelled"); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([&finish_status] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ReadSuccessThenFailure"); - return finish_status; - }); + .WillOnce(ResultMock(response)) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()) + .WillOnce(ResultMock(Status(StatusCode::kCancelled, "cancelled"))); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); // The first call to NextRow() yields a row but the second gives an error. internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ReadSuccessThenFailure"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); EXPECT_THAT((*reader)->NextRow(), IsValidAndEquals(MakeTestRow({{"AnInt", spanner::Value(80)}}))); auto row = (*reader)->NextRow(); - EXPECT_THAT(row, StatusIs(StatusCode::kCancelled)); + EXPECT_THAT(row, StatusIs(StatusCode::kCancelled, "cancelled")); } /// @test Verify the behavior when the first response does not contain metadata. TEST(PartialResultSetSourceTest, MissingMetadata) { auto grpc_reader = absl::make_unique(); - spanner_proto::PartialResultSet response; - EXPECT_CALL(*grpc_reader, Read()).WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingMetadata"); - return response; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingMetadata"); - return Status(); - }); + EXPECT_CALL(*grpc_reader, Read()) + .WillOnce(ResultMock(spanner_proto::PartialResultSet())); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingMetadata"); - }); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce(VoidMock()); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = - CreatePartialResultSetSource(std::move(grpc_reader), "MissingMetadata"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_THAT(reader, StatusIs(StatusCode::kInternal, "response contained no metadata")); } @@ -166,26 +162,13 @@ TEST(PartialResultSetSourceTest, MissingRowTypeNoData) { spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingRowTypeNoData"); - return response; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingRowTypeNoData"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingRowTypeNoData"); - return Status(); - }); + .WillOnce(ResultMock(response)) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "MissingRowTypeNoData"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); ASSERT_STATUS_OK(reader); EXPECT_THAT((*reader)->NextRow(), IsValidAndEquals(spanner::Row{})); } @@ -201,25 +184,13 @@ TEST(PartialResultSetSourceTest, MissingRowTypeWithData) { values: { string_value: "10" })pb"; spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); - EXPECT_CALL(*grpc_reader, Read()).WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingRowTypeWithData"); - return response; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingRowTypeWithData"); - return Status(); - }); + EXPECT_CALL(*grpc_reader, Read()).WillOnce(ResultMock(response)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MissingRowTypeWithData"); - }); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce(VoidMock()); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "MissingRowTypeWithData"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); ASSERT_STATUS_OK(reader); StatusOr row = (*reader)->NextRow(); EXPECT_THAT(row, StatusIs(StatusCode::kInternal, @@ -268,25 +239,13 @@ TEST(PartialResultSetSourceTest, SingleResponse) { spanner_proto::PartialResultSet response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "SingleResponse"); - return response; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "SingleResponse"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), "SingleResponse"); - return Status(); - }); + .WillOnce(ResultMock(response)) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = - CreatePartialResultSetSource(std::move(grpc_reader), "SingleResponse"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Verify the returned metadata is correct. @@ -400,46 +359,17 @@ TEST(PartialResultSetSourceTest, MultipleResponses) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return response[1]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return response[2]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return response[3]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return response[4]; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "MultipleResponses"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])) + .WillOnce(ResultMock(response[2])) + .WillOnce(ResultMock(response[3])) + .WillOnce(ResultMock(response[4])) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = - CreatePartialResultSetSource(std::move(grpc_reader), "MultipleResponses"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Verify the returned rows are correct. @@ -488,36 +418,15 @@ TEST(PartialResultSetSourceTest, ResponseWithNoValues) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ResponseWithNoValues"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ResponseWithNoValues"); - return response[1]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ResponseWithNoValues"); - return response[2]; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ResponseWithNoValues"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ResponseWithNoValues"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])) + .WillOnce(ResultMock(response[2])) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ResponseWithNoValues"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Verify the returned row is correct. @@ -572,46 +481,17 @@ TEST(PartialResultSetSourceTest, ChunkedStringValueWellFormed) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return response[1]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return response[2]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return response[3]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return response[4]; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedStringValueWellFormed"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])) + .WillOnce(ResultMock(response[2])) + .WillOnce(ResultMock(response[3])) + .WillOnce(ResultMock(response[4])) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ChunkedStringValueWellFormed"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Verify the returned values are correct. @@ -653,30 +533,14 @@ TEST(PartialResultSetSourceTest, ChunkedValueSetNoValue) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoValue"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoValue"); - return response[1]; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoValue"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoValue"); - }); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce(VoidMock()); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ChunkedValueSetNoValue"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -713,30 +577,14 @@ TEST(PartialResultSetSourceTest, ChunkedValueSetNoFollowingValue) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoFollowingValue"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoFollowingValue"); - return response[1]; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoFollowingValue"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetNoFollowingValue"); - }); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce(VoidMock()); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ChunkedValueSetNoFollowingValue"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -773,31 +621,14 @@ TEST(PartialResultSetSourceTest, ChunkedValueSetAtEndOfStream) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetAtEndOfStream"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetAtEndOfStream"); - return response[1]; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetAtEndOfStream"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueSetAtEndOfStream"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ChunkedValueSetAtEndOfStream"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -837,35 +668,15 @@ TEST(PartialResultSetSourceTest, ChunkedValueMergeFailure) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueMergeFailure"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueMergeFailure"); - return response[1]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueMergeFailure"); - return response[2]; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueMergeFailure"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])) + .WillOnce(ResultMock(response[2])); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); // The destructor should try to cancel the RPC to avoid deadlocks. - EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ChunkedValueMergeFailure"); - }); + EXPECT_CALL(*grpc_reader, TryCancel()).WillOnce(VoidMock()); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ChunkedValueMergeFailure"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Trying to read the next row should fail. @@ -929,46 +740,17 @@ TEST(PartialResultSetSourceTest, ErrorOnIncompleteRow) { ASSERT_TRUE(TextFormat::ParseFromString(text[i], &response[i])); } EXPECT_CALL(*grpc_reader, Read()) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return response[0]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return response[1]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return response[2]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return response[3]; - }) - .WillOnce([&response] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return response[4]; - }) - .WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return absl::optional{}; - }); - EXPECT_CALL(*grpc_reader, Finish()).WillOnce([] { - EXPECT_EQ(internal::CurrentOptions().get(), - "ErrorOnIncompleteRow"); - return Status(); - }); + .WillOnce(ResultMock(response[0])) + .WillOnce(ResultMock(response[1])) + .WillOnce(ResultMock(response[2])) + .WillOnce(ResultMock(response[3])) + .WillOnce(ResultMock(response[4])) + .WillOnce(ResultMock(absl::nullopt)); + EXPECT_CALL(*grpc_reader, Finish()).WillOnce(ResultMock(Status())); EXPECT_CALL(*grpc_reader, TryCancel()).Times(0); internal::OptionsSpan overlay(Options{}.set("uh-oh")); - auto reader = CreatePartialResultSetSource(std::move(grpc_reader), - "ErrorOnIncompleteRow"); + auto reader = CreatePartialResultSetSource(std::move(grpc_reader)); EXPECT_STATUS_OK(reader.status()); // Verify the first two rows are correct.