From 260ae90b308b5b2c84598e70c0683efc563c3d3a Mon Sep 17 00:00:00 2001
From: Shourya Dutta Biswas <114977491+shourya035@users.noreply.github.com>
Date: Wed, 2 Aug 2023 19:22:56 +0530
Subject: [PATCH] Fixing flaky testStatsOnShardUnassigned IT

Signed-off-by: Shourya Dutta Biswas <114977491+shourya035@users.noreply.github.com>
---
 .../remotestore/RemoteStoreStatsIT.java       | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java
index 1c7f14701b3e7..ce95ffd9bdc3b 100644
--- a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreStatsIT.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.remotestore;
 
+import org.junit.After;
 import org.junit.Before;
 import org.opensearch.action.admin.cluster.remotestore.restore.RestoreRemoteStoreRequest;
 import org.opensearch.action.admin.cluster.remotestore.stats.RemoteStoreStats;
@@ -46,6 +47,11 @@ public void setup() {
         setupRepo();
     }
 
+    @After
+    public void removeAllIndices() {
+        assertAcked(client().admin().indices().prepareDelete("*").get());
+    }
+
     public void testStatsResponseFromAllNodes() {
 
         // Step 1 - We create cluster, create an index, and then index documents into. We also do multiple refreshes/flushes
@@ -432,22 +438,32 @@ public void testStatsOnShardRelocation() {
         assertEquals(destinationNode, replicaShardStat.getShardRouting().currentNodeId());
     }
 
-    public void testStatsOnShardUnassigned() throws IOException {
+    public void testStatsOnShardUnassigned() throws Exception {
         // Scenario:
-        // - Create index with single primary and two replica shard
+        // - Create index with single primary and N-1 replica shard (N = number of data nodes)
         // - Index documents
         // - Stop one data node
         // - Assert:
         // a. Total shard Count in the response object is equal to the previous node count
         // b. Successful shard count in the response object is equal to the new node count
-        createIndex(INDEX_NAME, remoteStoreIndexSettings(2, 1));
+        int initialDataNodeCount = client().admin().cluster().prepareHealth().get().getNumberOfDataNodes();
+        createIndex(INDEX_NAME, remoteStoreIndexSettings(initialDataNodeCount - 1, 1));
         ensureGreen(INDEX_NAME);
         indexDocs();
-        int dataNodeCountBeforeStop = client().admin().cluster().prepareHealth().get().getNumberOfDataNodes();
-        internalCluster().stopRandomDataNode();
+
+        // Recording stats before stopping node
         RemoteStoreStatsResponse response = client().admin().cluster().prepareRemoteStoreStats(INDEX_NAME, "0").get();
+        int totalShardsBeforeStop = response.getTotalShards();
+
+        // Stop one data node
+        internalCluster().stopRandomDataNode();
+        ensureYellowAndNoInitializingShards(INDEX_NAME);
         int dataNodeCountAfterStop = client().admin().cluster().prepareHealth().get().getNumberOfDataNodes();
-        assertEquals(dataNodeCountBeforeStop, response.getTotalShards());
+        assertEquals(initialDataNodeCount - 1, dataNodeCountAfterStop);
+
+        // Recording stats again
+        response = client().admin().cluster().prepareRemoteStoreStats(INDEX_NAME, "0").get();
+        assertEquals(totalShardsBeforeStop, response.getTotalShards());
         assertEquals(dataNodeCountAfterStop, response.getSuccessfulShards());
     }