test fix

Signed-off-by: Shivansh Arora <[email protected]>
shiv0408 · May 7, 2024 · 15d06bc · 15d06bc
1 parent d5fc147
commit 15d06bc
Showing 1 changed file with 25 additions and 30 deletions.
diff --git a/server/src/test/java/org/opensearch/cluster/node/DiscoveryNodesTests.java b/server/src/test/java/org/opensearch/cluster/node/DiscoveryNodesTests.java
@@ -63,7 +63,6 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 import java.util.stream.StreamSupport;
 
 import static java.util.Collections.singletonMap;
@@ -518,37 +517,33 @@ public void testToXContentInAPIMode() throws IOException {
         String expectedNodeAPUXContent = "%1$s\"node_%2$d\" : {\n"
             + "%1$s  \"name\" : \"name_%2$d\",\n"
             + "%1$s  \"ephemeral_id\" : \"%3$s\",\n"
-            + "%1$s  \"transport_address\" : \"0.0.0.0:%4$d\",\n"
+            + "%1$s  \"transport_address\" : \"%4$s\",\n"
             + "%1$s  \"attributes\" : {%5$s}\n"
             + "%1$s}";
 
+        logger.info(nodes);
+
         verifyToXContentInContextMode(
             CONTEXT_MODE_API,
             nodes,
-            "{\n"
-                + "  \"nodes\" : {\n"
-                + IntStream.range(0, nodes.getSize())
-                    .mapToObj(
-                        i -> String.format(
-                            Locale.ROOT,
-                            expectedNodeAPUXContent,
-                            "    ",
-                            i,
-                            nodes.get("node_" + i).getEphemeralId(),
-                            i + 1,
-                            nodes.get("node_" + i).getAttributes().isEmpty()
-                                ? " "
-                                : "\n" + "        \"custom\" : \"" + nodes.get("node_" + i).getAttributes().get("custom") + "\"\n      "
-                        )
-                    )
-                    .collect(Collectors.joining(",\n"))
-                + "\n"
-                + "  }\n"
-                + "}"
+            "{\n" + "  \"nodes\" : {\n" + nodes.getNodes().entrySet().stream().map(entry -> {
+                int id = Integer.parseInt(entry.getKey().split("_")[1]);
+                return String.format(
+                    Locale.ROOT,
+                    expectedNodeAPUXContent,
+                    "    ",
+                    id,
+                    entry.getValue().getEphemeralId(),
+                    entry.getValue().getAddress().toString(),
+                    entry.getValue().getAttributes().isEmpty()
+                        ? " "
+                        : "\n" + "        \"custom\" : \"" + entry.getValue().getAttributes().get("custom") + "\"\n      "
+                );
+            }).collect(Collectors.joining(",\n")) + "\n" + "  }\n" + "}"
         );
     }
 
-    public void testXContentInGatewayMode() throws IOException {
+    public void testToXContentInGatewayMode() throws IOException {
         DiscoveryNodes nodes = buildDiscoveryNodes();
         String expectedXContent = getExpectedXContentInGatewayMode(nodes);
 
@@ -579,28 +574,28 @@ private String getExpectedXContentInGatewayMode(DiscoveryNodes nodes) {
         String expectedNodeAPUXContent = "%1$s\"node_%2$d\" : {\n"
             + "%1$s  \"name\" : \"name_%2$d\",\n"
             + "%1$s  \"ephemeral_id\" : \"%3$s\",\n"
-            + "%1$s  \"transport_address\" : \"0.0.0.0:%4$d\",\n"
+            + "%1$s  \"transport_address\" : \"%4$s\",\n"
             + "%1$s  \"attributes\" : {%5$s},\n"
             + "%1$s  \"host_name\" : \"0.0.0.0\",\n"
             + "%1$s  \"host_address\" : \"0.0.0.0\",\n"
             + "%1$s  \"version\" : \"%6$s\",\n"
             + "%1$s  \"roles\" : [%7$s]\n"
             + "%1$s}";
 
-        return "{\n" + "  \"nodes\" : {\n" + IntStream.range(0, nodes.getSize()).mapToObj(i -> {
-            String nodeId = "node_" + i;
-            DiscoveryNode node = nodes.get(nodeId);
+        return "{\n" + "  \"nodes\" : {\n" + nodes.getNodes().entrySet().stream().map(entry -> {
+            int id = Integer.parseInt(entry.getKey().split("_")[1]);
+            DiscoveryNode node = entry.getValue();
             String indent = "    ";
             return String.format(
                 Locale.ROOT,
                 expectedNodeAPUXContent,
                 indent,
-                i,
+                id,
                 node.getEphemeralId(),
-                i + 1,
+                entry.getValue().getAddress().toString(),
                 node.getAttributes().isEmpty()
                     ? " "
-                    : "\n" + indent + "    \"custom\" : \"" + nodes.get("node_" + i).getAttributes().get("custom") + "\"\n  " + indent,
+                    : "\n" + indent + "    \"custom\" : \"" + node.getAttributes().get("custom") + "\"\n  " + indent,
                 node.getVersion(),
                 node.getRoles().isEmpty()
                     ? " "