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

@@ -197,6 +197,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

@@ -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;
@@ -79,6 +81,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
@@ -1230,7 +1233,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

@@ -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;
@@ -76,6 +78,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.UUID;
 import java.util.logging.Logger;
 import org.apache.avro.Conversions;
@@ -1211,7 +1214,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(),