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

Example 11 with Operation

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

the class IntrinsicMapTaskExecutorTest method testNoReadOperation.

@Test
public void testNoReadOperation() throws Exception {
    // Test MapTaskExecutor without ReadOperation.
    List<Operation> operations = Arrays.<Operation>asList(createOperation("o1", 1), createOperation("o2", 2));
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
        thrown.expect(IllegalStateException.class);
        thrown.expectMessage("is not a ReadOperation");
        executor.getReadOperation();
    }
}

Example 12 with Operation

use of org.apache.beam.runners.dataflow.worker.util.common.worker.Operation 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 13 with Operation

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

the class IntrinsicMapTaskExecutorTest method testExceptionInStartAbortsAllOperations.

@Test
public void testExceptionInStartAbortsAllOperations() throws Exception {
    Operation o1 = Mockito.mock(Operation.class);
    Operation o2 = Mockito.mock(Operation.class);
    Operation o3 = Mockito.mock(Operation.class);
    Mockito.doThrow(new Exception("in start")).when(o2).start();
    ExecutionStateTracker stateTracker = ExecutionStateTracker.newForTest();
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(Arrays.<Operation>asList(o1, o2, o3), counterSet, stateTracker)) {
        executor.execute();
        fail("Should have thrown");
    } catch (Exception e) {
        InOrder inOrder = Mockito.inOrder(o1, o2, o3);
        inOrder.verify(o3).start();
        inOrder.verify(o2).start();
        // Order of abort doesn't matter
        Mockito.verify(o1).abort();
        Mockito.verify(o2).abort();
        Mockito.verify(o3).abort();
        Mockito.verifyNoMoreInteractions(o1, o2, o3);
    }
}

Example 14 with Operation

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

the class IntrinsicMapTaskExecutorTest method testExecuteMapTaskExecutor.

@Test
public void testExecuteMapTaskExecutor() throws Exception {
    Operation o1 = Mockito.mock(Operation.class);
    Operation o2 = Mockito.mock(Operation.class);
    Operation o3 = Mockito.mock(Operation.class);
    List<Operation> operations = Arrays.asList(new Operation[] { o1, o2, o3 });
    ExecutionStateTracker stateTracker = Mockito.mock(ExecutionStateTracker.class);
    try (IntrinsicMapTaskExecutor executor = IntrinsicMapTaskExecutor.withSharedCounterSet(operations, counterSet, stateTracker)) {
        executor.execute();
    }
    InOrder inOrder = Mockito.inOrder(stateTracker, o1, o2, o3);
    inOrder.verify(o3).start();
    inOrder.verify(o2).start();
    inOrder.verify(o1).start();
    inOrder.verify(o1).finish();
    inOrder.verify(o2).finish();
    inOrder.verify(o3).finish();
}

Example 15 with Operation

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

the class NodesTest method testOperationNode.

@Test
public void testOperationNode() {
    Operation param = mock(Operation.class);
    assertSame(param, OperationNode.create(param).getOperation());
    assertNotEquals(OperationNode.create(param), OperationNode.create(param));
}