List of usage examples for org.apache.hadoop.security Credentials Credentials
public Credentials()
From source file:org.apache.tez.dag.api.DAG.java
License:Apache License
@Private public synchronized DAGPlan createDag(Configuration tezConf, Credentials extraCredentials, Map<String, LocalResource> tezJarResources, LocalResource binaryConfig, boolean tezLrsAsArchive, Map<String, String> additionalConfigs) { verify(true);/* w ww.j a va2 s. co m*/ DAGPlan.Builder dagBuilder = DAGPlan.newBuilder(); dagBuilder.setName(this.name); if (this.dagInfo != null && !this.dagInfo.isEmpty()) { dagBuilder.setDagInfo(this.dagInfo); } if (!vertexGroups.isEmpty()) { for (VertexGroup av : vertexGroups) { GroupInfo groupInfo = av.getGroupInfo(); PlanVertexGroupInfo.Builder groupBuilder = PlanVertexGroupInfo.newBuilder(); groupBuilder.setGroupName(groupInfo.getGroupName()); for (Vertex v : groupInfo.getMembers()) { groupBuilder.addGroupMembers(v.getName()); } groupBuilder.addAllOutputs(groupInfo.outputs); for (Map.Entry<String, InputDescriptor> entry : groupInfo.edgeMergedInputs.entrySet()) { groupBuilder.addEdgeMergedInputs( PlanGroupInputEdgeInfo.newBuilder().setDestVertexName(entry.getKey()) .setMergedInput(DagTypeConverters.convertToDAGPlan(entry.getValue()))); } dagBuilder.addVertexGroups(groupBuilder); } } Credentials dagCredentials = new Credentials(); if (extraCredentials != null) { dagCredentials.mergeAll(extraCredentials); } dagCredentials.mergeAll(credentials); if (!commonTaskLocalFiles.isEmpty()) { dagBuilder.addAllLocalResource(DagTypeConverters.convertToDAGPlan(commonTaskLocalFiles)); } Preconditions.checkArgument(topologicalVertexStack.size() == vertices.size(), "size of topologicalVertexStack is:" + topologicalVertexStack.size() + " while size of vertices is:" + vertices.size() + ", make sure they are the same in order to sort the vertices"); while (!topologicalVertexStack.isEmpty()) { Vertex vertex = vertices.get(topologicalVertexStack.pop()); // infer credentials, resources and parallelism from data source Resource vertexTaskResource = vertex.getTaskResource(); if (vertexTaskResource == null) { vertexTaskResource = Resource.newInstance( tezConf.getInt(TezConfiguration.TEZ_TASK_RESOURCE_MEMORY_MB, TezConfiguration.TEZ_TASK_RESOURCE_MEMORY_MB_DEFAULT), tezConf.getInt(TezConfiguration.TEZ_TASK_RESOURCE_CPU_VCORES, TezConfiguration.TEZ_TASK_RESOURCE_CPU_VCORES_DEFAULT)); } Map<String, LocalResource> vertexLRs = Maps.newHashMap(); vertexLRs.putAll(vertex.getTaskLocalFiles()); List<DataSourceDescriptor> dataSources = vertex.getDataSources(); for (DataSourceDescriptor dataSource : dataSources) { if (dataSource.getCredentials() != null) { dagCredentials.addAll(dataSource.getCredentials()); } if (dataSource.getAdditionalLocalFiles() != null) { TezCommonUtils.addAdditionalLocalResources(dataSource.getAdditionalLocalFiles(), vertexLRs, "Vertex " + vertex.getName()); } } if (tezJarResources != null) { TezCommonUtils.addAdditionalLocalResources(tezJarResources, vertexLRs, "Vertex " + vertex.getName()); } if (binaryConfig != null) { vertexLRs.put(TezConstants.TEZ_PB_BINARY_CONF_NAME, binaryConfig); } int vertexParallelism = vertex.getParallelism(); VertexLocationHint vertexLocationHint = vertex.getLocationHint(); if (dataSources.size() == 1) { DataSourceDescriptor dataSource = dataSources.get(0); if (vertexParallelism == -1 && dataSource.getNumberOfShards() > -1) { vertexParallelism = dataSource.getNumberOfShards(); } if (vertexLocationHint == null && dataSource.getLocationHint() != null) { vertexLocationHint = dataSource.getLocationHint(); } } if (vertexParallelism == -1) { Preconditions.checkState(vertexLocationHint == null, "Cannot specify vertex location hint without specifying vertex parallelism. Vertex: " + vertex.getName()); } else if (vertexLocationHint != null) { Preconditions.checkState(vertexParallelism == vertexLocationHint.getTaskLocationHints().size(), "vertex task location hint must equal vertex parallelism. Vertex: " + vertex.getName()); } for (DataSinkDescriptor dataSink : vertex.getDataSinks()) { if (dataSink.getCredentials() != null) { dagCredentials.addAll(dataSink.getCredentials()); } } VertexPlan.Builder vertexBuilder = VertexPlan.newBuilder(); vertexBuilder.setName(vertex.getName()); vertexBuilder.setType(PlanVertexType.NORMAL); // vertex type is implicitly NORMAL until TEZ-46. vertexBuilder .setProcessorDescriptor(DagTypeConverters.convertToDAGPlan(vertex.getProcessorDescriptor())); if (vertex.getInputs().size() > 0) { for (RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor> input : vertex.getInputs()) { vertexBuilder.addInputs(DagTypeConverters.convertToDAGPlan(input)); } } if (vertex.getOutputs().size() > 0) { for (RootInputLeafOutput<OutputDescriptor, OutputCommitterDescriptor> output : vertex .getOutputs()) { vertexBuilder.addOutputs(DagTypeConverters.convertToDAGPlan(output)); } } if (vertex.getConf() != null && vertex.getConf().size() > 0) { ConfigurationProto.Builder confBuilder = ConfigurationProto.newBuilder(); for (Map.Entry<String, String> entry : vertex.getConf().entrySet()) { PlanKeyValuePair.Builder keyValueBuilder = PlanKeyValuePair.newBuilder(); keyValueBuilder.setKey(entry.getKey()); keyValueBuilder.setValue(entry.getValue()); confBuilder.addConfKeyValues(keyValueBuilder); } vertexBuilder.setVertexConf(confBuilder); } //task config PlanTaskConfiguration.Builder taskConfigBuilder = PlanTaskConfiguration.newBuilder(); taskConfigBuilder.setNumTasks(vertexParallelism); taskConfigBuilder.setMemoryMb(vertexTaskResource.getMemory()); taskConfigBuilder.setVirtualCores(vertexTaskResource.getVirtualCores()); taskConfigBuilder.setJavaOpts( TezClientUtils.addDefaultsToTaskLaunchCmdOpts(vertex.getTaskLaunchCmdOpts(), tezConf)); taskConfigBuilder.setTaskModule(vertex.getName()); if (!vertexLRs.isEmpty()) { taskConfigBuilder.addAllLocalResource(DagTypeConverters.convertToDAGPlan(vertexLRs)); } Map<String, String> taskEnv = Maps.newHashMap(vertex.getTaskEnvironment()); TezYARNUtils.setupDefaultEnv(taskEnv, tezConf, TezConfiguration.TEZ_TASK_LAUNCH_ENV, TezConfiguration.TEZ_TASK_LAUNCH_ENV_DEFAULT, tezLrsAsArchive); for (Map.Entry<String, String> entry : taskEnv.entrySet()) { PlanKeyValuePair.Builder envSettingBuilder = PlanKeyValuePair.newBuilder(); envSettingBuilder.setKey(entry.getKey()); envSettingBuilder.setValue(entry.getValue()); taskConfigBuilder.addEnvironmentSetting(envSettingBuilder); } if (vertexLocationHint != null) { if (vertexLocationHint.getTaskLocationHints() != null) { for (TaskLocationHint hint : vertexLocationHint.getTaskLocationHints()) { PlanTaskLocationHint.Builder taskLocationHintBuilder = PlanTaskLocationHint.newBuilder(); // we can allow this later on if needed if (hint.getAffinitizedTask() != null) { throw new TezUncheckedException( "Task based affinity may not be specified via the DAG API"); } if (hint.getHosts() != null) { taskLocationHintBuilder.addAllHost(hint.getHosts()); } if (hint.getRacks() != null) { taskLocationHintBuilder.addAllRack(hint.getRacks()); } vertexBuilder.addTaskLocationHint(taskLocationHintBuilder); } } } if (vertex.getVertexManagerPlugin() != null) { vertexBuilder.setVertexManagerPlugin( DagTypeConverters.convertToDAGPlan(vertex.getVertexManagerPlugin())); } for (Edge inEdge : vertex.getInputEdges()) { vertexBuilder.addInEdgeId(inEdge.getId()); } for (Edge outEdge : vertex.getOutputEdges()) { vertexBuilder.addOutEdgeId(outEdge.getId()); } vertexBuilder.setTaskConfig(taskConfigBuilder); dagBuilder.addVertex(vertexBuilder); } for (Edge edge : edges) { EdgePlan.Builder edgeBuilder = EdgePlan.newBuilder(); edgeBuilder.setId(edge.getId()); edgeBuilder.setInputVertexName(edge.getInputVertex().getName()); edgeBuilder.setOutputVertexName(edge.getOutputVertex().getName()); edgeBuilder.setDataMovementType( DagTypeConverters.convertToDAGPlan(edge.getEdgeProperty().getDataMovementType())); edgeBuilder.setDataSourceType( DagTypeConverters.convertToDAGPlan(edge.getEdgeProperty().getDataSourceType())); edgeBuilder.setSchedulingType( DagTypeConverters.convertToDAGPlan(edge.getEdgeProperty().getSchedulingType())); edgeBuilder.setEdgeSource(DagTypeConverters.convertToDAGPlan(edge.getEdgeProperty().getEdgeSource())); edgeBuilder.setEdgeDestination( DagTypeConverters.convertToDAGPlan(edge.getEdgeProperty().getEdgeDestination())); if (edge.getEdgeProperty().getDataMovementType() == DataMovementType.CUSTOM) { if (edge.getEdgeProperty().getEdgeManagerDescriptor() != null) { edgeBuilder.setEdgeManager( DagTypeConverters.convertToDAGPlan(edge.getEdgeProperty().getEdgeManagerDescriptor())); } // else the AM will deal with this. } dagBuilder.addEdge(edgeBuilder); } ConfigurationProto.Builder confProtoBuilder = ConfigurationProto.newBuilder(); if (dagAccessControls != null) { Configuration aclConf = new Configuration(false); dagAccessControls.serializeToConfiguration(aclConf); Iterator<Entry<String, String>> aclConfIter = aclConf.iterator(); while (aclConfIter.hasNext()) { Entry<String, String> entry = aclConfIter.next(); PlanKeyValuePair.Builder kvp = PlanKeyValuePair.newBuilder(); kvp.setKey(entry.getKey()); kvp.setValue(entry.getValue()); TezConfiguration.validateProperty(entry.getKey(), Scope.DAG); confProtoBuilder.addConfKeyValues(kvp); } } if (additionalConfigs != null && !additionalConfigs.isEmpty()) { for (Entry<String, String> entry : additionalConfigs.entrySet()) { PlanKeyValuePair.Builder kvp = PlanKeyValuePair.newBuilder(); kvp.setKey(entry.getKey()); kvp.setValue(entry.getValue()); TezConfiguration.validateProperty(entry.getKey(), Scope.DAG); confProtoBuilder.addConfKeyValues(kvp); } } if (this.dagConf != null && !this.dagConf.isEmpty()) { for (Entry<String, String> entry : this.dagConf.entrySet()) { PlanKeyValuePair.Builder kvp = PlanKeyValuePair.newBuilder(); kvp.setKey(entry.getKey()); kvp.setValue(entry.getValue()); confProtoBuilder.addConfKeyValues(kvp); } } dagBuilder.setDagConf(confProtoBuilder); if (dagCredentials != null) { dagBuilder.setCredentialsBinary(DagTypeConverters.convertCredentialsToProto(dagCredentials)); TezCommonUtils.logCredentials(LOG, dagCredentials, "dag"); } return dagBuilder.build(); }
From source file:org.apache.tez.dag.api.DagTypeConverters.java
License:Apache License
public static Credentials convertByteStringToCredentials(ByteString byteString) { if (byteString == null) { return null; }/*from w w w .j av a 2 s.co m*/ DataInputByteBuffer dib = new DataInputByteBuffer(); dib.reset(byteString.asReadOnlyByteBuffer()); Credentials credentials = new Credentials(); try { credentials.readTokenStorageStream(dib); return credentials; } catch (IOException e) { throw new TezUncheckedException("Failed to deserialize Credentials", e); } }
From source file:org.apache.tez.dag.api.TestDAGPlan.java
License:Apache License
@Test(timeout = 5000) public void testCredentialsSerde() { DAG dag = DAG.create("testDag"); ProcessorDescriptor pd1 = ProcessorDescriptor.create("processor1") .setUserPayload(UserPayload.create(ByteBuffer.wrap("processor1Bytes".getBytes()))); ProcessorDescriptor pd2 = ProcessorDescriptor.create("processor2") .setUserPayload(UserPayload.create(ByteBuffer.wrap("processor2Bytes".getBytes()))); Vertex v1 = Vertex.create("v1", pd1, 10, Resource.newInstance(1024, 1)); Vertex v2 = Vertex.create("v2", pd2, 1, Resource.newInstance(1024, 1)); v1.setTaskLaunchCmdOpts("").setTaskEnvironment(new HashMap<String, String>()) .addTaskLocalFiles(new HashMap<String, LocalResource>()); v2.setTaskLaunchCmdOpts("").setTaskEnvironment(new HashMap<String, String>()) .addTaskLocalFiles(new HashMap<String, LocalResource>()); InputDescriptor inputDescriptor = InputDescriptor.create("input") .setUserPayload(UserPayload.create(ByteBuffer.wrap("inputBytes".getBytes()))); OutputDescriptor outputDescriptor = OutputDescriptor.create("output") .setUserPayload(UserPayload.create(ByteBuffer.wrap("outputBytes".getBytes()))); Edge edge = Edge.create(v1, v2, EdgeProperty.create(DataMovementType.SCATTER_GATHER, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, outputDescriptor, inputDescriptor)); dag.addVertex(v1).addVertex(v2).addEdge(edge); Credentials dagCredentials = new Credentials(); Token<TokenIdentifier> token1 = new Token<TokenIdentifier>(); Token<TokenIdentifier> token2 = new Token<TokenIdentifier>(); dagCredentials.addToken(new Text("Token1"), token1); dagCredentials.addToken(new Text("Token2"), token2); dag.setCredentials(dagCredentials);/*from w w w .j a va 2s . c o m*/ DAGPlan dagProto = dag.createDag(new TezConfiguration(), null, null, null, true); assertTrue(dagProto.hasCredentialsBinary()); Credentials fetchedCredentials = DagTypeConverters .convertByteStringToCredentials(dagProto.getCredentialsBinary()); assertEquals(2, fetchedCredentials.numberOfTokens()); assertNotNull(fetchedCredentials.getToken(new Text("Token1"))); assertNotNull(fetchedCredentials.getToken(new Text("Token2"))); }
From source file:org.apache.tez.dag.app.dag.app.TestTezTaskCommunicatorManager.java
License:Apache License
@Test(timeout = 5000) public void testContainerAliveOnGetTask() throws IOException { TaskCommunicatorContext context = mock(TaskCommunicatorContext.class); Configuration conf = new Configuration(false); UserPayload userPayload = TezUtils.createUserPayloadFromConf(conf); ApplicationId appId = ApplicationId.newInstance(1000, 1); ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1); ContainerId containerId = createContainerId(appId, 1); doReturn(appAttemptId).when(context).getApplicationAttemptId(); doReturn(userPayload).when(context).getInitialUserPayload(); doReturn(new Credentials()).when(context).getAMCredentials(); TezTaskCommunicatorImpl taskComm = new TezTaskCommunicatorImpl(context); ContainerContext containerContext = new ContainerContext(containerId.toString()); taskComm.registerRunningContainer(containerId, "fakehost", 0); ContainerTask containerTask = taskComm.getUmbilical().getTask(containerContext); assertNull(containerTask);/*from w ww .j av a 2 s.co m*/ verify(context).containerAlive(containerId); }
From source file:org.apache.tez.dag.app.dag.impl.TestCommit.java
License:Apache License
@SuppressWarnings({ "unchecked", "rawtypes" }) public void setupDAG(DAGPlan dagPlan) { conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false); appAttemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(100, 1), 1); dagId = TezDAGID.getInstance(appAttemptId.getApplicationId(), 1); Assert.assertNotNull(dagId);/*from ww w . j av a2 s . c o m*/ dispatcher = new DrainDispatcher(); fsTokens = new Credentials(); appContext = mock(AppContext.class); when(appContext.getHadoopShim()).thenReturn(new DefaultHadoopShim()); rawExecutor = Executors.newCachedThreadPool( new ThreadFactoryBuilder().setDaemon(true).setNameFormat("App Shared Pool - " + "#%d").build()); execService = MoreExecutors.listeningDecorator(rawExecutor); doReturn(execService).when(appContext).getExecService(); historyEventHandler = new MockHistoryEventHandler(appContext); aclManager = new ACLManager("amUser"); doReturn(conf).when(appContext).getAMConf(); doReturn(appAttemptId).when(appContext).getApplicationAttemptId(); doReturn(appAttemptId.getApplicationId()).when(appContext).getApplicationID(); doReturn(dagId).when(appContext).getCurrentDAGID(); doReturn(historyEventHandler).when(appContext).getHistoryHandler(); doReturn(aclManager).when(appContext).getAMACLManager(); dag = new DAGImpl(dagId, conf, dagPlan, dispatcher.getEventHandler(), taskCommunicatorManagerInterface, fsTokens, clock, "user", thh, appContext); doReturn(dag).when(appContext).getCurrentDAG(); doReturn(dispatcher.getEventHandler()).when(appContext).getEventHandler(); ClusterInfo clusterInfo = new ClusterInfo(Resource.newInstance(8192, 10)); doReturn(clusterInfo).when(appContext).getClusterInfo(); dispatcher.register(CallableEventType.class, new CallableEventDispatcher()); taskEventDispatcher = new TaskEventDispatcher(); dispatcher.register(TaskEventType.class, taskEventDispatcher); taskAttemptEventDispatcher = new TaskAttemptEventDispatcher(); dispatcher.register(TaskAttemptEventType.class, taskAttemptEventDispatcher); vertexEventDispatcher = new VertexEventDispatcher(); dispatcher.register(VertexEventType.class, vertexEventDispatcher); dagEventDispatcher = new DagEventDispatcher(); dispatcher.register(DAGEventType.class, dagEventDispatcher); dagFinishEventHandler = new DAGFinishEventHandler(); dispatcher.register(DAGAppMasterEventType.class, dagFinishEventHandler); dispatcher.init(conf); dispatcher.start(); }
From source file:org.apache.tez.dag.app.dag.impl.TestDAGImpl.java
License:Apache License
@SuppressWarnings({ "unchecked", "rawtypes" }) @Before//w w w . ja va 2 s . c o m public void setup() { conf = new Configuration(); conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false); appAttemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(100, 1), 1); dagId = TezDAGID.getInstance(appAttemptId.getApplicationId(), 1); Assert.assertNotNull(dagId); dagPlan = createTestDAGPlan(); dispatcher = new DrainDispatcher(); fsTokens = new Credentials(); appContext = mock(AppContext.class); execService = mock(ListeningExecutorService.class); final ListenableFuture<Void> mockFuture = mock(ListenableFuture.class); Mockito.doAnswer(new Answer() { public ListenableFuture<Void> answer(InvocationOnMock invocation) { Object[] args = invocation.getArguments(); CallableEvent e = (CallableEvent) args[0]; dispatcher.getEventHandler().handle(e); return mockFuture; } }).when(execService).submit((Callable<Void>) any()); doReturn(execService).when(appContext).getExecService(); historyEventHandler = mock(HistoryEventHandler.class); aclManager = new ACLManager("amUser"); doReturn(conf).when(appContext).getAMConf(); doReturn(appAttemptId).when(appContext).getApplicationAttemptId(); doReturn(appAttemptId.getApplicationId()).when(appContext).getApplicationID(); doReturn(dagId).when(appContext).getCurrentDAGID(); doReturn(historyEventHandler).when(appContext).getHistoryHandler(); doReturn(aclManager).when(appContext).getAMACLManager(); dag = new DAGImpl(dagId, conf, dagPlan, dispatcher.getEventHandler(), taskAttemptListener, fsTokens, clock, "user", thh, appContext); doReturn(dag).when(appContext).getCurrentDAG(); mrrAppContext = mock(AppContext.class); doReturn(aclManager).when(mrrAppContext).getAMACLManager(); doReturn(execService).when(mrrAppContext).getExecService(); mrrDagId = TezDAGID.getInstance(appAttemptId.getApplicationId(), 2); mrrDagPlan = createTestMRRDAGPlan(); mrrDag = new DAGImpl(mrrDagId, conf, mrrDagPlan, dispatcher.getEventHandler(), taskAttemptListener, fsTokens, clock, "user", thh, mrrAppContext); doReturn(conf).when(mrrAppContext).getAMConf(); doReturn(mrrDag).when(mrrAppContext).getCurrentDAG(); doReturn(appAttemptId).when(mrrAppContext).getApplicationAttemptId(); doReturn(appAttemptId.getApplicationId()).when(mrrAppContext).getApplicationID(); doReturn(historyEventHandler).when(mrrAppContext).getHistoryHandler(); groupAppContext = mock(AppContext.class); doReturn(aclManager).when(groupAppContext).getAMACLManager(); doReturn(execService).when(groupAppContext).getExecService(); groupDagId = TezDAGID.getInstance(appAttemptId.getApplicationId(), 3); groupDagPlan = createGroupDAGPlan(); groupDag = new DAGImpl(groupDagId, conf, groupDagPlan, dispatcher.getEventHandler(), taskAttemptListener, fsTokens, clock, "user", thh, groupAppContext); doReturn(conf).when(groupAppContext).getAMConf(); doReturn(groupDag).when(groupAppContext).getCurrentDAG(); doReturn(appAttemptId).when(groupAppContext).getApplicationAttemptId(); doReturn(appAttemptId.getApplicationId()).when(groupAppContext).getApplicationID(); doReturn(historyEventHandler).when(groupAppContext).getHistoryHandler(); // reset totalCommitCounter to 0 TotalCountingOutputCommitter.totalCommitCounter = 0; dispatcher.register(CallableEventType.class, new CallableEventDispatcher()); taskEventDispatcher = new TaskEventDispatcher(); dispatcher.register(TaskEventType.class, taskEventDispatcher); taskAttemptEventDispatcher = new TaskAttemptEventDispatcher(); dispatcher.register(TaskAttemptEventType.class, taskAttemptEventDispatcher); vertexEventDispatcher = new VertexEventDispatcher(); dispatcher.register(VertexEventType.class, vertexEventDispatcher); dagEventDispatcher = new DagEventDispatcher(); dispatcher.register(DAGEventType.class, dagEventDispatcher); dagFinishEventHandler = new DAGFinishEventHandler(); dispatcher.register(DAGAppMasterEventType.class, dagFinishEventHandler); dispatcher.register(TaskEventType.class, new TaskEventHandler()); dispatcher.init(conf); dispatcher.start(); }
From source file:org.apache.tez.dag.app.dag.impl.TestDAGRecovery.java
License:Apache License
@Before public void setUp() { mockAppContext = mock(AppContext.class, RETURNS_DEEP_STUBS); when(mockAppContext.getCurrentDAG().getDagUGI()).thenReturn(null); mockEventHandler = mock(EventHandler.class); tezCounters.findCounter("grp_1", "counter_1").increment(1); DAGPlan dagPlan = TestDAGImpl.createTestDAGPlan(); dag = new DAGImpl(dagId, new Configuration(), dagPlan, mockEventHandler, mock(TaskAttemptListener.class), new Credentials(), new SystemClock(), user, mock(TaskHeartbeatHandler.class), mockAppContext); }
From source file:org.apache.tez.dag.app.dag.impl.TestTaskAttempt.java
License:Apache License
private static ContainerContext createFakeContainerContext() { return new ContainerContext(new HashMap<String, LocalResource>(), new Credentials(), new HashMap<String, String>(), ""); }
From source file:org.apache.tez.dag.app.dag.impl.TestTaskImpl.java
License:Apache License
@Before public void setup() { conf = new Configuration(); taskAttemptListener = mock(TaskAttemptListener.class); taskHeartbeatHandler = mock(TaskHeartbeatHandler.class); credentials = new Credentials(); clock = new SystemClock(); locationHint = TaskLocationHint.createTaskLocationHint(null, null); appId = ApplicationId.newInstance(System.currentTimeMillis(), 1); dagId = TezDAGID.getInstance(appId, 1); vertexId = TezVertexID.getInstance(dagId, 1); appContext = mock(AppContext.class, RETURNS_DEEP_STUBS); mockContainerId = mock(ContainerId.class); mockContainer = mock(Container.class); mockAMContainer = mock(AMContainer.class); mockNodeId = mock(NodeId.class); when(mockContainer.getId()).thenReturn(mockContainerId); when(mockContainer.getNodeId()).thenReturn(mockNodeId); when(mockAMContainer.getContainer()).thenReturn(mockContainer); when(appContext.getAllContainers().get(mockContainerId)).thenReturn(mockAMContainer); taskResource = Resource.newInstance(1024, 1); localResources = new HashMap<String, LocalResource>(); environment = new HashMap<String, String>(); javaOpts = ""; leafVertex = false;//ww w.j ava 2s . c om containerContext = new ContainerContext(localResources, credentials, environment, javaOpts); Vertex vertex = mock(Vertex.class); eventHandler = new TestEventHandler(); mockTask = new MockTaskImpl(vertexId, partition, eventHandler, conf, taskAttemptListener, clock, taskHeartbeatHandler, appContext, leafVertex, locationHint, taskResource, containerContext, vertex); }
From source file:org.apache.tez.dag.app.dag.impl.TestVertexImpl.java
License:Apache License
@SuppressWarnings({ "unchecked", "rawtypes" }) public void setupPostDagCreation() throws AMUserCodeException { String dagName = "dag0"; dispatcher = new DrainDispatcher(); appContext = mock(AppContext.class); thh = mock(TaskHeartbeatHandler.class); historyEventHandler = mock(HistoryEventHandler.class); TaskSchedulerEventHandler taskScheduler = mock(TaskSchedulerEventHandler.class); UserGroupInformation ugi;/* w ww.ja v a 2 s . com*/ try { ugi = UserGroupInformation.getCurrentUser(); } catch (IOException e) { throw new RuntimeException(e); } DAG dag = mock(DAG.class); doReturn(ugi).when(dag).getDagUGI(); doReturn(dagName).when(dag).getName(); doReturn(appAttemptId).when(appContext).getApplicationAttemptId(); doReturn(appAttemptId.getApplicationId()).when(appContext).getApplicationID(); doReturn(dag).when(appContext).getCurrentDAG(); execService = mock(ListeningExecutorService.class); final ListenableFuture<Void> mockFuture = mock(ListenableFuture.class); Mockito.doAnswer(new Answer() { public ListenableFuture<Void> answer(InvocationOnMock invocation) { Object[] args = invocation.getArguments(); CallableEvent e = (CallableEvent) args[0]; dispatcher.getEventHandler().handle(e); return mockFuture; } }).when(execService).submit((Callable<Void>) any()); doReturn(execService).when(appContext).getExecService(); doReturn(conf).when(appContext).getAMConf(); doReturn(new Credentials()).when(dag).getCredentials(); doReturn(DAGPlan.getDefaultInstance()).when(dag).getJobPlan(); doReturn(dagId).when(appContext).getCurrentDAGID(); doReturn(dagId).when(dag).getID(); doReturn(taskScheduler).when(appContext).getTaskScheduler(); doReturn(Resource.newInstance(102400, 60)).when(taskScheduler).getTotalResources(); doReturn(historyEventHandler).when(appContext).getHistoryHandler(); doReturn(dispatcher.getEventHandler()).when(appContext).getEventHandler(); vertexGroups = Maps.newHashMap(); for (PlanVertexGroupInfo groupInfo : dagPlan.getVertexGroupsList()) { vertexGroups.put(groupInfo.getGroupName(), new VertexGroupInfo(groupInfo)); } updateTracker = new StateChangeNotifier(appContext.getCurrentDAG()); setupVertices(); when(dag.getVertex(any(TezVertexID.class))).thenAnswer(new Answer<Vertex>() { @Override public Vertex answer(InvocationOnMock invocation) throws Throwable { Object[] args = invocation.getArguments(); if (args.length != 1) { return null; } TezVertexID vId = (TezVertexID) args[0]; return vertexIdMap.get(vId); } }); when(dag.getVertex(any(String.class))).thenAnswer(new Answer<Vertex>() { @Override public Vertex answer(InvocationOnMock invocation) throws Throwable { Object[] args = invocation.getArguments(); if (args.length != 1) { return null; } String vId = (String) args[0]; return vertices.get(vId); } }); // TODO - this test logic is tightly linked to impl DAGImpl code. edges = new HashMap<String, Edge>(); for (EdgePlan edgePlan : dagPlan.getEdgeList()) { EdgeProperty edgeProperty = DagTypeConverters.createEdgePropertyMapFromDAGPlan(edgePlan); edges.put(edgePlan.getId(), new Edge(edgeProperty, dispatcher.getEventHandler())); } parseVertexEdges(); for (Edge edge : edges.values()) { edge.initialize(); } dispatcher.register(CallableEventType.class, new CallableEventDispatcher()); taskAttemptEventDispatcher = new TaskAttemptEventDispatcher(); dispatcher.register(TaskAttemptEventType.class, taskAttemptEventDispatcher); taskEventDispatcher = new TaskEventDispatcher(); dispatcher.register(TaskEventType.class, taskEventDispatcher); vertexEventDispatcher = new VertexEventDispatcher(); dispatcher.register(VertexEventType.class, vertexEventDispatcher); dagEventDispatcher = new DagEventDispatcher(); dispatcher.register(DAGEventType.class, dagEventDispatcher); dispatcher.init(conf); dispatcher.start(); }