From 0ddb48a3038d9c5132efd9f48110f9d23c53ac67 Mon Sep 17 00:00:00 2001
From: Szilard Nemeth <snemeth@apache.org>
Date: Wed, 9 Oct 2019 13:09:09 +0200
Subject: [PATCH] YARN-9356. Add more tests to ratio method in
 TestResourceCalculator. Contributed by Zoltan Siegl

(cherry picked from commit 35f093f5b3f310af49a8bb476e62173862f5af51)
---
 .../util/resource/TestResourceCalculator.java | 47 +++++++++++++++++--
 1 file changed, 44 insertions(+), 3 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceCalculator.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceCalculator.java
index 58f40be5b8b..1a7599fe309 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceCalculator.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceCalculator.java
@@ -34,6 +34,8 @@ import static org.junit.Assert.assertEquals;
 
 @RunWith(Parameterized.class)
 public class TestResourceCalculator {
+  private static final String EXTRA_RESOURCE_NAME = "test";
+
   private final ResourceCalculator resourceCalculator;
 
   @Parameterized.Parameters
@@ -53,7 +55,7 @@ public class TestResourceCalculator {
   private static void setupExtraResource() {
     Configuration conf = new Configuration();
 
-    conf.set(YarnConfiguration.RESOURCE_TYPES, "test");
+    conf.set(YarnConfiguration.RESOURCE_TYPES, EXTRA_RESOURCE_NAME);
     ResourceUtils.resetResourceTypes(conf);
   }
 
@@ -95,10 +97,10 @@ public class TestResourceCalculator {
     return res;
   }
 
-  private Resource newResource(long memory, int cpu, int test) {
+  private Resource newResource(long memory, int cpu, int extraResource) {
     Resource res = newResource(memory, cpu);
 
-    res.setResourceValue("test", test);
+    res.setResourceValue(EXTRA_RESOURCE_NAME, extraResource);
 
     return res;
   }
@@ -415,4 +417,43 @@ public class TestResourceCalculator {
       assertEquals(2, result.getVirtualCores());
     }
   }
+
+  @Test
+  public void testRatioWithNoExtraResource() {
+    //setup
+    Resource resource1 = newResource(1, 1);
+    Resource resource2 = newResource(2, 1);
+
+    //act
+    float ratio = resourceCalculator.ratio(resource1, resource2);
+
+    //assert
+    if (resourceCalculator instanceof DefaultResourceCalculator) {
+      double ratioOfMemories = 0.5;
+      assertEquals(ratioOfMemories, ratio, 0.00001);
+    } else if (resourceCalculator instanceof DominantResourceCalculator) {
+      double ratioOfCPUs = 1.0;
+      assertEquals(ratioOfCPUs, ratio, 0.00001);
+    }
+  }
+
+  @Test
+  public void testRatioWithExtraResource() {
+    //setup
+    setupExtraResource();
+    Resource resource1 = newResource(1, 1, 2);
+    Resource resource2 = newResource(2, 1, 1);
+
+    //act
+    float ratio = resourceCalculator.ratio(resource1, resource2);
+
+    //assert
+    if (resourceCalculator instanceof DefaultResourceCalculator) {
+      double ratioOfMemories = 0.5;
+      assertEquals(ratioOfMemories, ratio, 0.00001);
+    } else if (resourceCalculator instanceof DominantResourceCalculator) {
+      double ratioOfExtraResources = 2.0;
+      assertEquals(ratioOfExtraResources, ratio, 0.00001);
+    }
+  }
 }
\ No newline at end of file