Adding UT to assert indicesSegments

Signed-off-by: Harsh Garg <[email protected]>

server/src/main/java/org/opensearch/action/admin/indices/segments/IndicesSegmentResponse.java

@@ -101,13 +101,7 @@ public Map<String, IndexSegments> getIndices() {
             }
         }
         // Add the last shardSegment from shards list which would have got missed in the loop above
-        indicesSegments.put(
-            shards[shards.length - 1].getShardRouting().getIndexName(),
-            new IndexSegments(
-                shards[shards.length - 1].getShardRouting().getIndexName(),
-                Arrays.copyOfRange(shards, startIndexPos, shards.length)
-            )
-        );
+        indicesSegments.put(startIndexName, new IndexSegments(startIndexName, Arrays.copyOfRange(shards, startIndexPos, shards.length)));
         this.indicesSegments = indicesSegments;
         return indicesSegments;
     }

server/src/test/java/org/opensearch/action/admin/indices/segments/IndicesSegmentResponseTests.java

@@ -42,7 +42,10 @@
 import org.opensearch.index.engine.Segment;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 
@@ -68,4 +71,49 @@ public void testToXContentSerialiationWithSortedFields() throws Exception {
             response.toXContent(builder, ToXContent.EMPTY_PARAMS);
         }
     }
+
+    public void testGetIndices() {
+        final int totalIndices = 5;
+        final int shardsPerIndex = 3;
+        final int segmentsPerShard = 2;
+        // Preparing a ShardSegments list, which will have (totalIndices * shardsPerIndex) shardSegments.
+        // Indices will be named -> foo1, foo2, ..., foo{totalIndices}
+        List<ShardSegments> shardSegmentsList = new ArrayList<>();
+        for (int indexName = 0; indexName < totalIndices; indexName++) {
+            for (int shardId = 0; shardId < shardsPerIndex; shardId++) {
+                ShardRouting shardRouting = TestShardRouting.newShardRouting(
+                    "foo" + indexName,
+                    shardId,
+                    "node_id",
+                    true,
+                    ShardRoutingState.STARTED
+                );
+                List<Segment> segmentList = new ArrayList<>();
+                for (int segmentNum = 0; segmentNum < segmentsPerShard; segmentNum++) {
+                    segmentList.add(new Segment("foo" + indexName + shardId + segmentNum));
+                }
+                shardSegmentsList.add(new ShardSegments(shardRouting, segmentList));
+            }
+        }
+
+        // Prepare the IndicesSegmentResponse object and get the indicesSegments map
+        IndicesSegmentResponse response = new IndicesSegmentResponse(
+            shardSegmentsList.toArray(new ShardSegments[0]),
+            totalIndices * shardsPerIndex,
+            totalIndices * shardsPerIndex,
+            0,
+            Collections.emptyList()
+        );
+        Map<String, IndexSegments> indicesSegments = response.getIndices();
+
+        assertEquals(totalIndices, indicesSegments.size());
+        for (Map.Entry<String, IndexSegments> indexSegmentEntry : indicesSegments.entrySet()) {
+            assertEquals(shardsPerIndex, indexSegmentEntry.getValue().getShards().size());
+            for (IndexShardSegments indexShardSegment : indexSegmentEntry.getValue().getShards().values()) {
+                for (ShardSegments shardSegment : indexShardSegment.getShards()) {
+                    assertEquals(segmentsPerShard, shardSegment.getSegments().size());
+                }
+            }
+        }
+    }
 }