YARN-9134. No test coverage for redefining FPGA / GPU resource types in TestResourceUtils. Contributed by Peter Bacsko

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceUtils.java

@@ -28,9 +28,15 @@ import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.File;
+import java.io.IOException;
+import java.net.URL;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -38,6 +44,8 @@ import java.util.Map;
  * Test class to verify all resource utility methods.
  */
 public class TestResourceUtils {
+  private static final Logger LOG =
+      LoggerFactory.getLogger(TestResourceUtils.class);
 
   private File nodeResourcesFile;
   private File resourceTypesFile;
@@ -54,6 +62,9 @@ public class TestResourceUtils {
     }
   }
 
+  @Rule
+  public ExpectedException expexted = ExpectedException.none();
+
   public static void addNewTypesToResources(String... resourceTypes) {
     // Initialize resource map
     Map<String, ResourceInformation> riMap = new HashMap<>();
@@ -78,14 +89,60 @@ public class TestResourceUtils {
 
   @After
   public void teardown() {
-    if(nodeResourcesFile != null && nodeResourcesFile.exists()) {
+    if (nodeResourcesFile != null && nodeResourcesFile.exists()) {
       nodeResourcesFile.delete();
     }
-    if(resourceTypesFile != null && resourceTypesFile.exists()) {
+    if (resourceTypesFile != null && resourceTypesFile.exists()) {
       resourceTypesFile.delete();
     }
   }
 
+  public static String setupResourceTypes(Configuration conf, String filename)
+      throws Exception {
+    File source = new File(
+        conf.getClassLoader().getResource(filename).getFile());
+    File dest = new File(source.getParent(), "resource-types.xml");
+    FileUtils.copyFile(source, dest);
+    try {
+      ResourceUtils.getResourceTypes();
+    } catch (Exception e) {
+      if (!dest.delete()) {
+        LOG.error("Could not delete {}", dest);
+      }
+      throw e;
+    }
+    return dest.getAbsolutePath();
+  }
+
+  private Map<String, ResourceInformation> setupResourceTypesInternal(
+      Configuration conf, String srcFileName) throws IOException {
+    URL srcFileUrl = conf.getClassLoader().getResource(srcFileName);
+    if (srcFileUrl == null) {
+      throw new IllegalArgumentException(
+          "Source file does not exist: " + srcFileName);
+    }
+    File source = new File(srcFileUrl.getFile());
+    File dest = new File(source.getParent(), "resource-types.xml");
+    FileUtils.copyFile(source, dest);
+    this.resourceTypesFile = dest;
+    return ResourceUtils.getResourceTypes();
+  }
+
+  private Map<String, ResourceInformation> setupNodeResources(
+      Configuration conf, String srcFileName) throws IOException {
+    URL srcFileUrl = conf.getClassLoader().getResource(srcFileName);
+    if (srcFileUrl == null) {
+      throw new IllegalArgumentException(
+          "Source file does not exist: " + srcFileName);
+    }
+    File source = new File(srcFileUrl.getFile());
+    File dest = new File(source.getParent(), "node-resources.xml");
+    FileUtils.copyFile(source, dest);
+    this.nodeResourcesFile = dest;
+    return ResourceUtils
+        .getNodeResourceInformation(conf);
+  }
+
   private void testMemoryAndVcores(Map<String, ResourceInformation> res) {
     String memory = ResourceInformation.MEMORY_MB.getName();
     String vcores = ResourceInformation.VCORES.getName();
@@ -104,8 +161,7 @@ public class TestResourceUtils {
   }
 
   @Test
-  public void testGetResourceTypes() throws Exception {
+  public void testGetResourceTypes() {
-
     Map<String, ResourceInformation> res = ResourceUtils.getResourceTypes();
     Assert.assertEquals(2, res.size());
     testMemoryAndVcores(res);
@@ -113,7 +169,6 @@ public class TestResourceUtils {
 
   @Test
   public void testGetResourceTypesConfigs() throws Exception {
-
     Configuration conf = new YarnConfiguration();
 
     ResourceFileInformation testFile1 =
@@ -135,16 +190,11 @@ public class TestResourceUtils {
     Map<String, ResourceInformation> res;
     for (ResourceFileInformation testInformation : tests) {
       ResourceUtils.resetResourceTypes();
-      File source = new File(
+      res = setupResourceTypesInternal(conf, testInformation.filename);
-          conf.getClassLoader().getResource(testInformation.filename)
-              .getFile());
-      resourceTypesFile = new File(source.getParent(), "resource-types.xml");
-      FileUtils.copyFile(source, resourceTypesFile);
-      res = ResourceUtils.getResourceTypes();
       testMemoryAndVcores(res);
       Assert.assertEquals(testInformation.resourceCount, res.size());
-      for (Map.Entry<String, String> entry : testInformation.resourceNameUnitsMap
+      for (Map.Entry<String, String> entry :
-          .entrySet()) {
+          testInformation.resourceNameUnitsMap.entrySet()) {
         String resourceName = entry.getKey();
         Assert.assertTrue("Missing key " + resourceName,
             res.containsKey(resourceName));
@@ -154,7 +204,7 @@ public class TestResourceUtils {
   }
 
   @Test
-  public void testGetResourceTypesConfigErrors() throws Exception {
+  public void testGetResourceTypesConfigErrors() throws IOException {
     Configuration conf = new YarnConfiguration();
 
     String[] resourceFiles = {"resource-types-error-1.xml",
@@ -163,22 +213,16 @@ public class TestResourceUtils {
     for (String resourceFile : resourceFiles) {
       ResourceUtils.resetResourceTypes();
       try {
-        File source =
+        setupResourceTypesInternal(conf, resourceFile);
-            new File(conf.getClassLoader().getResource(resourceFile).getFile());
-        resourceTypesFile = new File(source.getParent(), "resource-types.xml");
-        FileUtils.copyFile(source, resourceTypesFile);
-        ResourceUtils.getResourceTypes();
         Assert.fail("Expected error with file " + resourceFile);
-      } catch (NullPointerException ne) {
+      } catch (YarnRuntimeException | IllegalArgumentException e) {
-        throw ne;
-      } catch (Exception e) {
         //Test passed
       }
     }
   }
 
   @Test
-  public void testInitializeResourcesMap() throws Exception {
+  public void testInitializeResourcesMap() {
     String[] empty = {"", ""};
     String[] res1 = {"resource1", "m"};
     String[] res2 = {"resource2", "G"};
@@ -239,8 +283,7 @@ public class TestResourceUtils {
   }
 
   @Test
-  public void testInitializeResourcesMapErrors() throws Exception {
+  public void testInitializeResourcesMapErrors() {
-
     String[] mem1 = {"memory-mb", ""};
     String[] vcores1 = {"vcores", "M"};
 
@@ -280,11 +323,9 @@ public class TestResourceUtils {
 
   @Test
   public void testGetResourceInformation() throws Exception {
-
     Configuration conf = new YarnConfiguration();
     Map<String, Resource> testRun = new HashMap<>();
-    setupResourceTypes(conf, "resource-types-4.xml");
+    setupResourceTypesInternal(conf, "resource-types-4.xml");
-    // testRun.put("node-resources-1.xml", Resource.newInstance(1024, 1));
     Resource test3Resources = Resource.newInstance(0, 0);
     test3Resources.setResourceInformation("resource1",
         ResourceInformation.newInstance("resource1", "Gi", 5L));
@@ -297,12 +338,8 @@ public class TestResourceUtils {
     for (Map.Entry<String, Resource> entry : testRun.entrySet()) {
       String resourceFile = entry.getKey();
       ResourceUtils.resetNodeResources();
-      File source = new File(
+      Map<String, ResourceInformation> actual = setupNodeResources(conf,
-          conf.getClassLoader().getResource(resourceFile).getFile());
+          resourceFile);
-      nodeResourcesFile = new File(source.getParent(), "node-resources.xml");
-      FileUtils.copyFile(source, nodeResourcesFile);
-      Map<String, ResourceInformation> actual = ResourceUtils
-          .getNodeResourceInformation(conf);
       Assert.assertEquals(actual.size(),
           entry.getValue().getResources().length);
       for (ResourceInformation resInfo : entry.getValue().getResources()) {
@@ -314,28 +351,40 @@ public class TestResourceUtils {
   @Test
   public void testGetNodeResourcesConfigErrors() throws Exception {
     Configuration conf = new YarnConfiguration();
-    Map<String, Resource> testRun = new HashMap<>();
+    setupResourceTypesInternal(conf, "resource-types-4.xml");
-    setupResourceTypes(conf, "resource-types-4.xml");
+    String[] invalidNodeResFiles = {"node-resources-error-1.xml"};
-    String invalidNodeResFiles[] = { "node-resources-error-1.xml"};
 
     for (String resourceFile : invalidNodeResFiles) {
       ResourceUtils.resetNodeResources();
       try {
-        File source = new File(conf.getClassLoader().getResource(resourceFile).getFile());
+        setupNodeResources(conf, resourceFile);
-        nodeResourcesFile = new File(source.getParent(), "node-resources.xml");
-        FileUtils.copyFile(source, nodeResourcesFile);
-        Map<String, ResourceInformation> actual = ResourceUtils.getNodeResourceInformation(conf);
         Assert.fail("Expected error with file " + resourceFile);
-      } catch (NullPointerException ne) {
+      } catch (YarnRuntimeException e) {
-        throw ne;
-      } catch (Exception e) {
         //Test passed
       }
     }
   }
 
   @Test
-  public void testResourceNameFormatValidation() throws Exception {
+  public void testGetNodeResourcesRedefineFpgaErrors() throws Exception {
+    Configuration conf = new YarnConfiguration();
+    expexted.expect(YarnRuntimeException.class);
+    expexted.expectMessage("Defined mandatory resource type=yarn.io/fpga");
+    setupResourceTypesInternal(conf,
+        "resource-types-error-redefine-fpga-unit.xml");
+  }
+
+  @Test
+  public void testGetNodeResourcesRedefineGpuErrors() throws Exception {
+    Configuration conf = new YarnConfiguration();
+    expexted.expect(YarnRuntimeException.class);
+    expexted.expectMessage("Defined mandatory resource type=yarn.io/gpu");
+    setupResourceTypesInternal(conf,
+        "resource-types-error-redefine-gpu-unit.xml");
+  }
+
+  @Test
+  public void testResourceNameFormatValidation() {
     String[] validNames = new String[] {
         "yarn.io/gpu",
         "gpu",
@@ -372,10 +421,9 @@ public class TestResourceUtils {
 
   @Test
   public void testGetResourceInformationWithDiffUnits() throws Exception {
-
     Configuration conf = new YarnConfiguration();
     Map<String, Resource> testRun = new HashMap<>();
-    setupResourceTypes(conf, "resource-types-4.xml");
+    setupResourceTypesInternal(conf, "resource-types-4.xml");
     Resource test3Resources = Resource.newInstance(0, 0);
 
     //Resource 'resource1' has been passed as 5T
@@ -394,12 +442,8 @@ public class TestResourceUtils {
     for (Map.Entry<String, Resource> entry : testRun.entrySet()) {
       String resourceFile = entry.getKey();
       ResourceUtils.resetNodeResources();
-      File source = new File(
+      Map<String, ResourceInformation> actual = setupNodeResources(conf,
-          conf.getClassLoader().getResource(resourceFile).getFile());
+          resourceFile);
-      nodeResourcesFile = new File(source.getParent(), "node-resources.xml");
-      FileUtils.copyFile(source, nodeResourcesFile);
-      Map<String, ResourceInformation> actual = ResourceUtils
-          .getNodeResourceInformation(conf);
       Assert.assertEquals(actual.size(),
           entry.getValue().getResources().length);
       for (ResourceInformation resInfo : entry.getValue().getResources()) {
@@ -407,14 +451,4 @@ public class TestResourceUtils {
       }
     }
   }
-
-  public static String setupResourceTypes(Configuration conf, String filename)
-      throws Exception {
-    File source = new File(
-        conf.getClassLoader().getResource(filename).getFile());
-    File dest = new File(source.getParent(), "resource-types.xml");
-    FileUtils.copyFile(source, dest);
-    ResourceUtils.getResourceTypes();
-    return dest.getAbsolutePath();
-  }
 }