fix one more test

googleapis · Mar 2, 2021 · d276d9c · d276d9c
1 parent 6c4dc17
commit d276d9c
Showing 1 changed file with 16 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());
       }
     }
   }