Fixed flaky test: ResourceAwareTasksTests.testTaskIdPersistsInThreadC…

…ontext (#4484) Signed-off-by: Ketan Verma <ketan9495@gmail.com>
opensearch-project · Sep 13, 2022 · 763a89f · 763a89f
1 parent 5dc3d2e
commit 763a89f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -59,6 +59,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Segment Replication] Fix timeout issue by calculating time needed to process getSegmentFiles ([#4426](https://github.com/opensearch-project/OpenSearch/pull/4426))
 - [Bug]: gradle check failing with java heap OutOfMemoryError (([#4328](https://github.com/opensearch-project/OpenSearch/
 - `opensearch.bat` fails to execute when install path includes spaces ([#4362](https://github.com/opensearch-project/OpenSearch/pull/4362))
+- Fixed flaky test `ResourceAwareTasksTests.testTaskIdPersistsInThreadContext` ([#4484](https://github.com/opensearch-project/OpenSearch/pull/4484))
 
 ### Security
 - CVE-2022-25857 org.yaml:snakeyaml DOS vulnerability ([#4341](https://github.com/opensearch-project/OpenSearch/pull/4341))

diff --git a/...src/test/java/org/opensearch/action/admin/cluster/node/tasks/ResourceAwareTasksTests.java b/...src/test/java/org/opensearch/action/admin/cluster/node/tasks/ResourceAwareTasksTests.java
@@ -42,6 +42,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
@@ -616,6 +617,10 @@ public void onFailure(Exception e) {
 
         taskTestContext.requestCompleteLatch.await();
 
+        // It is possible for the MockTaskManagerListener to be called after the response is sent already.
+        // Wait enough time for taskId to be added to taskIdsRemovedFromThreadContext before performing validations.
+        waitUntil(() -> taskIdsAddedToThreadContext.size() == taskIdsRemovedFromThreadContext.size(), 5, TimeUnit.SECONDS);
+
         assertEquals(expectedTaskIdInThreadContext.get(), actualTaskIdInThreadContext.get());
         assertThat(taskIdsAddedToThreadContext, containsInAnyOrder(taskIdsRemovedFromThreadContext.toArray()));
     }