cleanup(flake): add retry loop to mitigate flakiness (#9752)

Author: scotthart

google/cloud/bigquery/integration_tests/bigquery_read_integration_test.cc

@@ -13,9 +13,11 @@
 // limitations under the License.
 
 #include "google/cloud/bigquery/bigquery_read_client.h"
+#include "google/cloud/bigquery/bigquery_read_connection_idempotency_policy.h"
 #include "google/cloud/bigquery/bigquery_read_options.h"
 #include "google/cloud/bigquery/internal/bigquery_read_stub_factory.h"
 #include "google/cloud/common_options.h"
+#include "google/cloud/idempotency.h"
 #include "google/cloud/internal/getenv.h"
 #include "google/cloud/log.h"
 #include "google/cloud/testing_util/integration_test.h"
@@ -35,6 +37,31 @@ using ::testing::HasSubstr;
 using ::testing::Not;
 using ::testing::NotNull;
 
+/// In order to reduce flakiness, allow retries on CreateReadSession and
+/// SplitReadStream.
+class RelaxedTestIdempotencyPolicy
+    : public BigQueryReadConnectionIdempotencyPolicy {
+ public:
+  ~RelaxedTestIdempotencyPolicy() override = default;
+
+  std::unique_ptr<BigQueryReadConnectionIdempotencyPolicy> clone()
+      const override {
+    return absl::make_unique<RelaxedTestIdempotencyPolicy>();
+  }
+
+  google::cloud::Idempotency CreateReadSession(
+      google::cloud::bigquery::storage::v1::CreateReadSessionRequest const&)
+      override {
+    return google::cloud::Idempotency::kIdempotent;
+  }
+
+  google::cloud::Idempotency SplitReadStream(
+      google::cloud::bigquery::storage::v1::SplitReadStreamRequest const&)
+      override {
+    return google::cloud::Idempotency::kIdempotent;
+  }
+};
+
 class BigQueryReadIntegrationTest
     : public ::google::cloud::testing_util::IntegrationTest {
  protected:
@@ -131,8 +158,15 @@ TEST_F(BigQueryReadIntegrationTest, SplitReadStreamProtoFailure) {
 }
 
 TEST_F(BigQueryReadIntegrationTest, CreateReadSessionSuccess) {
-  auto client =
-      BigQueryReadClient(MakeBigQueryReadConnection(TestSuccessOptions()));
+  auto client = BigQueryReadClient(MakeBigQueryReadConnection(
+      TestSuccessOptions()
+          .set<bigquery::BigQueryReadRetryPolicyOption>(
+              absl::make_unique<BigQueryReadLimitedErrorCountRetryPolicy>(4))
+          .set<bigquery::BigQueryReadBackoffPolicyOption>(
+              absl::make_unique<ExponentialBackoffPolicy>(
+                  std::chrono::seconds(15), std::chrono::seconds(60), 2.0))
+          .set<bigquery::BigQueryReadConnectionIdempotencyPolicyOption>(
+              absl::make_unique<RelaxedTestIdempotencyPolicy>())));
   ::google::cloud::bigquery::storage::v1::ReadSession read_session;
   read_session.set_data_format(
       google::cloud::bigquery::storage::v1::DataFormat::AVRO);