diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml index 5d2607f2ac0e..cb4f796140c6 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml +++ b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml @@ -197,6 +197,11 @@ opentelemetry-sdk-trace test + + org.awaitility + awaitility + test + junit junit diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java index 04191a38a45e..4e631baebe99 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta1/it/ITBigQueryStorageTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertWithMessage; +import static org.awaitility.Awaitility.await; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -26,6 +27,7 @@ import static org.junit.jupiter.api.Assertions.fail; import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.api.gax.rpc.NotFoundException; import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.UnauthenticatedException; import com.google.auth.oauth2.ServiceAccountCredentials; @@ -81,6 +83,7 @@ import java.util.List; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import java.util.logging.Logger; import org.apache.avro.Conversions; import org.apache.avro.LogicalTypes; @@ -1230,7 +1233,19 @@ private void ProcessRowsAtSnapshot( TableReadOptions.newBuilder().setRowRestriction(filter).build()); } - ReadSession session = client.createReadSession(createSessionRequestBuilder.build()); + CreateReadSessionRequest request = createSessionRequestBuilder.build(); + AtomicReference sessionRef = new AtomicReference<>(); + await() + .atMost(Duration.ofSeconds(10)) + .pollInterval(Duration.ofSeconds(1)) + // retry if the newly-created table has not yet fully propagated + .ignoreException(NotFoundException.class) + .until( + () -> { + sessionRef.set(client.createReadSession(request)); + return true; + }); + ReadSession session = sessionRef.get(); assertEquals( 1, session.getStreamsCount(), diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java index f48d41acbd70..88b2f48b8b7f 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java +++ b/java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1beta2/it/ITBigQueryStorageTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertWithMessage; +import static org.awaitility.Awaitility.await; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -26,6 +27,7 @@ import static org.junit.jupiter.api.Assertions.fail; import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.api.gax.rpc.NotFoundException; import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.UnauthenticatedException; import com.google.auth.oauth2.ServiceAccountCredentials; @@ -77,6 +79,7 @@ import java.util.Iterator; import java.util.List; import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; import java.util.logging.Logger; import org.apache.avro.Conversions; import org.apache.avro.LogicalTypes; @@ -1211,7 +1214,19 @@ private void ProcessRowsAtSnapshot( .setReadOptions(TableReadOptions.newBuilder().setRowRestriction(filter).build()); } - ReadSession session = client.createReadSession(createSessionRequestBuilder.build()); + CreateReadSessionRequest request = createSessionRequestBuilder.build(); + AtomicReference sessionRef = new AtomicReference<>(); + await() + .atMost(Duration.ofSeconds(10)) + .pollInterval(Duration.ofSeconds(1)) + // retry if the newly-created table has not yet fully propagated + .ignoreException(NotFoundException.class) + .until( + () -> { + sessionRef.set(client.createReadSession(request)); + return true; + }); + ReadSession session = sessionRef.get(); assertEquals( 1, session.getStreamsCount(),