YARN-9123. Clean up and split testcases in TestNMWebServices for GPU support. Contributed by Szilard Nemeth.

Signed-off-by: Wei-Chiu Chuang <weichiu@apache.org>

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/webapp/TestNMWebServices.java

@@ -18,6 +18,7 @@
 
 package org.apache.hadoop.yarn.server.nodemanager.webapp;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.inject.Guice;
 import com.google.inject.servlet.ServletModule;
 import com.sun.jersey.api.client.ClientResponse;
@@ -74,7 +75,6 @@
 import org.codehaus.jettison.json.JSONException;
 import org.codehaus.jettison.json.JSONObject;
 import org.junit.AfterClass;
-import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.w3c.dom.Document;
@@ -93,13 +93,13 @@
 import java.net.URI;
 import java.net.URL;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 import static org.apache.hadoop.yarn.webapp.WebServicesTestUtils.assertResponseStatusCode;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
@@ -110,6 +110,7 @@
  */
 public class TestNMWebServices extends JerseyTestBase {
 
+  private static final long NM_RESOURCE_VALUE = 1000L;
   private static NodeManager.NMContext nmContext;
   private static ResourceView resourceView;
   private static ApplicationACLsManager aclsManager;
@@ -189,6 +190,89 @@ public boolean isPmemCheckEnabled() {
         Guice.createInjector(new WebServletModule()));
   }
 
+  private void setupMockPluginsWithNmResourceInfo() throws YarnException {
+    ResourcePlugin mockPlugin1 = mock(ResourcePlugin.class);
+    NMResourceInfo nmResourceInfo1 = new NMResourceInfo() {
+      private long a = NM_RESOURCE_VALUE;
+
+      public long getA() {
+        return a;
+      }
+    };
+    when(mockPlugin1.getNMResourceInfo()).thenReturn(nmResourceInfo1);
+
+    ResourcePluginManager pluginManager = createResourceManagerWithPlugins(
+        ImmutableMap.<String, ResourcePlugin>builder()
+            .put("resource-1", mockPlugin1)
+            .put("yarn.io/resource-1", mockPlugin1)
+            .put("resource-2", mock(ResourcePlugin.class))
+            .build()
+    );
+
+    nmContext.setResourcePluginManager(pluginManager);
+  }
+
+  private void setupMockPluginsWithGpuResourceInfo() throws YarnException {
+    GpuDeviceInformation gpuDeviceInformation = new GpuDeviceInformation();
+    gpuDeviceInformation.setDriverVersion("1.2.3");
+    gpuDeviceInformation.setGpus(Arrays.asList(new PerGpuDeviceInformation()));
+
+    ResourcePlugin mockPlugin1 = mock(ResourcePlugin.class);
+    List<GpuDevice> totalGpuDevices = Arrays.asList(
+        new GpuDevice(1, 1), new GpuDevice(2, 2), new GpuDevice(3, 3));
+    List<AssignedGpuDevice> assignedGpuDevices = Arrays.asList(
+        new AssignedGpuDevice(2, 2, createContainerId(1)),
+        new AssignedGpuDevice(3, 3, createContainerId(2)));
+    NMResourceInfo nmResourceInfo1 = new NMGpuResourceInfo(gpuDeviceInformation,
+        totalGpuDevices,
+        assignedGpuDevices);
+    when(mockPlugin1.getNMResourceInfo()).thenReturn(nmResourceInfo1);
+
+    ResourcePluginManager pluginManager = createResourceManagerWithPlugins(
+        ImmutableMap.<String, ResourcePlugin>builder()
+            .put("resource-1", mockPlugin1)
+            .put("yarn.io/resource-1", mockPlugin1)
+            .put("resource-2", mock(ResourcePlugin.class))
+            .build()
+    );
+
+    nmContext.setResourcePluginManager(pluginManager);
+  }
+
+  private ResourcePluginManager createResourceManagerWithPlugins(
+      Map<String, ResourcePlugin> plugins) {
+    ResourcePluginManager pluginManager = mock(ResourcePluginManager.class);
+    when(pluginManager.getNameToPlugins()).thenReturn(plugins);
+    return pluginManager;
+  }
+
+  private void assertNMResourceInfoResponse(ClientResponse response, long value)
+      throws JSONException {
+    assertEquals("MediaType of the response is not the expected!",
+        MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
+        response.getType().toString());
+    JSONObject json = response.getEntity(JSONObject.class);
+    assertEquals("Unexpected value in the json response!", (int) value,
+        json.get("a"));
+  }
+
+  private void assertEmptyNMResourceInfo(ClientResponse response) {
+    assertEquals("MediaType of the response is not the expected!",
+        MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
+        response.getType().toString());
+    JSONObject json = response.getEntity(JSONObject.class);
+    assertEquals("Unexpected value in the json response!",
+        0, json.length());
+  }
+
+  private ClientResponse getNMResourceResponse(WebResource resource,
+      String resourceName) {
+    return resource.path("ws").path("v1").path("node").path("resources")
+        .path(resourceName).accept(MediaType.APPLICATION_JSON)
+        .get(ClientResponse.class);
+  }
+
+
   @Before
   @Override
   public void setUp() throws Exception {
@@ -431,109 +515,71 @@ public void testNMRedirect() {
   }
 
   @Test
-  public void testGetNMResourceInfo()
+  public void testGetNMResourceInfoSuccessful()
-      throws YarnException, InterruptedException, JSONException {
+      throws YarnException, JSONException {
-    ResourcePluginManager rpm = mock(ResourcePluginManager.class);
+    setupMockPluginsWithNmResourceInfo();
-    Map<String, ResourcePlugin> namesToPlugins = new HashMap<>();
-    ResourcePlugin mockPlugin1 = mock(ResourcePlugin.class);
-    NMResourceInfo nmResourceInfo1 = new NMResourceInfo() {
-      public long a = 1000L;
-    };
-    when(mockPlugin1.getNMResourceInfo()).thenReturn(nmResourceInfo1);
-    namesToPlugins.put("resource-1", mockPlugin1);
-    namesToPlugins.put("yarn.io/resource-1", mockPlugin1);
-    ResourcePlugin mockPlugin2 = mock(ResourcePlugin.class);
-    namesToPlugins.put("resource-2", mockPlugin2);
-    when(rpm.getNameToPlugins()).thenReturn(namesToPlugins);
-
-    nmContext.setResourcePluginManager(rpm);
 
     WebResource r = resource();
-    ClientResponse response = r.path("ws").path("v1").path("node").path(
+    ClientResponse response = getNMResourceResponse(r, "resource-1");
-        "resources").path("resource-2").accept(MediaType.APPLICATION_JSON).get(
+    assertNMResourceInfoResponse(response, NM_RESOURCE_VALUE);
-        ClientResponse.class);
+  }
-    assertEquals(MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
-        response.getType().toString());
 
-    // Access resource-2 should fail (empty NMResourceInfo returned).
+  @Test
-    JSONObject json = response.getEntity(JSONObject.class);
+  public void testGetNMResourceInfoEncodedIsSuccessful()
-    Assert.assertEquals(0, json.length());
+      throws YarnException, JSONException {
+    setupMockPluginsWithNmResourceInfo();
 
-    // Access resource-3 should fail (unknown plugin)
+    //test encoded yarn.io/resource-1 path
-    response = r.path("ws").path("v1").path("node").path(
+    WebResource r = resource();
-        "resources").path("resource-3").accept(MediaType.APPLICATION_JSON).get(
+    ClientResponse response = getNMResourceResponse(r, "yarn.io%2Fresource-1");
-        ClientResponse.class);
+    assertNMResourceInfoResponse(response, NM_RESOURCE_VALUE);
-    assertEquals(MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
+  }
-        response.getType().toString());
-    json = response.getEntity(JSONObject.class);
-    Assert.assertEquals(0, json.length());
 
-    // Access resource-1 should success
+  @Test
-    response = r.path("ws").path("v1").path("node").path(
+  public void testGetNMResourceInfoFailBecauseOfEmptyResourceInfo()
-        "resources").path("resource-1").accept(MediaType.APPLICATION_JSON).get(
+      throws YarnException {
-        ClientResponse.class);
+    setupMockPluginsWithNmResourceInfo();
-    assertEquals(MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
-        response.getType().toString());
-    json = response.getEntity(JSONObject.class);
-    Assert.assertEquals(1000, json.get("a"));
 
-    // Access resource-1 should success (encoded yarn.io/Fresource-1).
+    WebResource r = resource();
-    response = r.path("ws").path("v1").path("node").path("resources").path(
+    ClientResponse response = getNMResourceResponse(r, "resource-2");
-        "yarn.io%2Fresource-1").accept(MediaType.APPLICATION_JSON).get(
+    assertEmptyNMResourceInfo(response);
-        ClientResponse.class);
+  }
-    assertEquals(MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
+
-        response.getType().toString());
+  @Test
-    json = response.getEntity(JSONObject.class);
+  public void testGetNMResourceInfoWhenPluginIsUnknown()
-    Assert.assertEquals(1000, json.get("a"));
+      throws YarnException {
+    setupMockPluginsWithNmResourceInfo();
+
+    WebResource r = resource();
+    ClientResponse response = getNMResourceResponse(r, "resource-3");
+    assertEmptyNMResourceInfo(response);
   }
 
   private ContainerId createContainerId(int id) {
     ApplicationId appId = ApplicationId.newInstance(0, 0);
     ApplicationAttemptId appAttemptId =
         ApplicationAttemptId.newInstance(appId, 1);
-    ContainerId containerId = ContainerId.newContainerId(appAttemptId, id);
+    return ContainerId.newContainerId(appAttemptId, id);
-    return containerId;
   }
 
   @Test
   public void testGetYarnGpuResourceInfo()
-      throws YarnException, InterruptedException, JSONException {
+      throws YarnException, JSONException {
-    ResourcePluginManager rpm = mock(ResourcePluginManager.class);
+    setupMockPluginsWithGpuResourceInfo();
-    Map<String, ResourcePlugin> namesToPlugins = new HashMap<>();
-    ResourcePlugin mockPlugin1 = mock(ResourcePlugin.class);
-    GpuDeviceInformation gpuDeviceInformation = new GpuDeviceInformation();
-    gpuDeviceInformation.setDriverVersion("1.2.3");
-    gpuDeviceInformation.setGpus(Arrays.asList(new PerGpuDeviceInformation()));
-    NMResourceInfo nmResourceInfo1 = new NMGpuResourceInfo(gpuDeviceInformation,
-        Arrays.asList(new GpuDevice(1, 1), new GpuDevice(2, 2),
-            new GpuDevice(3, 3)), Arrays
-        .asList(new AssignedGpuDevice(2, 2, createContainerId(1)),
-            new AssignedGpuDevice(3, 3, createContainerId(2))));
-    when(mockPlugin1.getNMResourceInfo()).thenReturn(nmResourceInfo1);
-    namesToPlugins.put("resource-1", mockPlugin1);
-    namesToPlugins.put("yarn.io/resource-1", mockPlugin1);
-    ResourcePlugin mockPlugin2 = mock(ResourcePlugin.class);
-    namesToPlugins.put("resource-2", mockPlugin2);
-    when(rpm.getNameToPlugins()).thenReturn(namesToPlugins);
-
-    nmContext.setResourcePluginManager(rpm);
 
     WebResource r = resource();
-    ClientResponse response;
+    ClientResponse response = getNMResourceResponse(r, "resource-1");
-    JSONObject json;
+    assertEquals("MediaType of the response is not the expected!",
-
+        MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
-    // Access resource-1 should success
-    response = r.path("ws").path("v1").path("node").path(
-        "resources").path("resource-1").accept(MediaType.APPLICATION_JSON).get(
-        ClientResponse.class);
-    assertEquals(MediaType.APPLICATION_JSON + "; " + JettyUtils.UTF_8,
         response.getType().toString());
-    json = response.getEntity(JSONObject.class);
+    JSONObject json = response.getEntity(JSONObject.class);
-    Assert.assertEquals("1.2.3",
+    assertEquals("Unexpected driverVersion in the json response!",
+        "1.2.3",
         json.getJSONObject("gpuDeviceInformation").get("driverVersion"));
-    Assert.assertEquals(3, json.getJSONArray("totalGpuDevices").length());
+    assertEquals("Unexpected totalGpuDevices in the json response!",
-    Assert.assertEquals(2, json.getJSONArray("assignedGpuDevices").length());
+        3, json.getJSONArray("totalGpuDevices").length());
-    Assert.assertEquals(2, json.getJSONArray("assignedGpuDevices").length());
+    assertEquals("Unexpected assignedGpuDevices in the json response!",
+        2, json.getJSONArray("assignedGpuDevices").length());
   }
 
   private void testContainerLogs(WebResource r, ContainerId containerId)
@@ -713,7 +759,7 @@ private void testContainerLogs(WebResource r, ContainerId containerId)
     }
     // After container is completed, it is removed from nmContext
     nmContext.getContainers().remove(containerId);
-    Assert.assertNull(nmContext.getContainers().get(containerId));
+    assertNull(nmContext.getContainers().get(containerId));
     response =
         r.path(filename).accept(MediaType.TEXT_PLAIN)
             .get(ClientResponse.class);