diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/ParameterizedSchedulerTestBase.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/ParameterizedSchedulerTestBase.java index 3f2bfc9f941..c29e7034d83 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/ParameterizedSchedulerTestBase.java +++ b/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} + * + *

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}.

+ * + * @return The scheduler configured by + * {@code YarnConfiguration.RM_SCHEDULER} + */ + public AbstractYarnScheduler getScheduler() { + return scheduler; + } } diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/resourcetracker/TestNMReconnect.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/resourcetracker/TestNMReconnect.java index cc31f62c505..2c926d915cb 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/resourcetracker/TestNMReconnect.java +++ b/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);