Examples with TestReader org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader used on opensource projects

Example 1 with TestReader

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader in project beam by apache.

the class MapTaskExecutorTest method testPerElementProcessingTimeCounters.

/**
 * Verify counts for the per-element-output-time counter are correct.
 */
@Test
public void testPerElementProcessingTimeCounters() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(DataflowPipelineDebugOptions.class).setExperiments(Lists.newArrayList(DataflowElementExecutionTracker.TIME_PER_ELEMENT_EXPERIMENT));
    ExecutionStateSampler stateSampler = ExecutionStateSampler.newForTest();
    DataflowExecutionStateTracker stateTracker = new DataflowExecutionStateTracker(stateSampler, new TestDataflowExecutionState(NameContext.forStage("test-stage"), "other", null, /* requestingStepName */
    null, /* sideInputIndex */
    null, /* metricsContainer */
    NoopProfileScope.NOOP), counterSet, options, "test-work-item-id");
    NameContext parDoName = nameForStep("s1");
    // Wire a read operation with 3 elements to a ParDoOperation and assert that we count
    // the correct number of elements.
    ReadOperation read = ReadOperation.forTest(new TestReader("a", "b", "c"), new OutputReceiver(), TestOperationContext.create(counterSet, nameForStep("s0"), null, stateTracker));
    ParDoOperation parDo = new ParDoOperation(new NoopParDoFn(), new OutputReceiver[0], TestOperationContext.create(counterSet, parDoName, null, stateTracker));
    parDo.attachInput(read, 0);
    List<Operation> operations = Lists.newArrayList(read, parDo);
    try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) {
        executor.execute();
    }
    stateSampler.doSampling(100L);
    CounterName counterName = CounterName.named("per-element-processing-time").withOriginalName(parDoName);
    Counter<Long, CounterDistribution> counter = (Counter<Long, CounterDistribution>) counterSet.getExistingCounter(counterName);
    assertThat(counter.getAggregate().getCount(), equalTo(3L));
}

Example 2 with TestReader

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader in project beam by apache.

the class ReadOperationTest method testAbortBetweenStartAndFinish.

@Test
public void testAbortBetweenStartAndFinish() throws Exception {
    MockOutputReceiver receiver = new MockOutputReceiver();
    TestReader reader = new TestReader();
    ReadOperation op = ReadOperation.forTest(reader, receiver, context);
    op.start();
    op.abort();
    assertTrue(reader.aborted);
}

Example 3 with TestReader

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader in project beam by apache.

the class IntrinsicMapTaskExecutorTest method testPerElementProcessingTimeCounters.

/**
 * Verify counts for the per-element-output-time counter are correct.
 */
@Test
public void testPerElementProcessingTimeCounters() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(DataflowPipelineDebugOptions.class).setExperiments(Lists.newArrayList(DataflowElementExecutionTracker.TIME_PER_ELEMENT_EXPERIMENT));
    DataflowExecutionStateTracker stateTracker = new DataflowExecutionStateTracker(ExecutionStateSampler.newForTest(), new TestDataflowExecutionState(NameContext.forStage("test-stage"), "other", null, /* requestingStepName */
    null, /* sideInputIndex */
    null, /* metricsContainer */
    NoopProfileScope.NOOP), counterSet, options, "test-work-item-id");
    NameContext parDoName = nameForStep("s1");
    // Wire a read operation with 3 elements to a ParDoOperation and assert that we count
    // the correct number of elements.
    ReadOperation read = ReadOperation.forTest(new TestReader("a", "b", "c"), new OutputReceiver(), TestOperationContext.create(counterSet, nameForStep("s0"), null, stateTracker));
    ParDoOperation parDo = new ParDoOperation(new NoopParDoFn(), new OutputReceiver[0], TestOperationContext.create(counterSet, parDoName, null, stateTracker));
    parDo.attachInput(read, 0);
    List<Operation> operations = Lists.newArrayList(read, parDo);
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
        executor.execute();
    }
    CounterName counterName = CounterName.named("per-element-processing-time").withOriginalName(parDoName);
    Counter<Long, CounterDistribution> counter = (Counter<Long, CounterDistribution>) counterSet.getExistingCounter(counterName);
    assertThat(counter.getAggregate().getCount(), equalTo(3L));
}

Example 4 with TestReader

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader in project beam by apache.

the class ReadOperationTest method testRunReadOperation.

/**
 * Tests that a {@link ReadOperation} has expected counters, and that their values are reasonable.
 */
@Test
@SuppressWarnings("unchecked")
public void testRunReadOperation() throws Exception {
    TestReader reader = new TestReader("hi", "there", "", "bob");
    TestOutputReceiver receiver = new TestOutputReceiver(counterSet, NameContext.create("test", "test_receiver", "test_receiver", "test_receiver"));
    ReadOperation readOperation = ReadOperation.forTest(reader, receiver, context);
    readOperation.start();
    readOperation.finish();
    assertThat(receiver.outputElems, containsInAnyOrder((Object) "hi", "there", "", "bob"));
    CounterUpdateExtractor<?> updateExtractor = mock(CounterUpdateExtractor.class);
    counterSet.extractUpdates(false, updateExtractor);
    verify(updateExtractor).longSum(eq(getObjectCounterName("test_receiver_out")), anyBoolean(), eq(4L));
    verify(updateExtractor).longMean(eq(getMeanByteCounterName("test_receiver_out")), anyBoolean(), eq(LongCounterMean.ZERO.addValue(14L, 4)));
    verify(updateExtractor).longSum(eq(named("ReadOperation-ByteCount")), anyBoolean(), eq(10L));
    verifyNoMoreInteractions(updateExtractor);
}

Example 5 with TestReader

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils.TestReader in project beam by apache.

the class ReadOperationTest method testAbortAfterFinish.

@Test
public void testAbortAfterFinish() throws Exception {
    MockOutputReceiver receiver = new MockOutputReceiver();
    TestReader reader = new TestReader();
    ReadOperation op = ReadOperation.forTest(reader, receiver, context);
    op.start();
    op.finish();
    op.abort();
    assertTrue(reader.closed);
    assertTrue(reader.aborted);
}