Fix DataStream test timeout

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datastream/DatastreamResourceManager.java

@@ -69,6 +69,7 @@
 import java.util.Set;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.arrow.util.VisibleForTesting;
 import org.apache.beam.it.common.ResourceManager;
 import org.apache.beam.it.common.utils.ExceptionUtils;
@@ -376,7 +377,9 @@ public synchronized DestinationConfig buildGCSDestinationConfig(
         DestinationConfig.newBuilder().setDestinationConnectionProfile(connectionProfile.getName());
 
     GcsDestinationConfig.Builder gcsDestinationConfigBuilder =
-        GcsDestinationConfig.newBuilder().setPath(path);
+        GcsDestinationConfig.newBuilder()
+            .setPath(path)
+            .setFileRotationInterval(Duration.newBuilder().setSeconds(60).build());
 
     if (destinationOutputFormat == DestinationOutputFormat.AVRO_FILE_FORMAT) {
       gcsDestinationConfigBuilder.setAvroFileFormat(AvroFileFormat.getDefaultInstance());
@@ -637,11 +640,12 @@ public synchronized Stream pauseStream(Stream stream) {
   }
 
   public synchronized void cleanupAll() {
-    LOG.info("Cleaning up Datastream resource manager.");
-    boolean producedError = false;
+    LOG.info("Cleaning up Datastream resource manager at {}.", java.time.Instant.now());
+    AtomicBoolean producedError = new AtomicBoolean(false);
 
     for (String stream : createdStreamIds) {
       try {
+        LOG.info("Starting deleteStreamAsync get for {} at {}", stream, java.time.Instant.now());
         Failsafe.with(retryOnException())
             .get(
                 () ->
@@ -651,55 +655,70 @@ public synchronized void cleanupAll() {
                                 .setName(StreamName.format(projectId, location, stream))
                                 .build())
                         .get());
+        LOG.info("Finished deleteStreamAsync get for {} at {}", stream, java.time.Instant.now());
       } catch (Exception e) {
         if (ExceptionUtils.containsType(e, NotFoundException.class)
             || ExceptionUtils.containsMessage(e, "NOT_FOUND")) {
           LOG.warn(
               "Stream {} not found in project {}. Assuming already deleted.", stream, projectId);
         } else {
           LOG.error("Failed to delete stream {}.", stream, e);
-          producedError = true;
+          producedError.set(true);
         }
       }
     }
-    LOG.info("Successfully deleted stream(s). ");
-
-    for (String connectionProfile : createdConnectionProfileIds) {
-      try {
-        Failsafe.with(retryOnException())
-            .get(
-                () ->
-                    datastreamClient
-                        .deleteConnectionProfileAsync(
-                            DeleteConnectionProfileRequest.newBuilder()
-                                .setName(
-                                    ConnectionProfileName.format(
-                                        projectId, location, connectionProfile))
-                                .build())
-                        .get());
-      } catch (Exception e) {
-        if (ExceptionUtils.containsType(e, NotFoundException.class)
-            || ExceptionUtils.containsMessage(e, "NOT_FOUND")) {
-          LOG.warn(
-              "Connection Profile {} not found in project {}. Assuming already deleted.",
-              connectionProfile,
-              projectId);
-        } else {
-          LOG.error("Failed to delete connection profile {}.", connectionProfile, e);
-          producedError = true;
-        }
-      }
-    }
-    LOG.info("Successfully deleted connection profile(s). ");
+    LOG.info("Successfully deleted stream(s) at {}. ", java.time.Instant.now());
+
+    createdConnectionProfileIds.parallelStream()
+        .forEach(
+            connectionProfile -> {
+              try {
+                LOG.info(
+                    "Starting deleteConnectionProfileAsync get for {} at {}",
+                    connectionProfile,
+                    java.time.Instant.now());
+                Failsafe.with(retryOnException())
+                    .get(
+                        () ->
+                            datastreamClient
+                                .deleteConnectionProfileAsync(
+                                    DeleteConnectionProfileRequest.newBuilder()
+                                        .setName(
+                                            ConnectionProfileName.format(
+                                                projectId, location, connectionProfile))
+                                        .build())
+                                .get());
+                LOG.info(
+                    "Finished deleteConnectionProfileAsync get for {} at {}",
+                    connectionProfile,
+                    java.time.Instant.now());
+              } catch (Exception e) {
+                if (ExceptionUtils.containsType(e, NotFoundException.class)
+                    || ExceptionUtils.containsMessage(e, "NOT_FOUND")) {
+                  LOG.warn(
+                      "Connection Profile {} not found in project {}. Assuming already deleted.",
+                      connectionProfile,
+                      projectId);
+                } else {
+                  LOG.error(
+                      "Failed to delete connection profile {} at {}.",
+                      connectionProfile,
+                      java.time.Instant.now(),
+                      e);
+                  producedError.set(true);
+                }
+              }
+            });
+    LOG.info("Successfully deleted connection profile(s) at {}. ", java.time.Instant.now());
 
     try {
       datastreamClient.close();
     } catch (Exception e) {
       LOG.error("Failed to close datastream client. ");
-      producedError = true;
+      producedError.set(true);
     }
 
-    if (producedError) {
+    if (producedError.get()) {
       throw new DatastreamResourceManagerException(
           "Failed to delete resources. Check above for errors.");
     }

v2/datastream-to-bigquery/src/test/java/com/google/cloud/teleport/v2/templates/DataStreamToBigQueryIT.java

@@ -67,6 +67,7 @@
 import org.apache.commons.lang3.RandomStringUtils;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -103,16 +104,19 @@ enum JDBCType {
 
   private CloudSqlResourceManager cloudSqlResourceManager;
   private PubsubResourceManager pubsubResourceManager;
-  private DatastreamResourceManager datastreamResourceManager;
+  private static DatastreamResourceManager datastreamResourceManager;
   private BigQueryResourceManager bigQueryResourceManager;
 
   @Before
   public void setUp() throws IOException {
-    datastreamResourceManager =
-        DatastreamResourceManager.builder(testName, PROJECT, REGION)
-            .setCredentialsProvider(credentialsProvider)
-            .setPrivateConnectivity("datastream-connect-2")
-            .build();
+    if (datastreamResourceManager == null) {
+      datastreamResourceManager =
+          DatastreamResourceManager.builder(
+                  DataStreamToBigQueryIT.class.getSimpleName(), PROJECT, REGION)
+              .setCredentialsProvider(credentialsProvider)
+              .setPrivateConnectivity("datastream-connect-2")
+              .build();
+    }
 
     bigQueryResourceManager =
         BigQueryResourceManager.builder(testName, PROJECT, credentials).build();
@@ -143,10 +147,12 @@ public void setUp() throws IOException {
   @After
   public void cleanUp() {
     ResourceManagerUtils.cleanResources(
-        cloudSqlResourceManager,
-        pubsubResourceManager,
-        datastreamResourceManager,
-        bigQueryResourceManager);
+        cloudSqlResourceManager, pubsubResourceManager, bigQueryResourceManager);
+  }
+
+  @AfterClass
+  public static void tearDownClass() {
+    ResourceManagerUtils.cleanResources(datastreamResourceManager);
   }
 
   @Test