Fixed flaky test: ResourceAwareTasksTests.testTaskIdPersistsInThreadC…

…ontext (#4484) (#4500)

Signed-off-by: Daniel (dB.) Doubrovkine <dblock@amazon.com>
Co-authored-by: Ketan Verma <ketan9495@gmail.com>

CHANGELOG.md

@@ -29,6 +29,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Bugs for dependabot changelog verifier workflow ([#4364](https://github.com/opensearch-project/OpenSearch/pull/4364))
 - [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))

server/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()));
     }