diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/MockQueueHierarchyBuilder.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/MockQueueHierarchyBuilder.java new file mode 100644 index 0000000000..e2e6ad2877 --- /dev/null +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/MockQueueHierarchyBuilder.java @@ -0,0 +1,163 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hadoop.yarn.server.resourcemanager.placement; + +import com.google.common.collect.Maps; +import org.apache.commons.compress.utils.Lists; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.*; + +import java.util.List; +import java.util.Map; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class MockQueueHierarchyBuilder { + private static final String ROOT = "root"; + private static final String QUEUE_SEP = "."; + private List queuePaths = Lists.newArrayList(); + private List managedParentQueues = Lists.newArrayList(); + private CapacitySchedulerQueueManager queueManager; + + public static MockQueueHierarchyBuilder create() { + return new MockQueueHierarchyBuilder(); + } + + public MockQueueHierarchyBuilder withQueueManager( + CapacitySchedulerQueueManager queueManager) { + this.queueManager = queueManager; + return this; + } + + public MockQueueHierarchyBuilder withQueue(String queue) { + this.queuePaths.add(queue); + return this; + } + + public MockQueueHierarchyBuilder withManagedParentQueue( + String managedQueue) { + this.managedParentQueues.add(managedQueue); + return this; + } + + public void build() { + if (this.queueManager == null) { + throw new IllegalStateException( + "QueueManager instance is not provided!"); + } + + for (String managedParentQueue : managedParentQueues) { + if (!queuePaths.contains(managedParentQueue)) { + queuePaths.add(managedParentQueue); + } else { + throw new IllegalStateException("Cannot add a managed parent " + + "and a simple queue with the same path"); + } + } + + Map queues = Maps.newHashMap(); + for (String queuePath : queuePaths) { + addQueues(queues, queuePath); + } + } + + private void addQueues(Map queues, + String queuePath) { + final String[] pathComponents = queuePath.split("\\" + QUEUE_SEP); + + String currentQueuePath = ""; + for (int i = 0; i < pathComponents.length; ++i) { + boolean isLeaf = i == pathComponents.length - 1; + String queueName = pathComponents[i]; + String parentPath = currentQueuePath; + currentQueuePath += currentQueuePath.equals("") ? + queueName : QUEUE_SEP + queueName; + + if (managedParentQueues.contains(parentPath) && !isLeaf) { + throw new IllegalStateException("Cannot add a queue under " + + "managed parent"); + } + if (!queues.containsKey(currentQueuePath)) { + ParentQueue parentQueue = (ParentQueue) queues.get(parentPath); + AbstractCSQueue queue = createQueue(parentQueue, queueName, + currentQueuePath, isLeaf); + queues.put(currentQueuePath, queue); + } + } + } + + private AbstractCSQueue createQueue(ParentQueue parentQueue, + String queueName, String currentQueuePath, boolean isLeaf) { + if (queueName.equals(ROOT)) { + return createRootQueue(ROOT); + } else if (managedParentQueues.contains(currentQueuePath)) { + return addManagedParentQueueAsChildOf(parentQueue, queueName); + } else if (isLeaf) { + return addLeafQueueAsChildOf(parentQueue, queueName); + } else { + return addParentQueueAsChildOf(parentQueue, queueName); + } + } + + private AbstractCSQueue createRootQueue(String rootQueueName) { + ParentQueue root = mock(ParentQueue.class); + when(root.getQueuePath()).thenReturn(rootQueueName); + when(queueManager.getQueue(rootQueueName)).thenReturn(root); + when(queueManager.getQueueByFullName(rootQueueName)).thenReturn(root); + return root; + } + + private AbstractCSQueue addParentQueueAsChildOf(ParentQueue parent, + String queueName) { + ParentQueue queue = mock(ParentQueue.class); + setQueueFields(parent, queue, queueName); + return queue; + } + + private AbstractCSQueue addManagedParentQueueAsChildOf(ParentQueue parent, + String queueName) { + ManagedParentQueue queue = mock(ManagedParentQueue.class); + setQueueFields(parent, queue, queueName); + return queue; + } + + private AbstractCSQueue addLeafQueueAsChildOf(ParentQueue parent, + String queueName) { + LeafQueue queue = mock(LeafQueue.class); + setQueueFields(parent, queue, queueName); + return queue; + } + + private void setQueueFields(ParentQueue parent, AbstractCSQueue newQueue, + String queueName) { + String fullPathOfQueue = parent.getQueuePath() + QUEUE_SEP + queueName; + addQueueToQueueManager(queueName, newQueue, fullPathOfQueue); + + when(newQueue.getParent()).thenReturn(parent); + when(newQueue.getQueuePath()).thenReturn(fullPathOfQueue); + when(newQueue.getQueueName()).thenReturn(queueName); + } + + private void addQueueToQueueManager(String queueName, AbstractCSQueue queue, + String fullPathOfQueue) { + when(queueManager.getQueue(queueName)).thenReturn(queue); + when(queueManager.getQueue(fullPathOfQueue)).thenReturn(queue); + when(queueManager.getQueueByFullName(fullPathOfQueue)).thenReturn(queue); + } +} diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java index 98fba14787..0b0a195e8e 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java @@ -24,11 +24,7 @@ import static org.mockito.Mockito.isNull; import java.util.Arrays; -import java.util.List; -import java.util.Map; -import com.google.common.collect.Maps; -import org.apache.commons.compress.utils.Lists; import org.apache.hadoop.fs.CommonConfigurationKeys; import org.apache.hadoop.security.GroupMappingServiceProvider; import org.apache.hadoop.security.Groups; @@ -39,11 +35,7 @@ import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.MappingType; import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.QueueMappingBuilder; import org.apache.hadoop.yarn.server.resourcemanager.placement.TestUserGroupMappingPlacementRule.QueueMappingTestData.QueueMappingTestDataBuilder; -import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.AbstractCSQueue; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerQueueManager; -import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.LeafQueue; -import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ManagedParentQueue; -import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ParentQueue; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.PrimaryGroupMapping; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.SimpleGroupsMapping; import org.apache.hadoop.yarn.util.Records; @@ -57,141 +49,6 @@ public class TestUserGroupMappingPlacementRule { private static final Logger LOG = LoggerFactory.getLogger(TestUserGroupMappingPlacementRule.class); - private static class MockQueueHierarchyBuilder { - private static final String ROOT = "root"; - private static final String QUEUE_SEP = "."; - private List queuePaths = Lists.newArrayList(); - private List managedParentQueues = Lists.newArrayList(); - private CapacitySchedulerQueueManager queueManager; - - public static MockQueueHierarchyBuilder create() { - return new MockQueueHierarchyBuilder(); - } - - public MockQueueHierarchyBuilder withQueueManager( - CapacitySchedulerQueueManager queueManager) { - this.queueManager = queueManager; - return this; - } - - public MockQueueHierarchyBuilder withQueue(String queue) { - this.queuePaths.add(queue); - return this; - } - - public MockQueueHierarchyBuilder withManagedParentQueue( - String managedQueue) { - this.managedParentQueues.add(managedQueue); - return this; - } - - public void build() { - if (this.queueManager == null) { - throw new IllegalStateException( - "QueueManager instance is not provided!"); - } - - for (String managedParentQueue : managedParentQueues) { - if (!queuePaths.contains(managedParentQueue)) { - queuePaths.add(managedParentQueue); - } else { - throw new IllegalStateException("Cannot add a managed parent " + - "and a simple queue with the same path"); - } - } - - Map queues = Maps.newHashMap(); - for (String queuePath : queuePaths) { - LOG.info("Processing queue path: " + queuePath); - addQueues(queues, queuePath); - } - } - - private void addQueues(Map queues, - String queuePath) { - final String[] pathComponents = queuePath.split("\\" + QUEUE_SEP); - - String currentQueuePath = ""; - for (int i = 0; i < pathComponents.length; ++i) { - boolean isLeaf = i == pathComponents.length - 1; - String queueName = pathComponents[i]; - String parentPath = currentQueuePath; - currentQueuePath += currentQueuePath.equals("") ? - queueName : QUEUE_SEP + queueName; - - if (managedParentQueues.contains(parentPath) && !isLeaf) { - throw new IllegalStateException("Cannot add a queue under " + - "managed parent"); - } - if (!queues.containsKey(currentQueuePath)) { - ParentQueue parentQueue = (ParentQueue) queues.get(parentPath); - AbstractCSQueue queue = createQueue(parentQueue, queueName, - currentQueuePath, isLeaf); - queues.put(currentQueuePath, queue); - } - } - } - - private AbstractCSQueue createQueue(ParentQueue parentQueue, - String queueName, String currentQueuePath, boolean isLeaf) { - if (queueName.equals(ROOT)) { - return createRootQueue(ROOT); - } else if (managedParentQueues.contains(currentQueuePath)) { - return addManagedParentQueueAsChildOf(parentQueue, queueName); - } else if (isLeaf) { - return addLeafQueueAsChildOf(parentQueue, queueName); - } else { - return addParentQueueAsChildOf(parentQueue, queueName); - } - } - - private AbstractCSQueue createRootQueue(String rootQueueName) { - ParentQueue root = mock(ParentQueue.class); - when(root.getQueuePath()).thenReturn(rootQueueName); - when(queueManager.getQueue(rootQueueName)).thenReturn(root); - when(queueManager.getQueueByFullName(rootQueueName)).thenReturn(root); - return root; - } - - private AbstractCSQueue addParentQueueAsChildOf(ParentQueue parent, - String queueName) { - ParentQueue queue = mock(ParentQueue.class); - setQueueFields(parent, queue, queueName); - return queue; - } - - private AbstractCSQueue addManagedParentQueueAsChildOf(ParentQueue parent, - String queueName) { - ManagedParentQueue queue = mock(ManagedParentQueue.class); - setQueueFields(parent, queue, queueName); - return queue; - } - - private AbstractCSQueue addLeafQueueAsChildOf(ParentQueue parent, - String queueName) { - LeafQueue queue = mock(LeafQueue.class); - setQueueFields(parent, queue, queueName); - return queue; - } - - private void setQueueFields(ParentQueue parent, AbstractCSQueue newQueue, - String queueName) { - String fullPathOfQueue = parent.getQueuePath() + QUEUE_SEP + queueName; - addQueueToQueueManager(queueName, newQueue, fullPathOfQueue); - - when(newQueue.getParent()).thenReturn(parent); - when(newQueue.getQueuePath()).thenReturn(fullPathOfQueue); - when(newQueue.getQueueName()).thenReturn(queueName); - } - - private void addQueueToQueueManager(String queueName, AbstractCSQueue queue, - String fullPathOfQueue) { - when(queueManager.getQueue(queueName)).thenReturn(queue); - when(queueManager.getQueue(fullPathOfQueue)).thenReturn(queue); - when(queueManager.getQueueByFullName(fullPathOfQueue)).thenReturn(queue); - } - } - YarnConfiguration conf = new YarnConfiguration(); @Before