From 6c4dc17320d9101844187ee6b9369b9d471f5314 Mon Sep 17 00:00:00 2001
From: yirutang <yiru@google.com>
Date: Tue, 2 Mar 2021 12:07:40 -0800
Subject: [PATCH 1/3] fix:test failure testAppendWhileShutdownSuccess, a race
 special to test execution

---
 .../cloud/bigquery/storage/v1beta2/StreamWriterTest.java  | 8 ++++++++
 1 file changed, 8 insertions(+)

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 73bf22f122..723ba393f6 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
@@ -739,6 +739,14 @@ public Throwable call() {
                 }
               }
               return null;
+            } catch (ExecutionException ex) {
+              // Some wiredness in test presubmit runs, it seems this thread is always started after
+              // the main thread.
+              if (ex.getCause().getClass() == AbortedException.class) {
+                return null;
+              } else {
+                return ex;
+              }
             } catch (Exception e) {
               return e;
             }

From d276d9ca0d23be39f94a548965028c36c69ecaca Mon Sep 17 00:00:00 2001
From: yirutang <yiru@google.com>
Date: Tue, 2 Mar 2021 13:09:49 -0800
Subject: [PATCH 2/3] fix one more test

---
 .../storage/v1beta2/StreamWriterTest.java     | 28 +++++++++++--------
 1 file changed, 16 insertions(+), 12 deletions(-)

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 723ba393f6..e6e006af63 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
@@ -471,19 +471,23 @@ public void testBatchIsFlushedWithError() throws Exception {
               .build());
 
       ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
-      ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});
-      ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[] {"C"});
-      try {
-        appendFuture2.get();
-      } catch (ExecutionException ex) {
-        assertEquals(DataLossException.class, ex.getCause().getClass());
-      }
-      assertFalse(appendFuture3.isDone());
-      writer.shutdown();
       try {
-        appendFuture3.get();
-      } catch (ExecutionException ex) {
-        assertEquals(AbortedException.class, ex.getCause().getClass());
+        ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[]{"B"});
+        ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[]{"C"});
+        try {
+          appendFuture2.get();
+        } catch (ExecutionException ex) {
+          assertEquals(DataLossException.class, ex.getCause().getClass());
+        }
+        assertFalse(appendFuture3.isDone());
+        writer.shutdown();
+        try {
+          appendFuture3.get();
+        } catch (ExecutionException ex) {
+          assertEquals(AbortedException.class, ex.getCause().getClass());
+        }
+      } catch (IllegalStateException ex) {
+        assertEquals("Stream already failed.", ex.getMessage());
       }
     }
   }

From 5a9f2c13ef2c44c95d80ad70a5e07075bde58d94 Mon Sep 17 00:00:00 2001
From: yirutang <yiru@google.com>
Date: Tue, 2 Mar 2021 13:10:24 -0800
Subject: [PATCH 3/3] fix another test

---
 .../cloud/bigquery/storage/v1beta2/StreamWriterTest.java      | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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 e6e006af63..aee22284ad 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
@@ -472,8 +472,8 @@ public void testBatchIsFlushedWithError() throws Exception {
 
       ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
       try {
-        ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[]{"B"});
-        ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[]{"C"});
+        ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});
+        ApiFuture<AppendRowsResponse> appendFuture3 = sendTestMessage(writer, new String[] {"C"});
         try {
           appendFuture2.get();
         } catch (ExecutionException ex) {