Examples with TaskCommunicatorContext org.apache.tez.serviceplugins.api.TaskCommunicatorContext used on opensource projects

Example 1 with TaskCommunicatorContext

use of org.apache.tez.serviceplugins.api.TaskCommunicatorContext in project tez by apache.

the class TestTaskCommunicatorContextImpl method testIsKnownContainer.

@Test(timeout = 5000)
public void testIsKnownContainer() {
    AppContext appContext = mock(AppContext.class);
    when(appContext.getAMConf()).thenReturn(new Configuration());
    TaskCommunicatorManager tal = mock(TaskCommunicatorManager.class);
    AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class), tal, mock(ContainerSignatureMatcher.class), appContext);
    doReturn(amContainerMap).when(appContext).getAllContainers();
    ContainerId containerId01 = mock(ContainerId.class);
    Container container01 = mock(Container.class);
    doReturn(containerId01).when(container01).getId();
    ContainerId containerId11 = mock(ContainerId.class);
    Container container11 = mock(Container.class);
    doReturn(containerId11).when(container11).getId();
    amContainerMap.addContainerIfNew(container01, 0, 0, 0);
    amContainerMap.addContainerIfNew(container11, 1, 1, 1);
    TaskCommunicatorContext taskCommContext0 = new TaskCommunicatorContextImpl(appContext, tal, null, 0);
    TaskCommunicatorContext taskCommContext1 = new TaskCommunicatorContextImpl(appContext, tal, null, 1);
    assertTrue(taskCommContext0.isKnownContainer(containerId01));
    assertFalse(taskCommContext0.isKnownContainer(containerId11));
    assertFalse(taskCommContext1.isKnownContainer(containerId01));
    assertTrue(taskCommContext1.isKnownContainer(containerId11));
    taskCommContext0.containerAlive(containerId01);
    verify(tal).containerAlive(containerId01);
    reset(tal);
    taskCommContext0.containerAlive(containerId11);
    verify(tal, never()).containerAlive(containerId11);
    reset(tal);
    taskCommContext1.containerAlive(containerId01);
    verify(tal, never()).containerAlive(containerId01);
    reset(tal);
    taskCommContext1.containerAlive(containerId11);
    verify(tal).containerAlive(containerId11);
    reset(tal);
    taskCommContext1.containerAlive(containerId01);
    verify(tal, never()).containerAlive(containerId01);
    reset(tal);
}

Example 2 with TaskCommunicatorContext

use of org.apache.tez.serviceplugins.api.TaskCommunicatorContext in project tez by apache.

the class TestTezTaskCommunicatorManager method testContainerAliveOnGetTask.

@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);
    verify(context).containerAlive(containerId);
}