Overload safeGet and safeAwait to also accept timeouts

Author: gmarouli

test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java

@@ -2477,9 +2477,19 @@ public static void safeAcquire(int permits, Semaphore semaphore) {
      * @return The value with which the {@code listener} was completed.
      */
     public static <T> T safeAwait(SubscribableListener<T> listener) {
+        return safeAwait(listener, SAFE_AWAIT_TIMEOUT.getMillis(), TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Wait for the successful completion of the given {@link SubscribableListener}, respecting the provided timeout,
+     * preserving the thread's interrupt status flag and converting all exceptions into an {@link AssertionError} to trigger a test failure.
+     *
+     * @return The value with which the {@code listener} was completed.
+     */
+    public static <T> T safeAwait(SubscribableListener<T> listener, long timeout, TimeUnit unit) {
         final var future = new TestPlainActionFuture<T>();
         listener.addListener(future);
-        return safeGet(future);
+        return safeGet(future, timeout, unit);
     }
 
     /**
@@ -2509,8 +2519,18 @@ public static <T extends ActionResponse> T safeExecute(ElasticsearchClient clien
      * @return The value with which the {@code future} was completed.
      */
     public static <T> T safeGet(Future<T> future) {
+        return safeGet(future, SAFE_AWAIT_TIMEOUT.millis(), TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Wait for the successful completion of the given {@link Future}, respecting the provided timeout, preserving the
+     * thread's interrupt status flag and converting all exceptions into an {@link AssertionError} to trigger a test failure.
+     *
+     * @return The value with which the {@code future} was completed.
+     */
+    public static <T> T safeGet(Future<T> future, long timeout, TimeUnit unit) {
         try {
-            return future.get(SAFE_AWAIT_TIMEOUT.millis(), TimeUnit.MILLISECONDS);
+            return future.get(timeout, unit);
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
             throw new AssertionError("safeGet: interrupted waiting for SubscribableListener", e);

x-pack/plugin/downsample/src/internalClusterTest/java/org/elasticsearch/xpack/downsample/DataStreamLifecycleDownsampleDisruptionIT.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.action.admin.indices.rollover.RolloverAction;
 import org.elasticsearch.action.admin.indices.rollover.RolloverRequest;
 import org.elasticsearch.action.downsample.DownsampleConfig;
-import org.elasticsearch.action.support.TestPlainActionFuture;
 import org.elasticsearch.cluster.metadata.DataStreamLifecycle;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.service.ClusterService;
@@ -101,8 +100,7 @@ public void testDataStreamLifecycleDownsampleRollingRestart() throws Exception {
         );
 
         logger.info("-> Starting the disruption.");
-        internalCluster().rollingRestart(new InternalTestCluster.RestartCallback() {
-        });
+        internalCluster().rollingRestart(new InternalTestCluster.RestartCallback());
 
         ensureDownsamplingStatus(targetIndex, Set.of(IndexMetadata.DownsampleTaskStatus.SUCCESS), TimeValue.timeValueSeconds(120));
         ensureGreen(targetIndex);
@@ -115,28 +113,18 @@ private void ensureDownsamplingStatus(
         TimeValue timeout
     ) {
         final var clusterService = internalCluster().getCurrentMasterNodeInstance(ClusterService.class);
-        final IndexMetadata.DownsampleTaskStatus[] downsamplingStatus = new IndexMetadata.DownsampleTaskStatus[1];
         final var listener = ClusterServiceUtils.addTemporaryStateListener(clusterService, clusterState -> {
             final var indexMetadata = clusterState.metadata().getProject().index(downsampledIndex);
             if (indexMetadata == null) {
                 return false;
             }
-            downsamplingStatus[0] = INDEX_DOWNSAMPLE_STATUS.get(indexMetadata.getSettings());
-            return expectedStatuses.contains(downsamplingStatus[0]);
-        });
-        try {
-            final var future = new TestPlainActionFuture<Void>();
-            listener.addListener(future);
-            future.get(timeout.getMillis(), TimeUnit.MILLISECONDS);
-            logger.info("-> Downsampling status for index [{}] is [{}]", downsampledIndex, downsamplingStatus[0]);
-        } catch (Exception e) {
-            if (e instanceof InterruptedException) {
-                Thread.currentThread().interrupt();
+            var downsamplingStatus = INDEX_DOWNSAMPLE_STATUS.get(indexMetadata.getSettings());
+            if (expectedStatuses.contains(downsamplingStatus)) {
+                logger.info("-> Downsampling status for index [{}] is [{}]", downsampledIndex, downsamplingStatus);
+                return true;
             }
-            throw new AssertionError(
-                "Error while waiting for " + expectedStatuses + " but found '" + downsamplingStatus[0] + "'. " + e.getMessage(),
-                e
-            );
-        }
+            return false;
+        });
+        safeAwait(listener, timeout.millis(), TimeUnit.MILLISECONDS);
     }
 }