refactor(bigtable): future-proof ReadRows mocking (#11725)

Add a new struct `bigtable::ReadRowsParams` to consolidate configuration parameters for `Table::ReadRows()`. This future-proofs the growing of the API. It allows the library to stay up to date with the latest and greatest capabilities from Bigtable. Adds a `DataConnection::ReadRowsFull(ReadRowsParams params)` for mocking. All calls to `Table::ReadRows(...)` flow through this method. Customers that mock the library should use `MockDataConnection::ReadRowsFull(...)` in their tests.
googleapis · May 25, 2023 · 3678c29 · 3678c29
1 parent 1260047
commit 3678c29
Show file tree

Hide file tree

Showing 11 changed files with 138 additions and 36 deletions.
diff --git a/google/cloud/bigtable/data_connection.cc b/google/cloud/bigtable/data_connection.cc
@@ -74,9 +74,17 @@ future<std::vector<FailedMutation>> DataConnection::AsyncBulkApply(
       Status(StatusCode::kUnimplemented, "not-implemented"), mut.size()));
 }
 
-RowReader DataConnection::ReadRows(
-    // NOLINTNEXTLINE(performance-unnecessary-value-param)
-    std::string const&, RowSet, std::int64_t, Filter) {
+RowReader DataConnection::ReadRows(std::string const& table_name,
+                                   RowSet row_set, std::int64_t rows_limit,
+                                   Filter filter) {
+  return ReadRowsFull(ReadRowsParams{
+      std::move(table_name),
+      google::cloud::internal::CurrentOptions().get<AppProfileIdOption>(),
+      std::move(row_set), rows_limit, std::move(filter)});
+}
+
+// NOLINTNEXTLINE(performance-unnecessary-value-param)
+RowReader DataConnection::ReadRowsFull(ReadRowsParams) {
   return MakeRowReader(std::make_shared<bigtable_internal::StatusOnlyRowReader>(
       Status(StatusCode::kUnimplemented, "not implemented")));
 }

diff --git a/google/cloud/bigtable/data_connection.h b/google/cloud/bigtable/data_connection.h
@@ -42,6 +42,15 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 namespace bigtable {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
+/// Wrap the arguments to `ReadRows()`.
+struct ReadRowsParams {
+  std::string table_name;
+  std::string app_profile_id;
+  RowSet row_set;
+  std::int64_t rows_limit;
+  Filter filter = Filter::PassAllFilter();
+};
+
 /**
  * A connection to the Cloud Bigtable Data API.
  *
@@ -93,9 +102,12 @@ class DataConnection {
   virtual future<std::vector<FailedMutation>> AsyncBulkApply(
       std::string const& table_name, BulkMutation mut);
 
+  /// Prefer to use `ReadRowsFull()` in mocks.
   virtual RowReader ReadRows(std::string const& table_name, RowSet row_set,
                              std::int64_t rows_limit, Filter filter);
 
+  virtual RowReader ReadRowsFull(ReadRowsParams params);
+
   virtual StatusOr<std::pair<bool, Row>> ReadRow(std::string const& table_name,
                                                  std::string row_key,
                                                  Filter filter);

diff --git a/google/cloud/bigtable/examples/howto_mock_data_api.cc b/google/cloud/bigtable/examples/howto_mock_data_api.cc
@@ -42,7 +42,7 @@ TEST(MockTableTest, ReadRowsSuccess) {
   // yield "r1" then "r2":
   //! [simulate-call]
   std::vector<cbt::Row> rows = {cbt::Row("r1", {}), cbt::Row("r2", {})};
-  EXPECT_CALL(*mock, ReadRows)
+  EXPECT_CALL(*mock, ReadRowsFull)
       .WillOnce(Return(ByMove(cbtm::MakeRowReader(rows))));
   //! [simulate-call]
 
@@ -72,7 +72,7 @@ TEST(MockTableTest, ReadRowsFailure) {
 
   // Return a `RowReader` that yields only a failing status (no rows).
   gc::Status final_status(gc::StatusCode::kPermissionDenied, "fail");
-  EXPECT_CALL(*mock, ReadRows)
+  EXPECT_CALL(*mock, ReadRowsFull)
       .WillOnce(Return(ByMove(cbtm::MakeRowReader({}, final_status))));
 
   cbt::Table table(mock, cbt::TableResource("project", "instance", "table"));

diff --git a/google/cloud/bigtable/internal/data_connection_impl.cc b/google/cloud/bigtable/internal/data_connection_impl.cc
@@ -173,13 +173,12 @@ DataConnectionImpl::AsyncBulkApply(std::string const& table_name,
                                   app_profile_id(), table_name, std::move(mut));
 }
 
-bigtable::RowReader DataConnectionImpl::ReadRows(std::string const& table_name,
-                                                 bigtable::RowSet row_set,
-                                                 std::int64_t rows_limit,
-                                                 bigtable::Filter filter) {
+bigtable::RowReader DataConnectionImpl::ReadRowsFull(
+    bigtable::ReadRowsParams params) {
   auto impl = std::make_shared<DefaultRowReader>(
-      stub_, app_profile_id(), table_name, std::move(row_set), rows_limit,
-      std::move(filter), retry_policy(), backoff_policy());
+      stub_, std::move(params.app_profile_id), std::move(params.table_name),
+      std::move(params.row_set), params.rows_limit, std::move(params.filter),
+      retry_policy(), backoff_policy());
   return MakeRowReader(std::move(impl));
 }
 
@@ -188,8 +187,9 @@ StatusOr<std::pair<bool, bigtable::Row>> DataConnectionImpl::ReadRow(
     bigtable::Filter filter) {
   bigtable::RowSet row_set(std::move(row_key));
   std::int64_t const rows_limit = 1;
-  auto reader =
-      ReadRows(table_name, std::move(row_set), rows_limit, std::move(filter));
+  auto reader = ReadRowsFull(
+      bigtable::ReadRowsParams{table_name, app_profile_id(), std::move(row_set),
+                               rows_limit, std::move(filter)});
 
   auto it = reader.begin();
   if (it == reader.end()) return std::make_pair(false, bigtable::Row("", {}));

diff --git a/google/cloud/bigtable/internal/data_connection_impl.h b/google/cloud/bigtable/internal/data_connection_impl.h
@@ -52,10 +52,7 @@ class DataConnectionImpl : public bigtable::DataConnection {
   future<std::vector<bigtable::FailedMutation>> AsyncBulkApply(
       std::string const& table_name, bigtable::BulkMutation mut) override;
 
-  bigtable::RowReader ReadRows(std::string const& table_name,
-                               bigtable::RowSet row_set,
-                               std::int64_t rows_limit,
-                               bigtable::Filter filter) override;
+  bigtable::RowReader ReadRowsFull(bigtable::ReadRowsParams params) override;
 
   StatusOr<std::pair<bool, bigtable::Row>> ReadRow(
       std::string const& table_name, std::string row_key,

diff --git a/google/cloud/bigtable/internal/data_connection_impl_test.cc b/google/cloud/bigtable/internal/data_connection_impl_test.cc
@@ -704,6 +704,30 @@ TEST(DataConnectionTest, ReadRows) {
   EXPECT_EQ(reader.begin(), reader.end());
 }
 
+// The DefaultRowReader is tested extensively in `default_row_reader_test.cc`.
+// In this test, we just verify that the configuration is passed along.
+TEST(DataConnectionTest, ReadRowsFull) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, ReadRows)
+      .WillOnce([](auto, google::bigtable::v2::ReadRowsRequest const& request) {
+        EXPECT_EQ(kAppProfile, request.app_profile_id());
+        EXPECT_EQ(kTableName, request.table_name());
+        EXPECT_EQ(42, request.rows_limit());
+        EXPECT_THAT(request, HasTestRowSet());
+        EXPECT_THAT(request.filter(), IsTestFilter());
+
+        auto stream = std::make_unique<MockReadRowsStream>();
+        EXPECT_CALL(*stream, Read).WillOnce(Return(Status()));
+        return stream;
+      });
+
+  auto conn = TestConnection(std::move(mock));
+  internal::OptionsSpan span(CallOptions());
+  auto reader = conn->ReadRowsFull(bigtable::ReadRowsParams{
+      kTableName, kAppProfile, TestRowSet(), 42, TestFilter()});
+  EXPECT_EQ(reader.begin(), reader.end());
+}
+
 TEST(DataConnectionTest, ReadRowEmpty) {
   auto mock = std::make_shared<MockBigtableStub>();
   EXPECT_CALL(*mock, ReadRows)

diff --git a/google/cloud/bigtable/internal/data_tracing_connection.cc b/google/cloud/bigtable/internal/data_tracing_connection.cc
@@ -90,14 +90,10 @@ class DataTracingConnection : public bigtable::DataConnection {
     });
   }
 
-  bigtable::RowReader ReadRows(std::string const& table_name,
-                               bigtable::RowSet row_set,
-                               std::int64_t rows_limit,
-                               bigtable::Filter filter) override {
+  bigtable::RowReader ReadRowsFull(bigtable::ReadRowsParams params) override {
     auto span = internal::MakeSpan("bigtable::Table::ReadRows");
     auto scope = opentelemetry::trace::Scope(span);
-    auto reader = child_->ReadRows(table_name, std::move(row_set), rows_limit,
-                                   std::move(filter));
+    auto reader = child_->ReadRowsFull(std::move(params));
     return MakeTracedRowReader(std::move(span), std::move(reader));
   }
 

diff --git a/google/cloud/bigtable/internal/data_tracing_connection_test.cc b/google/cloud/bigtable/internal/data_tracing_connection_test.cc
@@ -227,7 +227,7 @@ TEST(DataTracingConnection, ReadRows) {
   auto span_catcher = InstallSpanCatcher();
 
   auto mock = std::make_shared<MockDataConnection>();
-  EXPECT_CALL(*mock, ReadRows).WillOnce([] {
+  EXPECT_CALL(*mock, ReadRowsFull).WillOnce([] {
     EXPECT_TRUE(ThereIsAnActiveSpan());
     return bigtable_mocks::MakeRowReader({}, internal::AbortedError("fail"));
   });
@@ -249,6 +249,33 @@ TEST(DataTracingConnection, ReadRows) {
               SpanAttribute<int>("gcloud.status_code", kErrorCode)))));
 }
 
+TEST(DataTracingConnection, ReadRowsFull) {
+  auto span_catcher = InstallSpanCatcher();
+
+  auto mock = std::make_shared<MockDataConnection>();
+  EXPECT_CALL(*mock, ReadRowsFull).WillOnce([] {
+    EXPECT_TRUE(ThereIsAnActiveSpan());
+    return bigtable_mocks::MakeRowReader({}, internal::AbortedError("fail"));
+  });
+
+  auto under_test = MakeDataTracingConnection(mock);
+  auto reader = under_test->ReadRowsFull(
+      bigtable::ReadRowsParams{kTableName, "app-profile-id", bigtable::RowSet(),
+                               0, bigtable::Filter::PassAllFilter()});
+  auto it = reader.begin();
+  EXPECT_THAT(*it, StatusIs(StatusCode::kAborted));
+  EXPECT_EQ(++it, reader.end());
+
+  EXPECT_THAT(
+      span_catcher->GetSpans(),
+      ElementsAre(AllOf(
+          SpanHasInstrumentationScope(), SpanKindIsClient(),
+          SpanNamed("bigtable::Table::ReadRows"),
+          SpanWithStatus(opentelemetry::trace::StatusCode::kError, "fail"),
+          SpanHasAttributes(
+              SpanAttribute<int>("gcloud.status_code", kErrorCode)))));
+}
+
 TEST(DataTracingConnection, ReadRowFound) {
   auto span_catcher = InstallSpanCatcher();
 

diff --git a/google/cloud/bigtable/mocks/mock_data_connection.h b/google/cloud/bigtable/mocks/mock_data_connection.h
@@ -36,6 +36,20 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  */
 class MockDataConnection : public bigtable::DataConnection {
  public:
+  MockDataConnection() {
+    ON_CALL(*this, ReadRows)
+        .WillByDefault([this](std::string const& table_name,
+                              bigtable::RowSet const& row_set,
+                              std::int64_t rows_limit,
+                              bigtable::Filter const& filter) {
+          bigtable::ReadRowsParams params{
+              table_name,
+              internal::CurrentOptions().get<bigtable::AppProfileIdOption>(),
+              row_set, rows_limit, filter};
+          return this->ReadRowsFull(params);
+        });
+  }
+
   MOCK_METHOD(Options, options, (), (override));
 
   MOCK_METHOD(Status, Apply,
@@ -54,11 +68,15 @@ class MockDataConnection : public bigtable::DataConnection {
               (std::string const& table_name, bigtable::BulkMutation mut),
               (override));
 
+  /// Prefer to use `ReadRowsFull()`.
   MOCK_METHOD(bigtable::RowReader, ReadRows,
               (std::string const& table_name, bigtable::RowSet row_set,
                std::int64_t rows_limit, bigtable::Filter filter),
               (override));
 
+  MOCK_METHOD(bigtable::RowReader, ReadRowsFull, (bigtable::ReadRowsParams),
+              (override));
+
   MOCK_METHOD((StatusOr<std::pair<bool, bigtable::Row>>), ReadRow,
               (std::string const& table_name, std::string row_key,
                bigtable::Filter filter),

diff --git a/google/cloud/bigtable/mocks/mock_row_reader.h b/google/cloud/bigtable/mocks/mock_row_reader.h
@@ -46,7 +46,7 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  *   std::vector<cbt::Row> rows = {cbt::Row("r1", {}), cbt::Row("r2", {})};
  *
  *   auto mock = std::shared_ptr<cbtm::MockDataConnection>();
- *   EXPECT_CALL(*mock, ReadRows)
+ *   EXPECT_CALL(*mock, ReadRowsFull)
  *       .WillOnce(Return(cbtm::MakeRowReader(rows)));
  *
  *   auto table = cbt::Table(mock);

diff --git a/google/cloud/bigtable/table_test.cc b/google/cloud/bigtable/table_test.cc
@@ -252,17 +252,15 @@ TEST(TableTest, AsyncBulkApply) {
 
 TEST(TableTest, ReadRows) {
   auto mock = std::make_shared<MockDataConnection>();
-  EXPECT_CALL(*mock, ReadRows)
-      .WillOnce([](std::string const& table_name,
-                   bigtable::RowSet const& row_set, std::int64_t rows_limit,
-                   bigtable::Filter const& filter) {
-        CheckCurrentOptions();
-        EXPECT_EQ(kTableName, table_name);
-        EXPECT_THAT(row_set, IsTestRowSet());
-        EXPECT_EQ(rows_limit, RowReader::NO_ROWS_LIMIT);
-        EXPECT_THAT(filter, IsTestFilter());
-        return bigtable_mocks::MakeRowReader({}, PermanentError());
-      });
+  EXPECT_CALL(*mock, ReadRowsFull).WillOnce([](ReadRowsParams const& params) {
+    CheckCurrentOptions();
+    EXPECT_EQ(params.table_name, kTableName);
+    EXPECT_EQ(params.app_profile_id, kAppProfileId);
+    EXPECT_THAT(params.row_set, IsTestRowSet());
+    EXPECT_EQ(params.rows_limit, RowReader::NO_ROWS_LIMIT);
+    EXPECT_THAT(params.filter, IsTestFilter());
+    return bigtable_mocks::MakeRowReader({}, PermanentError());
+  });
 
   auto table = TestTable(std::move(mock));
   auto reader = table.ReadRows(TestRowSet(), TestFilter(), CallOptions());
@@ -273,6 +271,28 @@ TEST(TableTest, ReadRows) {
 
 TEST(TableTest, ReadRowsWithRowLimit) {
   auto mock = std::make_shared<MockDataConnection>();
+  EXPECT_CALL(*mock, ReadRowsFull).WillOnce([](ReadRowsParams const& params) {
+    CheckCurrentOptions();
+    EXPECT_EQ(params.table_name, kTableName);
+    EXPECT_EQ(params.app_profile_id, kAppProfileId);
+    EXPECT_THAT(params.row_set, IsTestRowSet());
+    EXPECT_EQ(params.rows_limit, 42);
+    EXPECT_THAT(params.filter, IsTestFilter());
+    return bigtable_mocks::MakeRowReader({}, PermanentError());
+  });
+
+  auto table = TestTable(std::move(mock));
+  auto reader = table.ReadRows(TestRowSet(), 42, TestFilter(), CallOptions());
+  auto it = reader.begin();
+  EXPECT_THAT(*it, StatusIs(StatusCode::kPermissionDenied));
+  EXPECT_EQ(++it, reader.end());
+}
+
+TEST(TableTest, ReadRowsMockBackwardsCompatibility) {
+  auto mock = std::make_shared<MockDataConnection>();
+  // Ensure that existing mocks which set expectations on the legacy `ReadRows`
+  // method continue to work. This is more a test of `MockDataConnection` than
+  // of `Table`.
   EXPECT_CALL(*mock, ReadRows)
       .WillOnce([](std::string const& table_name,
                    bigtable::RowSet const& row_set, std::int64_t rows_limit,