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.scheduler.capacity; 20 21 import static org.mockito.Matchers.any; 22 import static org.mockito.Mockito.doReturn; 23 import static org.mockito.Mockito.mock; 24 import static org.mockito.Mockito.spy; 25 import static org.mockito.Mockito.when; 26 27 import java.util.Set; 28 29 import org.apache.commons.logging.Log; 30 import org.apache.commons.logging.LogFactory; 31 import org.apache.hadoop.conf.Configuration; 32 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId; 33 import org.apache.hadoop.yarn.api.records.ApplicationId; 34 import org.apache.hadoop.yarn.api.records.Container; 35 import org.apache.hadoop.yarn.api.records.ContainerId; 36 import org.apache.hadoop.yarn.api.records.NodeId; 37 import org.apache.hadoop.yarn.api.records.Priority; 38 import org.apache.hadoop.yarn.api.records.Resource; 39 import org.apache.hadoop.yarn.api.records.ResourceRequest; 40 import org.apache.hadoop.yarn.event.Dispatcher; 41 import org.apache.hadoop.yarn.event.Event; 42 import org.apache.hadoop.yarn.event.EventHandler; 43 import org.apache.hadoop.yarn.factories.RecordFactory; 44 import org.apache.hadoop.yarn.factory.providers.RecordFactoryProvider; 45 import org.apache.hadoop.yarn.server.resourcemanager.RMContext; 46 import org.apache.hadoop.yarn.server.resourcemanager.RMContextImpl; 47 import org.apache.hadoop.yarn.server.resourcemanager.ahs.RMApplicationHistoryWriter; 48 import org.apache.hadoop.yarn.server.resourcemanager.metrics.SystemMetricsPublisher; 49 import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager; 50 import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.ContainerAllocationExpirer; 51 import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNode; 52 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerApp; 53 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerNode; 54 import org.apache.hadoop.yarn.server.resourcemanager.security.AMRMTokenSecretManager; 55 import org.apache.hadoop.yarn.server.resourcemanager.security.ClientToAMTokenSecretManagerInRM; 56 import org.apache.hadoop.yarn.server.resourcemanager.security.NMTokenSecretManagerInRM; 57 import org.apache.hadoop.yarn.server.resourcemanager.security.RMContainerTokenSecretManager; 58 import org.apache.hadoop.yarn.server.utils.BuilderUtils; 59 import org.apache.hadoop.yarn.util.resource.Resources; 60 import org.mockito.invocation.InvocationOnMock; 61 import org.mockito.stubbing.Answer; 62 63 public class TestUtils { 64 private static final Log LOG = LogFactory.getLog(TestUtils.class); 65 66 /** 67 * Get a mock {@link RMContext} for use in test cases. 68 * @return a mock {@link RMContext} for use in test cases 69 */ 70 @SuppressWarnings({ "rawtypes", "unchecked" }) getMockRMContext()71 public static RMContext getMockRMContext() { 72 // Null dispatcher 73 Dispatcher nullDispatcher = new Dispatcher() { 74 private final EventHandler handler = 75 new EventHandler() { 76 @Override 77 public void handle(Event event) { 78 } 79 }; 80 @Override 81 public void register(Class<? extends Enum> eventType, 82 EventHandler handler) { 83 } 84 @Override 85 public EventHandler getEventHandler() { 86 return handler; 87 } 88 }; 89 90 // No op 91 ContainerAllocationExpirer cae = 92 new ContainerAllocationExpirer(nullDispatcher); 93 94 Configuration conf = new Configuration(); 95 RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class); 96 RMContextImpl rmContext = 97 new RMContextImpl(nullDispatcher, cae, null, null, null, 98 new AMRMTokenSecretManager(conf, null), 99 new RMContainerTokenSecretManager(conf), 100 new NMTokenSecretManagerInRM(conf), 101 new ClientToAMTokenSecretManagerInRM()); 102 RMNodeLabelsManager nlm = mock(RMNodeLabelsManager.class); 103 when( 104 nlm.getQueueResource(any(String.class), any(Set.class), 105 any(Resource.class))).thenAnswer(new Answer<Resource>() { 106 @Override 107 public Resource answer(InvocationOnMock invocation) throws Throwable { 108 Object[] args = invocation.getArguments(); 109 return (Resource) args[2]; 110 } 111 }); 112 113 when(nlm.getResourceByLabel(any(String.class), any(Resource.class))) 114 .thenAnswer(new Answer<Resource>() { 115 @Override public Resource answer(InvocationOnMock invocation) 116 throws Throwable { 117 Object[] args = invocation.getArguments(); 118 return (Resource) args[1]; 119 } 120 }); 121 122 rmContext.setNodeLabelManager(nlm); 123 rmContext.setSystemMetricsPublisher(mock(SystemMetricsPublisher.class)); 124 rmContext.setRMApplicationHistoryWriter(mock(RMApplicationHistoryWriter.class)); 125 return rmContext; 126 } 127 128 /** 129 * Hook to spy on queues. 130 */ 131 static class SpyHook extends CapacityScheduler.QueueHook { 132 @Override hook(CSQueue queue)133 public CSQueue hook(CSQueue queue) { 134 return spy(queue); 135 } 136 } 137 public static SpyHook spyHook = new SpyHook(); 138 139 private static final RecordFactory recordFactory = 140 RecordFactoryProvider.getRecordFactory(null); 141 createMockPriority( int priority)142 public static Priority createMockPriority( int priority) { 143 // Priority p = mock(Priority.class); 144 // when(p.getPriority()).thenReturn(priority); 145 Priority p = recordFactory.newRecordInstance(Priority.class); 146 p.setPriority(priority); 147 return p; 148 } 149 createResourceRequest( String resourceName, int memory, int numContainers, boolean relaxLocality, Priority priority, RecordFactory recordFactory)150 public static ResourceRequest createResourceRequest( 151 String resourceName, int memory, int numContainers, boolean relaxLocality, 152 Priority priority, RecordFactory recordFactory) { 153 ResourceRequest request = 154 recordFactory.newRecordInstance(ResourceRequest.class); 155 Resource capability = Resources.createResource(memory, 1); 156 157 request.setNumContainers(numContainers); 158 request.setResourceName(resourceName); 159 request.setCapability(capability); 160 request.setRelaxLocality(relaxLocality); 161 request.setPriority(priority); 162 return request; 163 } 164 getMockApplicationId(int appId)165 public static ApplicationId getMockApplicationId(int appId) { 166 ApplicationId applicationId = mock(ApplicationId.class); 167 when(applicationId.getClusterTimestamp()).thenReturn(0L); 168 when(applicationId.getId()).thenReturn(appId); 169 return applicationId; 170 } 171 172 public static ApplicationAttemptId getMockApplicationAttemptId(int appId, int attemptId)173 getMockApplicationAttemptId(int appId, int attemptId) { 174 ApplicationId applicationId = BuilderUtils.newApplicationId(0l, appId); 175 ApplicationAttemptId applicationAttemptId = mock(ApplicationAttemptId.class); 176 when(applicationAttemptId.getApplicationId()).thenReturn(applicationId); 177 when(applicationAttemptId.getAttemptId()).thenReturn(attemptId); 178 return applicationAttemptId; 179 } 180 getMockNode( String host, String rack, int port, int capability)181 public static FiCaSchedulerNode getMockNode( 182 String host, String rack, int port, int capability) { 183 NodeId nodeId = mock(NodeId.class); 184 when(nodeId.getHost()).thenReturn(host); 185 when(nodeId.getPort()).thenReturn(port); 186 RMNode rmNode = mock(RMNode.class); 187 when(rmNode.getNodeID()).thenReturn(nodeId); 188 when(rmNode.getTotalCapability()).thenReturn( 189 Resources.createResource(capability, 1)); 190 when(rmNode.getNodeAddress()).thenReturn(host+":"+port); 191 when(rmNode.getHostName()).thenReturn(host); 192 when(rmNode.getRackName()).thenReturn(rack); 193 194 FiCaSchedulerNode node = spy(new FiCaSchedulerNode(rmNode, false)); 195 LOG.info("node = " + host + " avail=" + node.getAvailableResource()); 196 return node; 197 } 198 199 @SuppressWarnings("deprecation") getMockContainerId(FiCaSchedulerApp application)200 public static ContainerId getMockContainerId(FiCaSchedulerApp application) { 201 ContainerId containerId = mock(ContainerId.class); 202 doReturn(application.getApplicationAttemptId()). 203 when(containerId).getApplicationAttemptId(); 204 long id = application.getNewContainerId(); 205 doReturn((int)id).when(containerId).getId(); 206 doReturn(id).when(containerId).getContainerId(); 207 return containerId; 208 } 209 getMockContainer( ContainerId containerId, NodeId nodeId, Resource resource, Priority priority)210 public static Container getMockContainer( 211 ContainerId containerId, NodeId nodeId, 212 Resource resource, Priority priority) { 213 Container container = mock(Container.class); 214 when(container.getId()).thenReturn(containerId); 215 when(container.getNodeId()).thenReturn(nodeId); 216 when(container.getResource()).thenReturn(resource); 217 when(container.getPriority()).thenReturn(priority); 218 return container; 219 } 220 } 221