Fixed flaky test: ResourceAwareTasksTests.testTaskIdPersistsInThreadC…

…ontext (opensearch-project#4484) (opensearch-project#4500) Signed-off-by: Daniel (dB.) Doubrovkine <[email protected]> Co-authored-by: Ketan Verma <[email protected]>
reta · Sep 14, 2022 · 3612b24 · 3612b24
1 parent 3080dfd
commit 3612b24
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/CHANGELOG.md b/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))

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