From d3fbb3efe65c060f35734ee6a9b3d650b602da29 Mon Sep 17 00:00:00 2001
From: David Turner <david.turner@elastic.co>
Date: Tue, 29 Nov 2022 17:24:45 +0000
Subject: [PATCH] Fix PersistentTasksClusterServiceTests (#92002)

Today these tests don't wait for the assertions to run before shutting
down the master service, and today if the master service shuts down then
work (such as these assertions) is silently dropped. This commit ensures
that the assertions have run before the test completes.
---
 .../PersistentTasksClusterServiceTests.java     | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/server/src/test/java/org/elasticsearch/persistent/PersistentTasksClusterServiceTests.java b/server/src/test/java/org/elasticsearch/persistent/PersistentTasksClusterServiceTests.java
index 5e5d97a7bee24..f58ff9ae5ecdf 100644
--- a/server/src/test/java/org/elasticsearch/persistent/PersistentTasksClusterServiceTests.java
+++ b/server/src/test/java/org/elasticsearch/persistent/PersistentTasksClusterServiceTests.java
@@ -50,6 +50,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
@@ -527,7 +528,7 @@ public void testPeriodicRecheckOffMaster() {
         assertFalse(service.getPeriodicRechecker().isScheduled());
     }
 
-    public void testUnassignTask() {
+    public void testUnassignTask() throws InterruptedException {
         ClusterState clusterState = initialState();
         ClusterState.Builder builder = ClusterState.builder(clusterState);
         PersistentTasksCustomMetadata.Builder tasks = PersistentTasksCustomMetadata.builder(
@@ -545,14 +546,18 @@ public void testUnassignTask() {
         clusterState = builder.metadata(metadata).nodes(nodes).build();
         setState(clusterService, clusterState);
         PersistentTasksClusterService service = createService((params, candidateNodes, currentState) -> new Assignment("_node_2", "test"));
+        final var countDownLatch = new CountDownLatch(1);
         service.unassignPersistentTask(unassignedId, tasks.getLastAllocationId(), "unassignment test", ActionListener.wrap(task -> {
             assertThat(task.getAssignment().getExecutorNode(), is(nullValue()));
             assertThat(task.getId(), equalTo(unassignedId));
             assertThat(task.getAssignment().getExplanation(), equalTo("unassignment test"));
+            countDownLatch.countDown();
         }, e -> fail()));
+
+        assertTrue(countDownLatch.await(10, TimeUnit.SECONDS));
     }
 
-    public void testUnassignNonExistentTask() {
+    public void testUnassignNonExistentTask() throws InterruptedException {
         ClusterState clusterState = initialState();
         ClusterState.Builder builder = ClusterState.builder(clusterState);
         PersistentTasksCustomMetadata.Builder tasks = PersistentTasksCustomMetadata.builder(
@@ -568,12 +573,18 @@ public void testUnassignNonExistentTask() {
         clusterState = builder.metadata(metadata).nodes(nodes).build();
         setState(clusterService, clusterState);
         PersistentTasksClusterService service = createService((params, candidateNodes, currentState) -> new Assignment("_node_2", "test"));
+        final var countDownLatch = new CountDownLatch(1);
         service.unassignPersistentTask(
             "missing-task",
             tasks.getLastAllocationId(),
             "unassignment test",
-            ActionListener.wrap(task -> fail(), e -> assertThat(e, instanceOf(ResourceNotFoundException.class)))
+            ActionListener.wrap(task -> fail(), e -> {
+                assertThat(e, instanceOf(ResourceNotFoundException.class));
+                countDownLatch.countDown();
+            })
         );
+
+        assertTrue(countDownLatch.await(10, TimeUnit.SECONDS));
     }
 
     public void testTasksNotAssignedToShuttingDownNodes() {