diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java index 3a817e62c68..35f5c08cc20 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java @@ -165,37 +165,27 @@ public void testCapacitySchedulerInitWithCustomResourceType() resourceTypesFile = new File(source.getParent(), "resource-types.xml"); FileUtils.copyFile(source, resourceTypesFile); - CapacityScheduler cs = new CapacityScheduler(); - CapacityScheduler spyCS = spy(cs); - CapacitySchedulerConfiguration csConf = + CapacitySchedulerConfiguration newConf = (CapacitySchedulerConfiguration) TestUtils .getConfigurationWithMultipleQueues(conf); - csConf.setClass(CapacitySchedulerConfiguration.RESOURCE_CALCULATOR_CLASS, + newConf.setClass(CapacitySchedulerConfiguration.RESOURCE_CALCULATOR_CLASS, DominantResourceCalculator.class, ResourceCalculator.class); - spyCS.setConf(csConf); - - RMNodeLabelsManager nodeLabelsManager = new NullRMNodeLabelsManager(); - nodeLabelsManager.init(csConf); - PlacementManager pm = new PlacementManager(); - RMContext mockContext = mock(RMContext.class); - when(mockContext.getConfigurationProvider()).thenReturn( - new LocalConfigurationProvider()); - mockContext.setNodeLabelManager(nodeLabelsManager); - when(mockContext.getNodeLabelManager()).thenReturn(nodeLabelsManager); - when(mockContext.getQueuePlacementManager()).thenReturn(pm); - spyCS.setRMContext(mockContext); + //start RM + MockRM rm = new MockRM(newConf); + rm.start(); - spyCS.init(csConf); + CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler(); // Ensure the method can get custom resource type from // CapacitySchedulerConfiguration Assert.assertNotEquals(0, ResourceUtils - .fetchMaximumAllocationFromConfig(spyCS.getConfiguration()) + .fetchMaximumAllocationFromConfig(cs.getConfiguration()) .getResourceValue("yarn.io/gpu")); // Ensure custom resource type exists in queue's maximumAllocation Assert.assertNotEquals(0, - spyCS.getMaximumResourceCapability("a") + cs.getMaximumResourceCapability("a") .getResourceValue("yarn.io/gpu")); + rm.close(); } }