From 61bf2e0d977f3e28e09cbcb36856273b78b46437 Mon Sep 17 00:00:00 2001
From: Pranshu Shukla <pranshushukla06@gmail.com>
Date: Fri, 12 Jul 2024 19:56:42 +0530
Subject: [PATCH] Fix test logic

Signed-off-by: Pranshu Shukla <pranshushukla06@gmail.com>
---
 .../org/opensearch/nodestats/NodeStatsIT.java | 45 ++++++++++++++-----
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/server/src/internalClusterTest/java/org/opensearch/nodestats/NodeStatsIT.java b/server/src/internalClusterTest/java/org/opensearch/nodestats/NodeStatsIT.java
index bbb5398c89aa7..1fd958293dca6 100644
--- a/server/src/internalClusterTest/java/org/opensearch/nodestats/NodeStatsIT.java
+++ b/server/src/internalClusterTest/java/org/opensearch/nodestats/NodeStatsIT.java
@@ -348,11 +348,6 @@ public void testNodeIndicesStatsOptimizedResponse() {
             NodeIndicesStats.Fields.SHARDS,
             "unknown"
         );
-        String testToXContentLevel = randomFrom(
-            NodeIndicesStats.Fields.NODE,
-            NodeIndicesStats.Fields.INDICES,
-            NodeIndicesStats.Fields.SHARDS
-        );
         internalCluster().startNode();
         ensureGreen();
         String indexName = "test1";
@@ -373,9 +368,10 @@ public void testNodeIndicesStatsOptimizedResponse() {
         response.getNodes().forEach(nodeStats -> {
             try {
                 XContentBuilder builder = XContentFactory.jsonBuilder();
+
                 builder.startObject();
                 builder = nodeStats.getIndices()
-                    .toXContent(builder, new ToXContent.MapParams(Collections.singletonMap("level", testToXContentLevel)));
+                    .toXContent(builder, new ToXContent.MapParams(Collections.singletonMap("level", NodeIndicesStats.Fields.SHARDS)));
                 builder.endObject();
 
                 Map<String, Object> xContentMap = xContentBuilderToMap(builder);
@@ -384,18 +380,43 @@ public void testNodeIndicesStatsOptimizedResponse() {
                 LinkedHashMap shardStats = (LinkedHashMap) indicesStatsMap.get(NodeIndicesStats.Fields.SHARDS);
                 switch (testLevel) {
                     case NodeIndicesStats.Fields.SHARDS:
-                        assertFalse(shardStats == null || shardStats.isEmpty());
-                        assertFalse(indicesStats == null || indicesStats.isEmpty());
+                        assertFalse(shardStats.isEmpty());
+                        assertNull(indicesStats);
+                        break;
+                    case NodeIndicesStats.Fields.INDICES:
+                        assertTrue(shardStats.isEmpty());
+                        assertNull(indicesStats);
+                        break;
+                    case NodeIndicesStats.Fields.NODE:
+                    case "null":
+                    case "unknown":
+                        assertNull(shardStats);
+                        assertNull(indicesStats);
                         break;
+                }
+
+                builder = XContentFactory.jsonBuilder();
+
+                builder.startObject();
+                builder = nodeStats.getIndices()
+                    .toXContent(builder, new ToXContent.MapParams(Collections.singletonMap("level", NodeIndicesStats.Fields.INDICES)));
+                builder.endObject();
+
+                xContentMap = xContentBuilderToMap(builder);
+                indicesStatsMap = (LinkedHashMap) xContentMap.get(NodeIndicesStats.Fields.INDICES);
+                indicesStats = (LinkedHashMap) indicesStatsMap.get(NodeIndicesStats.Fields.INDICES);
+                shardStats = (LinkedHashMap) indicesStatsMap.get(NodeIndicesStats.Fields.SHARDS);
+                switch (testLevel) {
+                    case NodeIndicesStats.Fields.SHARDS:
                     case NodeIndicesStats.Fields.INDICES:
-                        assertTrue(shardStats == null || shardStats.isEmpty());
-                        assertFalse(indicesStats == null || indicesStats.isEmpty());
+                        assertNull(shardStats);
+                        assertFalse(indicesStats.isEmpty());
                         break;
                     case NodeIndicesStats.Fields.NODE:
                     case "null":
                     case "unknown":
-                        assertTrue(shardStats == null || shardStats.isEmpty());
-                        assertTrue(indicesStats == null || indicesStats.isEmpty());
+                        assertNull(shardStats);
+                        assertNull(indicesStats);
                         break;
                 }
             } catch (IOException e) {