diff --git a/google-cloud-jar-parent/pom.xml b/google-cloud-jar-parent/pom.xml index 20fa3e5ec4a7..e26fb46c8d96 100644 --- a/google-cloud-jar-parent/pom.xml +++ b/google-cloud-jar-parent/pom.xml @@ -71,6 +71,12 @@ + + org.awaitility + awaitility + 4.3.0 + test + junit junit diff --git a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml index fdc6eab9271f..d8f902e54fef 100644 --- a/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml +++ b/java-bigquerystorage/google-cloud-bigquerystorage/pom.xml @@ -193,6 +193,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..5c09d56d7c8b 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,15 @@ private void ProcessRowsAtSnapshot( TableReadOptions.newBuilder().setRowRestriction(filter).build()); } - ReadSession session = client.createReadSession(createSessionRequestBuilder.build()); + final CreateReadSessionRequest request = createSessionRequestBuilder.build(); + final 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) + .untilAsserted(() -> sessionRef.set(client.createReadSession(request))); + 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..871e774f579e 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,15 @@ private void ProcessRowsAtSnapshot( .setReadOptions(TableReadOptions.newBuilder().setRowRestriction(filter).build()); } - ReadSession session = client.createReadSession(createSessionRequestBuilder.build()); + final CreateReadSessionRequest request = createSessionRequestBuilder.build(); + final 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) + .untilAsserted(() -> sessionRef.set(client.createReadSession(request))); + ReadSession session = sessionRef.get(); assertEquals( 1, session.getStreamsCount(),