From d39443d51d2625e4b3aee59d1e593229e9e449d3 Mon Sep 17 00:00:00 2001
From: Yiru Tang <yiru@google.com>
Date: Wed, 3 Mar 2021 10:44:28 -0800
Subject: [PATCH] fix: testAppendWhileShutdownSuccess race (#907)

* fix:testAppendWhileShutdownSuccess race

* .

* .
---
 .../com/google/cloud/bigquery/storage/v1beta2/StreamWriter.java | 2 +-
 .../google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriter.java
index 652d87c6d0..cedee9c475 100644
--- a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriter.java
+++ b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriter.java
@@ -924,7 +924,7 @@ public void onResponse(AppendRowsResponse response) {
             IllegalStateException exception =
                 new IllegalStateException(
                     String.format(
-                        "The append result offset %s does not match " + "the expected offset %s.",
+                        "The append result offset %s does not match the expected offset %s.",
                         response.getAppendResult().getOffset().getValue(),
                         inflightBatch.getExpectedOffset()));
             inflightBatch.onFailure(exception);
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java
index aee22284ad..392565dba7 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java
@@ -678,6 +678,7 @@ public void testAppendWhileShutdownSuccess() throws Exception {
                     .toBuilder()
                     // When shutdown, we should have something in batch.
                     .setElementCountThreshold(3L)
+                    .setDelayThreshold(Duration.ofSeconds(1000))
                     .setFlowControlSettings(
                         StreamWriter.Builder.DEFAULT_FLOW_CONTROL_SETTINGS
                             .toBuilder()
@@ -1072,6 +1073,7 @@ public void testOffset() throws Exception {
                 StreamWriter.Builder.DEFAULT_BATCHING_SETTINGS
                     .toBuilder()
                     .setElementCountThreshold(2L)
+                    .setDelayThreshold(Duration.ofSeconds(1000))
                     .build())
             .build()) {