diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java index 1e408fadf33..1ca69fe975d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java @@ -58,6 +58,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.ThreadLocalRandom; import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; @@ -237,17 +238,24 @@ public void invalidDatabase() throws Exception { RemoteSpannerHelper helper = env.getTestHelper(); DatabaseClient invalidClient = helper.getClient().getDatabaseClient(DatabaseId.of(helper.getInstanceId(), "invalid")); - ApiFuture row = - invalidClient - .singleUse(TimestampBound.strong()) - .readRowAsync(TABLE_NAME, Key.of("k99"), ALL_COLUMNS); + Thread.sleep(ThreadLocalRandom.current().nextLong(100L)); try { + // The NOT_FOUND error can come from both the call to invalidClient.singleUse() as well as + // from the call to row.get(), which is why both need to be inside the try block. + ApiFuture row = + invalidClient + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k99"), ALL_COLUMNS); row.get(); fail("missing expected exception"); - } catch (ExecutionException e) { - assertThat(e.getCause()).isInstanceOf(SpannerException.class); - SpannerException se = (SpannerException) e.getCause(); - assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND); + } catch (ExecutionException | SpannerException thrownException) { + SpannerException spannerException; + if (thrownException instanceof ExecutionException) { + spannerException = (SpannerException) thrownException.getCause(); + } else { + spannerException = (SpannerException) thrownException; + } + assertEquals(ErrorCode.NOT_FOUND, spannerException.getErrorCode()); } }