Examples with TezCounters org.apache.tez.common.counters.TezCounters used on opensource projects

Example 86 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestMergeManager method createMockInputContext.

private InputContext createMockInputContext(String uniqueId, long mem) {
    InputContext inputContext = mock(InputContext.class);
    doReturn(new TezCounters()).when(inputContext).getCounters();
    doReturn(mem).when(inputContext).getTotalMemoryAvailableToTask();
    doReturn("srcVertexName").when(inputContext).getSourceVertexName();
    doReturn(uniqueId).when(inputContext).getUniqueIdentifier();
    return inputContext;
}

Example 87 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestShuffleInputEventHandlerOrderedGrouped method createTezInputContext.

private InputContext createTezInputContext() throws IOException {
    ApplicationId applicationId = ApplicationId.newInstance(1, 1);
    InputContext inputContext = mock(InputContext.class);
    doReturn(applicationId).when(inputContext).getApplicationId();
    doReturn("sourceVertex").when(inputContext).getSourceVertexName();
    when(inputContext.getCounters()).thenReturn(new TezCounters());
    ExecutionContext executionContext = new ExecutionContextImpl("localhost");
    doReturn(executionContext).when(inputContext).getExecutionContext();
    ByteBuffer shuffleBuffer = ByteBuffer.allocate(4).putInt(0, 4);
    doReturn(shuffleBuffer).when(inputContext).getServiceProviderMetaData(anyString());
    Token<JobTokenIdentifier> sessionToken = new Token<JobTokenIdentifier>(new JobTokenIdentifier(new Text("text")), new JobTokenSecretManager());
    ByteBuffer tokenBuffer = TezCommonUtils.serializeServiceData(sessionToken);
    doReturn(tokenBuffer).when(inputContext).getServiceConsumerMetaData(anyString());
    when(inputContext.createTezFrameworkExecutorService(anyInt(), anyString())).thenAnswer(new Answer<ExecutorService>() {

        @Override
        public ExecutorService answer(InvocationOnMock invocation) throws Throwable {
            return sharedExecutor.createExecutorService(invocation.getArgumentAt(0, Integer.class), invocation.getArgumentAt(1, String.class));
        }
    });
    return inputContext;
}

Example 88 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class TestOnFileSortedOutput method createTezOutputContext.

private OutputContext createTezOutputContext() throws IOException {
    String[] workingDirs = { workingDir.toString() };
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
    DataOutputBuffer serviceProviderMetaData = new DataOutputBuffer();
    serviceProviderMetaData.writeInt(PORT);
    TezCounters counters = new TezCounters();
    OutputStatisticsReporter reporter = mock(OutputStatisticsReporter.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            outputSize.set((Long) invocation.getArguments()[0]);
            return null;
        }
    }).when(reporter).reportDataSize(anyLong());
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            numRecords.set((Long) invocation.getArguments()[0]);
            return null;
        }
    }).when(reporter).reportItemsProcessed(anyLong());
    OutputContext context = mock(OutputContext.class);
    doReturn(counters).when(context).getCounters();
    doReturn(workingDirs).when(context).getWorkDirs();
    doReturn(payLoad).when(context).getUserPayload();
    doReturn(5 * 1024 * 1024l).when(context).getTotalMemoryAvailableToTask();
    doReturn(UniqueID).when(context).getUniqueIdentifier();
    doReturn("v1").when(context).getDestinationVertexName();
    doReturn(ByteBuffer.wrap(serviceProviderMetaData.getData())).when(context).getServiceProviderMetaData(conf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT));
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            long requestedSize = (Long) invocation.getArguments()[0];
            MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation.getArguments()[1];
            callback.memoryAssigned(requestedSize);
            return null;
        }
    }).when(context).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
    ExecutionContext ExecutionContext = mock(ExecutionContext.class);
    doReturn(HOST).when(ExecutionContext).getHostName();
    doReturn(reporter).when(context).getStatisticsReporter();
    doReturn(ExecutionContext).when(context).getExecutionContext();
    return context;
}

Example 89 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class OutputTestHelpers method createOutputContext.

static OutputContext createOutputContext(Configuration conf, Path workingDir) throws IOException {
    OutputContext ctx = mock(OutputContext.class);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            long requestedSize = (Long) invocation.getArguments()[0];
            MemoryUpdateCallbackHandler callback = (MemoryUpdateCallbackHandler) invocation.getArguments()[1];
            callback.memoryAssigned(requestedSize);
            return null;
        }
    }).when(ctx).requestInitialMemory(anyLong(), any(MemoryUpdateCallback.class));
    doReturn(TezUtils.createUserPayloadFromConf(conf)).when(ctx).getUserPayload();
    doReturn("destinationVertex").when(ctx).getDestinationVertexName();
    doReturn("UUID").when(ctx).getUniqueIdentifier();
    doReturn(new String[] { workingDir.toString() }).when(ctx).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(ctx).getTotalMemoryAvailableToTask();
    doReturn(new TezCounters()).when(ctx).getCounters();
    OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
    doReturn(statsReporter).when(ctx).getStatisticsReporter();
    doReturn(new ExecutionContextImpl("localhost")).when(ctx).getExecutionContext();
    return ctx;
}

Example 90 with TezCounters

use of org.apache.tez.common.counters.TezCounters in project tez by apache.

the class OutputTestHelpers method createOutputContext.

static OutputContext createOutputContext() throws IOException {
    OutputContext outputContext = mock(OutputContext.class);
    Configuration conf = new TezConfiguration();
    UserPayload payLoad = TezUtils.createUserPayloadFromConf(conf);
    String[] workingDirs = new String[] { "workDir1" };
    OutputStatisticsReporter statsReporter = mock(OutputStatisticsReporter.class);
    TezCounters counters = new TezCounters();
    doReturn("destinationVertex").when(outputContext).getDestinationVertexName();
    doReturn(payLoad).when(outputContext).getUserPayload();
    doReturn(workingDirs).when(outputContext).getWorkDirs();
    doReturn(200 * 1024 * 1024l).when(outputContext).getTotalMemoryAvailableToTask();
    doReturn(counters).when(outputContext).getCounters();
    doReturn(statsReporter).when(outputContext).getStatisticsReporter();
    return outputContext;
}