From b80183ea23c8b78611a42d22d8c62a4ba4904a80 Mon Sep 17 00:00:00 2001
From: Yiru Tang <yiru@google.com>
Date: Tue, 2 Mar 2021 14:15:29 -0800
Subject: [PATCH] fix:test failure testAppendWhileShutdownSuccess (#904)

* fix:test failure testAppendWhileShutdownSuccess, a race special to test execution

* fix one more test

* fix another test
---
 .../storage/v1beta2/StreamWriterTest.java     | 36 ++++++++++++-------
 1 file changed, 24 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 73bf22f122..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
@@ -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());
       }
     }
   }
@@ -739,6 +743,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;
             }