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

Example 1 with TaskScheduler

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

the class TestTaskSchedulerManager method testTaskSchedulerUserError.

@SuppressWarnings("unchecked")
@Test(timeout = 5000)
public void testTaskSchedulerUserError() {
    TaskScheduler taskScheduler = mock(TaskScheduler.class, new ExceptionAnswer());
    EventHandler eventHandler = mock(EventHandler.class);
    AppContext appContext = mock(AppContext.class, RETURNS_DEEP_STUBS);
    when(appContext.getEventHandler()).thenReturn(eventHandler);
    doReturn("testTaskScheduler").when(appContext).getTaskSchedulerName(0);
    String expectedId = "[0:testTaskScheduler]";
    Configuration conf = new Configuration(false);
    InetSocketAddress address = new InetSocketAddress(15222);
    DAGClientServer mockClientService = mock(DAGClientServer.class);
    doReturn(address).when(mockClientService).getBindAddress();
    TaskSchedulerManager taskSchedulerManager = new TaskSchedulerManager(taskScheduler, appContext, mock(ContainerSignatureMatcher.class), mockClientService, Executors.newFixedThreadPool(1)) {

        @Override
        protected void instantiateSchedulers(String host, int port, String trackingUrl, AppContext appContext) throws TezException {
        // Stubbed out since these are setup up front in the constructor used for testing
        }
    };
    try {
        taskSchedulerManager.init(conf);
        taskSchedulerManager.start();
        // Invoking a couple of random methods
        AMSchedulerEventTALaunchRequest launchRequest = new AMSchedulerEventTALaunchRequest(mock(TezTaskAttemptID.class), mock(Resource.class), mock(TaskSpec.class), mock(TaskAttempt.class), mock(TaskLocationHint.class), 0, mock(ContainerContext.class), 0, 0, 0);
        taskSchedulerManager.handleEvent(launchRequest);
        ArgumentCaptor<Event> argumentCaptor = ArgumentCaptor.forClass(Event.class);
        verify(eventHandler, times(1)).handle(argumentCaptor.capture());
        Event rawEvent = argumentCaptor.getValue();
        assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError);
        DAGAppMasterEventUserServiceFatalError event = (DAGAppMasterEventUserServiceFatalError) rawEvent;
        assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType());
        assertTrue(event.getError().getMessage().contains("TestException_" + "allocateTask"));
        assertTrue(event.getDiagnosticInfo().contains("Task Allocation"));
        assertTrue(event.getDiagnosticInfo().contains(expectedId));
        taskSchedulerManager.dagCompleted();
        argumentCaptor = ArgumentCaptor.forClass(Event.class);
        verify(eventHandler, times(2)).handle(argumentCaptor.capture());
        rawEvent = argumentCaptor.getAllValues().get(1);
        assertTrue(rawEvent instanceof DAGAppMasterEventUserServiceFatalError);
        event = (DAGAppMasterEventUserServiceFatalError) rawEvent;
        assertEquals(DAGAppMasterEventType.TASK_SCHEDULER_SERVICE_FATAL_ERROR, event.getType());
        assertTrue(event.getError().getMessage().contains("TestException_" + "dagComplete"));
        assertTrue(event.getDiagnosticInfo().contains("Dag Completion"));
        assertTrue(event.getDiagnosticInfo().contains(expectedId));
    } finally {
        taskSchedulerManager.stop();
    }
}