.

yirutang · yirutang · commit a04598a0d536 · 2021-02-16T11:22:26.000-08:00
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
@@ -1037,6 +1037,7 @@ public void testShutdownWithConnectionError() throws Exception {
     ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});
     ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});
     Thread.sleep(5000L);
+    // Move the needle for responses to be sent.
     fakeExecutor.advanceTime(Duration.ofSeconds(20));
     // Shutdown writer immediately and there will be some error happened when flushing the queue.
     writer.shutdown();
@@ -1045,7 +1046,7 @@ public void testShutdownWithConnectionError() throws Exception {
       appendFuture2.get();
       fail("Should fail with exception");
     } catch (java.util.concurrent.ExecutionException e) {
-      LOG.info("got: " + e.toString());
+      assertEquals("Request aborted due to previous failures", e.getCause().getMessage());
     }
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -1037,6 +1037,7 @@ public void testShutdownWithConnectionError() throws Exception {`
`1037`	`1037`	`ApiFuture<AppendRowsResponse> appendFuture1 = sendTestMessage(writer, new String[] {"A"});`
`1038`	`1038`	`ApiFuture<AppendRowsResponse> appendFuture2 = sendTestMessage(writer, new String[] {"B"});`
`1039`	`1039`	`Thread.sleep(5000L);`
	`1040`	`+ // Move the needle for responses to be sent.`
`1040`	`1041`	`fakeExecutor.advanceTime(Duration.ofSeconds(20));`
`1041`	`1042`	`// Shutdown writer immediately and there will be some error happened when flushing the queue.`
`1042`	`1043`	`writer.shutdown();`
`@@ -1045,7 +1046,7 @@ public void testShutdownWithConnectionError() throws Exception {`
`1045`	`1046`	`appendFuture2.get();`
`1046`	`1047`	`fail("Should fail with exception");`
`1047`	`1048`	`} catch (java.util.concurrent.ExecutionException e) {`
`1048`		`- LOG.info("got: " + e.toString());`
	`1049`	`+ assertEquals("Request aborted due to previous failures", e.getCause().getMessage());`
`1049`	`1050`	`}`
`1050`	`1051`	`}`
`1051`	`1052`	`}`