fix:test failure testAppendWhileShutdownSuccess (#904)

* fix:test failure testAppendWhileShutdownSuccess, a race special to test execution * fix one more test * fix another test
googleapis · Mar 2, 2021 · b80183e · b80183e
1 parent 2e49ce8
commit b80183e
Showing 1 changed file with 24 additions and 12 deletions.
diff --git a/...uerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/StreamWriterTest.java b/...uerystorage/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;
             }