YARN-8101. Add UT to verify node-attributes in RM nodes rest API. Contributed by Prabhu Joseph.

(cherry picked from commit 721d5c2a5f43b29c9cc4439980f21854d9e8ae27)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestRMWebServicesNodes.java

@@ -29,6 +29,8 @@ import java.io.StringReader;
 import java.util.ArrayList;
 import java.util.EnumSet;
 import java.util.Map;
+import java.util.Set;
+import java.util.HashSet;
 import java.util.TreeMap;
 import java.util.Iterator;
 
@@ -39,15 +41,22 @@ import javax.xml.parsers.DocumentBuilderFactory;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.http.JettyUtils;
 import org.apache.hadoop.yarn.api.records.ContainerStatus;
+import org.apache.hadoop.yarn.api.records.NodeAttribute;
+import org.apache.hadoop.yarn.api.records.NodeAttributeType;
 import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.api.records.NodeState;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceUtilization;
+import org.apache.hadoop.yarn.server.utils.BuilderUtils;
 import org.apache.hadoop.yarn.server.api.records.NodeHealthStatus;
 import org.apache.hadoop.yarn.server.api.records.NodeStatus;
 import org.apache.hadoop.yarn.server.api.records.OpportunisticContainersStatus;
+import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatRequest;
+import org.apache.hadoop.yarn.server.api.protocolrecords.RegisterNodeManagerRequest;
+import org.apache.hadoop.yarn.server.api.protocolrecords.RegisterNodeManagerResponse;
 import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
 import org.apache.hadoop.yarn.server.resourcemanager.ResourceManager;
+import org.apache.hadoop.yarn.server.resourcemanager.ResourceTrackerService;
 import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNode;
 import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeEvent;
 import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeEventType;
@@ -57,6 +66,7 @@ import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeStatusEvent;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerNodeReport;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.constraint.AllocationTagsManager;
+import org.apache.hadoop.yarn.util.Records;
 import org.apache.hadoop.yarn.util.RackResolver;
 import org.apache.hadoop.yarn.util.YarnVersionInfo;
 import org.apache.hadoop.yarn.webapp.GenericExceptionHandler;
@@ -879,6 +889,72 @@ public class TestRMWebServicesNodes extends JerseyTestBase {
     rm.stop();
   }
 
+  @Test
+  public void testNodeAttributesInfo() throws Exception {
+    ResourceTrackerService resourceTrackerService =
+        rm.getResourceTrackerService();
+    RegisterNodeManagerRequest registerReq =
+        Records.newRecord(RegisterNodeManagerRequest.class);
+    NodeId nodeId = NodeId.newInstance("host1", 1234);
+    Resource capability = BuilderUtils.newResource(1024, 1);
+    registerReq.setResource(capability);
+    registerReq.setNodeId(nodeId);
+    registerReq.setHttpPort(1234);
+    registerReq.setNMVersion(YarnVersionInfo.getVersion());
+    RegisterNodeManagerResponse registerResponse =
+        resourceTrackerService.registerNodeManager(registerReq);
+
+    Set<NodeAttribute> nodeAttributes = new HashSet<>();
+    nodeAttributes.add(NodeAttribute.newInstance(
+        NodeAttribute.PREFIX_DISTRIBUTED, "host",
+        NodeAttributeType.STRING, "host1"));
+    nodeAttributes.add(NodeAttribute.newInstance(
+        NodeAttribute.PREFIX_DISTRIBUTED, "rack",
+        NodeAttributeType.STRING, "rack1"));
+
+    NodeHeartbeatRequest heartbeatReq =
+        Records.newRecord(NodeHeartbeatRequest.class);
+    NodeStatus nodeStatus =
+        NodeStatus.newInstance(nodeId, 0, new ArrayList<ContainerStatus>(),
+        null, null, null, null, null);
+    heartbeatReq.setNodeStatus(nodeStatus);
+    heartbeatReq.setLastKnownNMTokenMasterKey(registerResponse
+        .getNMTokenMasterKey());
+    heartbeatReq.setLastKnownContainerTokenMasterKey(registerResponse
+        .getContainerTokenMasterKey());
+    heartbeatReq.setNodeAttributes(nodeAttributes);
+    resourceTrackerService.nodeHeartbeat(heartbeatReq);
+
+    WebResource r = resource();
+    ClientResponse response = r.path("ws").path("v1").path("cluster")
+        .path("nodes").accept("application/json").get(ClientResponse.class);
+
+    JSONObject nodesInfoJson = response.getEntity(JSONObject.class);
+    JSONArray nodes = nodesInfoJson.getJSONObject("nodes")
+        .getJSONArray("node");
+    JSONObject nodeJson = nodes.getJSONObject(0);
+    JSONArray nodeAttributesInfo = nodeJson.getJSONObject("nodeAttributesInfo")
+        .getJSONArray("nodeAttributeInfo");
+    assertEquals(nodeAttributes.size(), nodeAttributesInfo.length());
+
+    Iterator<NodeAttribute> it = nodeAttributes.iterator();
+    for (int j=0; j<nodeAttributesInfo.length(); j++) {
+      JSONObject nodeAttributeInfo = nodeAttributesInfo.getJSONObject(j);
+      NodeAttribute expectedNodeAttribute = it.next();
+      String expectedPrefix = expectedNodeAttribute.getAttributeKey()
+          .getAttributePrefix();
+      String expectedName = expectedNodeAttribute.getAttributeKey()
+          .getAttributeName();
+      String expectedType = expectedNodeAttribute.getAttributeType()
+          .toString();
+      String expectedValue = expectedNodeAttribute.getAttributeValue();
+      assertEquals(expectedPrefix, nodeAttributeInfo.getString("prefix"));
+      assertEquals(expectedName, nodeAttributeInfo.getString("name"));
+      assertEquals(expectedType, nodeAttributeInfo.getString("type"));
+      assertEquals(expectedValue, nodeAttributeInfo.getString("value"));
+    }
+  }
+
   private void verifyNodeAllocationTag(JSONObject json,
       Map<String, Map<String, Long>> expectedAllocationTags)
       throws JSONException {