From 6bf94473f890523565aeaaa8572da2955045a77f Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Thu, 7 Sep 2023 12:58:56 -0400
Subject: [PATCH] Fix flaky codec tests (#9889) (#9896)

(cherry picked from commit 731b13fe5830ec98022f650db52d8c011f7c56e7)

Signed-off-by: Sarthak Aggarwal <sarthagg@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../org/opensearch/index/codec/MultiCodecReindexIT.java   | 4 ++--
 .../opensearch/index/codec/CodecCompressionLevelIT.java   | 8 ++++----
 .../org/opensearch/index/codec/MultiCodecMergeIT.java     | 5 ++---
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/modules/reindex/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecReindexIT.java b/modules/reindex/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecReindexIT.java
index e1deb5f2457d9..940463ba10825 100644
--- a/modules/reindex/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecReindexIT.java
+++ b/modules/reindex/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecReindexIT.java
@@ -135,7 +135,7 @@ private void assertReindexingWithMultipleCodecs(String destCodec, String destCod
     }
 
     private void useCodec(String index, String codec) throws ExecutionException, InterruptedException {
-        assertAcked(client().admin().indices().prepareClose(index));
+        assertAcked(client().admin().indices().prepareClose(index).setWaitForActiveShards(1));
 
         assertAcked(
             client().admin()
@@ -144,7 +144,7 @@ private void useCodec(String index, String codec) throws ExecutionException, Int
                 .get()
         );
 
-        assertAcked(client().admin().indices().prepareOpen(index));
+        assertAcked(client().admin().indices().prepareOpen(index).setWaitForActiveShards(1));
     }
 
     private void flushAndRefreshIndex(String index) {
diff --git a/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/CodecCompressionLevelIT.java b/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/CodecCompressionLevelIT.java
index 9bef3e2d3d235..7810b5810fffb 100644
--- a/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/CodecCompressionLevelIT.java
+++ b/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/CodecCompressionLevelIT.java
@@ -97,7 +97,7 @@ public void testZStandardToLuceneCodecsWithCompressionLevel() throws ExecutionEx
         );
         ensureGreen(index);
 
-        assertAcked(client().admin().indices().prepareClose(index));
+        assertAcked(client().admin().indices().prepareClose(index).setWaitForActiveShards(1));
 
         Throwable executionException = expectThrows(
             ExecutionException.class,
@@ -128,7 +128,7 @@ public void testZStandardToLuceneCodecsWithCompressionLevel() throws ExecutionEx
                 .get()
         );
 
-        assertAcked(client().admin().indices().prepareOpen(index));
+        assertAcked(client().admin().indices().prepareOpen(index).setWaitForActiveShards(1));
         ensureGreen(index);
     }
 
@@ -148,7 +148,7 @@ public void testLuceneToZStandardCodecsWithCompressionLevel() throws ExecutionEx
         );
         ensureGreen(index);
 
-        assertAcked(client().admin().indices().prepareClose(index));
+        assertAcked(client().admin().indices().prepareClose(index).setWaitForActiveShards(1));
 
         Throwable executionException = expectThrows(
             ExecutionException.class,
@@ -181,7 +181,7 @@ public void testLuceneToZStandardCodecsWithCompressionLevel() throws ExecutionEx
                 .get()
         );
 
-        assertAcked(client().admin().indices().prepareOpen(index));
+        assertAcked(client().admin().indices().prepareOpen(index).setWaitForActiveShards(1));
         ensureGreen(index);
     }
 
diff --git a/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecMergeIT.java b/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecMergeIT.java
index fda9327c79dc1..bb508282e1952 100644
--- a/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecMergeIT.java
+++ b/plugins/custom-codecs/src/internalClusterTest/java/org/opensearch/index/codec/MultiCodecMergeIT.java
@@ -49,7 +49,6 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Collections.singletonList(CustomCodecPlugin.class);
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/9872")
     public void testForceMergeMultipleCodecs() throws ExecutionException, InterruptedException {
 
         Map<String, String> codecMap = Map.of(
@@ -119,7 +118,7 @@ private void forceMergeMultipleCodecs(String finalCodec, String finalCodecMode,
     }
 
     private void useCodec(String index, String codec) throws ExecutionException, InterruptedException {
-        assertAcked(client().admin().indices().prepareClose(index));
+        assertAcked(client().admin().indices().prepareClose(index).setWaitForActiveShards(1));
 
         assertAcked(
             client().admin()
@@ -128,7 +127,7 @@ private void useCodec(String index, String codec) throws ExecutionException, Int
                 .get()
         );
 
-        assertAcked(client().admin().indices().prepareOpen(index));
+        assertAcked(client().admin().indices().prepareOpen(index).setWaitForActiveShards(1));
     }
 
     private void ingestDocs(String index) throws InterruptedException {