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(),