test(bigquerystorage): create sessions using await() to address flakiness

Author: whowes

google-cloud-jar-parent/pom.xml

@@ -71,6 +71,12 @@
       </dependency>
 
       <!-- Test dependencies -->
+      <dependency>
+        <groupId>org.awaitility</groupId>
+        <artifactId>awaitility</artifactId>
+        <version>4.3.0</version>
+        <scope>test</scope>
+      </dependency>
       <dependency>
         <groupId>junit</groupId>
         <artifactId>junit</artifactId>

java-bigquerystorage/google-cloud-bigquerystorage/pom.xml

@@ -193,6 +193,11 @@
       <artifactId>opentelemetry-sdk-trace</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>junit</groupId>
       <artifactId>junit</artifactId>

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

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