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..7e023ee26edb 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
@@ -16,6 +16,10 @@
package com.google.cloud.bigquery.storage.v1beta1.it;
+import static org.awaitility.Awaitility.await;
+import com.google.api.gax.rpc.NotFoundException;
+import java.util.Objects;
+
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
@@ -1230,7 +1234,13 @@ private void ProcessRowsAtSnapshot(
TableReadOptions.newBuilder().setRowRestriction(filter).build());
}
- ReadSession session = client.createReadSession(createSessionRequestBuilder.build());
+ final CreateReadSessionRequest request = createSessionRequestBuilder.build();
+ ReadSession session =
+ await()
+ .atMost(Duration.ofSeconds(30))
+ .pollInterval(Duration.ofSeconds(1))
+ .ignoreException(NotFoundException.class)
+ .until(() -> client.createReadSession(request), Objects::nonNull);
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..e7cf800f3bef 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
@@ -16,6 +16,10 @@
package com.google.cloud.bigquery.storage.v1beta2.it;
+import static org.awaitility.Awaitility.await;
+import com.google.api.gax.rpc.NotFoundException;
+import java.util.Objects;
+
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
@@ -1211,7 +1215,13 @@ private void ProcessRowsAtSnapshot(
.setReadOptions(TableReadOptions.newBuilder().setRowRestriction(filter).build());
}
- ReadSession session = client.createReadSession(createSessionRequestBuilder.build());
+ final CreateReadSessionRequest request = createSessionRequestBuilder.build();
+ ReadSession session =
+ await()
+ .atMost(Duration.ofSeconds(30))
+ .pollInterval(Duration.ofSeconds(1))
+ .ignoreException(NotFoundException.class)
+ .until(() -> client.createReadSession(request), Objects::nonNull);
assertEquals(
1,
session.getStreamsCount(),