YARN-4996. Make TestNMReconnect.testCompareRMNodeAfterReconnect() scheduler agnostic (Kai Sasaki via Varun Saxena)

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

@@ -19,6 +19,7 @@
 package org.apache.hadoop.yarn.server.resourcemanager;
 
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacityScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.FairScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.FairSchedulerConfiguration;
@@ -38,6 +39,7 @@ public abstract class ParameterizedSchedulerTestBase {
 
   private SchedulerType schedulerType;
   private YarnConfiguration conf = null;
+  private AbstractYarnScheduler scheduler = null;
 
   public enum SchedulerType {
     CAPACITY, FAIR
@@ -58,8 +60,11 @@ public abstract class ParameterizedSchedulerTestBase {
     if (schedulerClass == FairScheduler.class) {
       schedulerType = SchedulerType.FAIR;
       configureFairScheduler(conf);
+      scheduler = new FairScheduler();
     } else if (schedulerClass == CapacityScheduler.class) {
       schedulerType = SchedulerType.CAPACITY;
+      scheduler = new CapacityScheduler();
+      ((CapacityScheduler)scheduler).setConf(conf);
     }
   }
 
@@ -87,4 +92,19 @@ public abstract class ParameterizedSchedulerTestBase {
   public SchedulerType getSchedulerType() {
     return schedulerType;
   }
+
+  /**
+   * Return a scheduler configured by {@code YarnConfiguration.RM_SCHEDULER}
+   *
+   * <p>The scheduler is configured by {@link #configureScheduler()}.
+   * Client test code can obtain the scheduler with this getter method.
+   * Schedulers supported by this class are {@link FairScheduler} or
+   * {@link CapacityScheduler}. </p>
+   *
+   * @return   The scheduler configured by
+   *           {@code YarnConfiguration.RM_SCHEDULER}
+   */
+  public AbstractYarnScheduler getScheduler() {
+    return scheduler;
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/resourcetracker/TestNMReconnect.java

@@ -38,6 +38,7 @@ import org.apache.hadoop.yarn.server.api.protocolrecords.RegisterNodeManagerResp
 import org.apache.hadoop.yarn.server.api.records.NodeAction;
 import org.apache.hadoop.yarn.server.resourcemanager.MockNM;
 import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
+import org.apache.hadoop.yarn.server.resourcemanager.ParameterizedSchedulerTestBase;
 import org.apache.hadoop.yarn.server.resourcemanager.NMLivelinessMonitor;
 import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager;
 import org.apache.hadoop.yarn.server.resourcemanager.NodesListManager;
@@ -47,7 +48,7 @@ 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;
-import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacityScheduler;
+import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.SchedulerEventType;
 import org.apache.hadoop.yarn.server.resourcemanager.security.NMTokenSecretManagerInRM;
 import org.apache.hadoop.yarn.server.resourcemanager.security.RMContainerTokenSecretManager;
@@ -57,7 +58,12 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
-public class TestNMReconnect {
+/**
+ * TestNMReconnect run tests against the scheduler set by
+ * {@link ParameterizedSchedulerTestBase} which is configured
+ * in {@link YarnConfiguration}.
+ */
+public class TestNMReconnect extends ParameterizedSchedulerTestBase {
   private static final RecordFactory recordFactory = 
       RecordFactoryProvider.getRecordFactory(null);
 
@@ -147,9 +153,8 @@ public class TestNMReconnect {
 
   @Test
   public void testCompareRMNodeAfterReconnect() throws Exception {
+    AbstractYarnScheduler scheduler = getScheduler();
     Configuration yarnConf = new YarnConfiguration();
-    CapacityScheduler scheduler = new CapacityScheduler();
-    scheduler.setConf(yarnConf);
     ConfigurationProvider configurationProvider =
         ConfigurationProviderFactory.getConfigurationProvider(yarnConf);
     configurationProvider.init(yarnConf);