Fix TransportSLMGetExpiredSnapshotsActionTests (elastic#105950)

Addresses test bug introduced in elastic#105721: we must consume all the `SnapshotInfo` instances before completing the final listener. Closes elastic#105922
kderusso · Mar 5, 2024 · 7c6120b · 7c6120b
1 parent fbbfbd5
commit 7c6120b
Showing 1 changed file with 14 additions and 12 deletions.
diff --git a/...src/test/java/org/elasticsearch/xpack/slm/TransportSLMGetExpiredSnapshotsActionTests.java b/...src/test/java/org/elasticsearch/xpack/slm/TransportSLMGetExpiredSnapshotsActionTests.java
@@ -11,6 +11,7 @@
 import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.action.support.RefCountingRunnable;
 import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.cluster.metadata.RepositoryMetadata;
 import org.elasticsearch.common.settings.Settings;
@@ -286,7 +287,7 @@ private static Repository createMockRepository(ThreadPool threadPool, List<Snaps
                 var repositoryData = RepositoryData.EMPTY;
                 for (SnapshotInfo snapshotInfo : snapshotInfos) {
                     var snapshotDetails = RepositoryData.SnapshotDetails.fromSnapshotInfo(snapshotInfo);
-                    if (rarely()) {
+                    if (randomBoolean()) {
                         snapshotDetails = new RepositoryData.SnapshotDetails(
                             snapshotDetails.getSnapshotState(),
                             snapshotDetails.getVersion(),
@@ -314,19 +315,20 @@ private static Repository createMockRepository(ThreadPool threadPool, List<Snaps
             final CheckedConsumer<SnapshotInfo, Exception> consumer = invocation.getArgument(3);
             final ActionListener<Void> listener = invocation.getArgument(4);
 
-            final Set<SnapshotId> snapshotIds = new HashSet<>(snapshotIdCollection);
-            for (SnapshotInfo snapshotInfo : snapshotInfos) {
-                if (snapshotIds.remove(snapshotInfo.snapshotId())) {
-                    threadPool.generic().execute(() -> {
-                        try {
-                            consumer.accept(snapshotInfo);
-                        } catch (Exception e) {
-                            fail(e);
-                        }
-                    });
+            try (var refs = new RefCountingRunnable(() -> listener.onResponse(null))) {
+                final Set<SnapshotId> snapshotIds = new HashSet<>(snapshotIdCollection);
+                for (SnapshotInfo snapshotInfo : snapshotInfos) {
+                    if (snapshotIds.remove(snapshotInfo.snapshotId())) {
+                        threadPool.generic().execute(ActionRunnable.run(refs.acquireListener(), () -> {
+                            try {
+                                consumer.accept(snapshotInfo);
+                            } catch (Exception e) {
+                                fail(e);
+                            }
+                        }));
+                    }
                 }
             }
-            listener.onResponse(null);
             return null;
         }).when(repository).getSnapshotInfo(any(), anyBoolean(), any(), any(), any());