Skip to content

Commit 4454286

Browse files
szilard-nemethszilard-nemeth
committed
YARN-10408. Extract MockQueueHierarchyBuilder to a separate class. Contributed by Gergely Pollak
1 parent 6e3d705 commit 4454286

File tree

2 files changed

+163
-143
lines changed

2 files changed

+163
-143
lines changed

‎hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/MockQueueHierarchyBuilder.java‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
/**
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.hadoop.yarn.server.resourcemanager.placement;
20+
21+
import com.google.common.collect.Maps;
22+
import org.apache.commons.compress.utils.Lists;
23+
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.*;
24+
25+
import java.util.List;
26+
import java.util.Map;
27+
28+
import static org.mockito.Mockito.mock;
29+
import static org.mockito.Mockito.when;
30+
31+
class MockQueueHierarchyBuilder {
32+
private static final String ROOT = "root";
33+
private static final String QUEUE_SEP = ".";
34+
private List<String> queuePaths = Lists.newArrayList();
35+
private List<String> managedParentQueues = Lists.newArrayList();
36+
private CapacitySchedulerQueueManager queueManager;
37+
38+
public static MockQueueHierarchyBuilder create() {
39+
return new MockQueueHierarchyBuilder();
40+
}
41+
42+
public MockQueueHierarchyBuilder withQueueManager(
43+
CapacitySchedulerQueueManager queueManager) {
44+
this.queueManager = queueManager;
45+
return this;
46+
}
47+
48+
public MockQueueHierarchyBuilder withQueue(String queue) {
49+
this.queuePaths.add(queue);
50+
return this;
51+
}
52+
53+
public MockQueueHierarchyBuilder withManagedParentQueue(
54+
String managedQueue) {
55+
this.managedParentQueues.add(managedQueue);
56+
return this;
57+
}
58+
59+
public void build() {
60+
if (this.queueManager == null) {
61+
throw new IllegalStateException(
62+
"QueueManager instance is not provided!");
63+
}
64+
65+
for (String managedParentQueue : managedParentQueues) {
66+
if (!queuePaths.contains(managedParentQueue)) {
67+
queuePaths.add(managedParentQueue);
68+
} else {
69+
throw new IllegalStateException("Cannot add a managed parent " +
70+
"and a simple queue with the same path");
71+
}
72+
}
73+
74+
Map<String, AbstractCSQueue> queues = Maps.newHashMap();
75+
for (String queuePath : queuePaths) {
76+
addQueues(queues, queuePath);
77+
}
78+
}
79+
80+
private void addQueues(Map<String, AbstractCSQueue> queues,
81+
String queuePath) {
82+
final String[] pathComponents = queuePath.split("\\" + QUEUE_SEP);
83+
84+
String currentQueuePath = "";
85+
for (int i = 0; i < pathComponents.length; ++i) {
86+
boolean isLeaf = i == pathComponents.length - 1;
87+
String queueName = pathComponents[i];
88+
String parentPath = currentQueuePath;
89+
currentQueuePath += currentQueuePath.equals("") ?
90+
queueName : QUEUE_SEP + queueName;
91+
92+
if (managedParentQueues.contains(parentPath) && !isLeaf) {
93+
throw new IllegalStateException("Cannot add a queue under " +
94+
"managed parent");
95+
}
96+
if (!queues.containsKey(currentQueuePath)) {
97+
ParentQueue parentQueue = (ParentQueue) queues.get(parentPath);
98+
AbstractCSQueue queue = createQueue(parentQueue, queueName,
99+
currentQueuePath, isLeaf);
100+
queues.put(currentQueuePath, queue);
101+
}
102+
}
103+
}
104+
105+
private AbstractCSQueue createQueue(ParentQueue parentQueue,
106+
String queueName, String currentQueuePath, boolean isLeaf) {
107+
if (queueName.equals(ROOT)) {
108+
return createRootQueue(ROOT);
109+
} else if (managedParentQueues.contains(currentQueuePath)) {
110+
return addManagedParentQueueAsChildOf(parentQueue, queueName);
111+
} else if (isLeaf) {
112+
return addLeafQueueAsChildOf(parentQueue, queueName);
113+
} else {
114+
return addParentQueueAsChildOf(parentQueue, queueName);
115+
}
116+
}
117+
118+
private AbstractCSQueue createRootQueue(String rootQueueName) {
119+
ParentQueue root = mock(ParentQueue.class);
120+
when(root.getQueuePath()).thenReturn(rootQueueName);
121+
when(queueManager.getQueue(rootQueueName)).thenReturn(root);
122+
when(queueManager.getQueueByFullName(rootQueueName)).thenReturn(root);
123+
return root;
124+
}
125+
126+
private AbstractCSQueue addParentQueueAsChildOf(ParentQueue parent,
127+
String queueName) {
128+
ParentQueue queue = mock(ParentQueue.class);
129+
setQueueFields(parent, queue, queueName);
130+
return queue;
131+
}
132+
133+
private AbstractCSQueue addManagedParentQueueAsChildOf(ParentQueue parent,
134+
String queueName) {
135+
ManagedParentQueue queue = mock(ManagedParentQueue.class);
136+
setQueueFields(parent, queue, queueName);
137+
return queue;
138+
}
139+
140+
private AbstractCSQueue addLeafQueueAsChildOf(ParentQueue parent,
141+
String queueName) {
142+
LeafQueue queue = mock(LeafQueue.class);
143+
setQueueFields(parent, queue, queueName);
144+
return queue;
145+
}
146+
147+
private void setQueueFields(ParentQueue parent, AbstractCSQueue newQueue,
148+
String queueName) {
149+
String fullPathOfQueue = parent.getQueuePath() + QUEUE_SEP + queueName;
150+
addQueueToQueueManager(queueName, newQueue, fullPathOfQueue);
151+
152+
when(newQueue.getParent()).thenReturn(parent);
153+
when(newQueue.getQueuePath()).thenReturn(fullPathOfQueue);
154+
when(newQueue.getQueueName()).thenReturn(queueName);
155+
}
156+
157+
private void addQueueToQueueManager(String queueName, AbstractCSQueue queue,
158+
String fullPathOfQueue) {
159+
when(queueManager.getQueue(queueName)).thenReturn(queue);
160+
when(queueManager.getQueue(fullPathOfQueue)).thenReturn(queue);
161+
when(queueManager.getQueueByFullName(fullPathOfQueue)).thenReturn(queue);
162+
}
163+
}

‎hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java‎

Lines changed: 0 additions & 143 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,7 @@
2424
import static org.mockito.Mockito.isNull;
2525

2626
import java.util.Arrays;
27-
import java.util.List;
28-
import java.util.Map;
2927

30-
import com.google.common.collect.Maps;
31-
import org.apache.commons.compress.utils.Lists;
3228
import org.apache.hadoop.fs.CommonConfigurationKeys;
3329
import org.apache.hadoop.security.GroupMappingServiceProvider;
3430
import org.apache.hadoop.security.Groups;
@@ -39,11 +35,7 @@
3935
import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.MappingType;
4036
import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.QueueMappingBuilder;
4137
import org.apache.hadoop.yarn.server.resourcemanager.placement.TestUserGroupMappingPlacementRule.QueueMappingTestData.QueueMappingTestDataBuilder;
42-
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.AbstractCSQueue;
4338
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerQueueManager;
44-
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.LeafQueue;
45-
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ManagedParentQueue;
46-
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ParentQueue;
4739
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.PrimaryGroupMapping;
4840
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.SimpleGroupsMapping;
4941
import org.apache.hadoop.yarn.util.Records;
@@ -57,141 +49,6 @@ public class TestUserGroupMappingPlacementRule {
5749
private static final Logger LOG =
5850
LoggerFactory.getLogger(TestUserGroupMappingPlacementRule.class);
5951

60-
private static class MockQueueHierarchyBuilder {
61-
private static final String ROOT = "root";
62-
private static final String QUEUE_SEP = ".";
63-
private List<String> queuePaths = Lists.newArrayList();
64-
private List<String> managedParentQueues = Lists.newArrayList();
65-
private CapacitySchedulerQueueManager queueManager;
66-
67-
public static MockQueueHierarchyBuilder create() {
68-
return new MockQueueHierarchyBuilder();
69-
}
70-
71-
public MockQueueHierarchyBuilder withQueueManager(
72-
CapacitySchedulerQueueManager queueManager) {
73-
this.queueManager = queueManager;
74-
return this;
75-
}
76-
77-
public MockQueueHierarchyBuilder withQueue(String queue) {
78-
this.queuePaths.add(queue);
79-
return this;
80-
}
81-
82-
public MockQueueHierarchyBuilder withManagedParentQueue(
83-
String managedQueue) {
84-
this.managedParentQueues.add(managedQueue);
85-
return this;
86-
}
87-
88-
public void build() {
89-
if (this.queueManager == null) {
90-
throw new IllegalStateException(
91-
"QueueManager instance is not provided!");
92-
}
93-
94-
for (String managedParentQueue : managedParentQueues) {
95-
if (!queuePaths.contains(managedParentQueue)) {
96-
queuePaths.add(managedParentQueue);
97-
} else {
98-
throw new IllegalStateException("Cannot add a managed parent " +
99-
"and a simple queue with the same path");
100-
}
101-
}
102-
103-
Map<String, AbstractCSQueue> queues = Maps.newHashMap();
104-
for (String queuePath : queuePaths) {
105-
LOG.info("Processing queue path: " + queuePath);
106-
addQueues(queues, queuePath);
107-
}
108-
}
109-
110-
private void addQueues(Map<String, AbstractCSQueue> queues,
111-
String queuePath) {
112-
final String[] pathComponents = queuePath.split("\\" + QUEUE_SEP);
113-
114-
String currentQueuePath = "";
115-
for (int i = 0; i < pathComponents.length; ++i) {
116-
boolean isLeaf = i == pathComponents.length - 1;
117-
String queueName = pathComponents[i];
118-
String parentPath = currentQueuePath;
119-
currentQueuePath += currentQueuePath.equals("") ?
120-
queueName : QUEUE_SEP + queueName;
121-
122-
if (managedParentQueues.contains(parentPath) && !isLeaf) {
123-
throw new IllegalStateException("Cannot add a queue under " +
124-
"managed parent");
125-
}
126-
if (!queues.containsKey(currentQueuePath)) {
127-
ParentQueue parentQueue = (ParentQueue) queues.get(parentPath);
128-
AbstractCSQueue queue = createQueue(parentQueue, queueName,
129-
currentQueuePath, isLeaf);
130-
queues.put(currentQueuePath, queue);
131-
}
132-
}
133-
}
134-
135-
private AbstractCSQueue createQueue(ParentQueue parentQueue,
136-
String queueName, String currentQueuePath, boolean isLeaf) {
137-
if (queueName.equals(ROOT)) {
138-
return createRootQueue(ROOT);
139-
} else if (managedParentQueues.contains(currentQueuePath)) {
140-
return addManagedParentQueueAsChildOf(parentQueue, queueName);
141-
} else if (isLeaf) {
142-
return addLeafQueueAsChildOf(parentQueue, queueName);
143-
} else {
144-
return addParentQueueAsChildOf(parentQueue, queueName);
145-
}
146-
}
147-
148-
private AbstractCSQueue createRootQueue(String rootQueueName) {
149-
ParentQueue root = mock(ParentQueue.class);
150-
when(root.getQueuePath()).thenReturn(rootQueueName);
151-
when(queueManager.getQueue(rootQueueName)).thenReturn(root);
152-
when(queueManager.getQueueByFullName(rootQueueName)).thenReturn(root);
153-
return root;
154-
}
155-
156-
private AbstractCSQueue addParentQueueAsChildOf(ParentQueue parent,
157-
String queueName) {
158-
ParentQueue queue = mock(ParentQueue.class);
159-
setQueueFields(parent, queue, queueName);
160-
return queue;
161-
}
162-
163-
private AbstractCSQueue addManagedParentQueueAsChildOf(ParentQueue parent,
164-
String queueName) {
165-
ManagedParentQueue queue = mock(ManagedParentQueue.class);
166-
setQueueFields(parent, queue, queueName);
167-
return queue;
168-
}
169-
170-
private AbstractCSQueue addLeafQueueAsChildOf(ParentQueue parent,
171-
String queueName) {
172-
LeafQueue queue = mock(LeafQueue.class);
173-
setQueueFields(parent, queue, queueName);
174-
return queue;
175-
}
176-
177-
private void setQueueFields(ParentQueue parent, AbstractCSQueue newQueue,
178-
String queueName) {
179-
String fullPathOfQueue = parent.getQueuePath() + QUEUE_SEP + queueName;
180-
addQueueToQueueManager(queueName, newQueue, fullPathOfQueue);
181-
182-
when(newQueue.getParent()).thenReturn(parent);
183-
when(newQueue.getQueuePath()).thenReturn(fullPathOfQueue);
184-
when(newQueue.getQueueName()).thenReturn(queueName);
185-
}
186-
187-
private void addQueueToQueueManager(String queueName, AbstractCSQueue queue,
188-
String fullPathOfQueue) {
189-
when(queueManager.getQueue(queueName)).thenReturn(queue);
190-
when(queueManager.getQueue(fullPathOfQueue)).thenReturn(queue);
191-
when(queueManager.getQueueByFullName(fullPathOfQueue)).thenReturn(queue);
192-
}
193-
}
194-
19552
YarnConfiguration conf = new YarnConfiguration();
19653

19754
@Before

0 commit comments

Comments
 (0)