Examples with DataflowStepContext org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext used on opensource projects

Example 1 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserParDoFnFactoryTest method testCleanupRegistered.

@Test
public void testCleanupRegistered() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    CounterSet counters = new CounterSet();
    DoFn<?, ?> initialFn = new TestStatefulDoFn();
    CloudObject cloudObject = getCloudObject(initialFn, WindowingStrategy.globalDefault().withWindowFn(FixedWindows.of(Duration.millis(10))));
    TimerInternals timerInternals = mock(TimerInternals.class);
    DataflowStepContext stepContext = mock(DataflowStepContext.class);
    when(stepContext.timerInternals()).thenReturn(timerInternals);
    DataflowExecutionContext<DataflowStepContext> executionContext = mock(DataflowExecutionContext.class);
    TestOperationContext operationContext = TestOperationContext.create(counters);
    when(executionContext.getStepContext(operationContext)).thenReturn(stepContext);
    when(executionContext.getSideInputReader(any(), any(), any())).thenReturn(NullSideInputReader.empty());
    ParDoFn parDoFn = factory.create(options, cloudObject, Collections.emptyList(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), executionContext, operationContext);
    Receiver rcvr = new OutputReceiver();
    parDoFn.startBundle(rcvr);
    IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(10));
    parDoFn.processElement(WindowedValue.of("foo", new Instant(1), firstWindow, PaneInfo.NO_FIRING));
    verify(stepContext).setStateCleanupTimer(SimpleParDoFn.CLEANUP_TIMER_ID, firstWindow, IntervalWindow.getCoder(), firstWindow.maxTimestamp().plus(Duration.millis(1L)), firstWindow.maxTimestamp().plus(Duration.millis(1L)));
}

Example 2 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class StreamingPCollectionViewWriterDoFnFactory method create.

@Override
public ParDoFn create(PipelineOptions options, CloudObject cloudUserFn, List<SideInputInfo> sideInputInfos, TupleTag<?> mainOutputTag, Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices, DataflowExecutionContext<?> executionContext, DataflowOperationContext operationContext) throws Exception {
    DataflowStepContext stepContext = executionContext.getStepContext(operationContext);
    checkArgument(stepContext instanceof StreamingModeExecutionContext.StreamingModeStepContext, "stepContext must be a StreamingModeStepContext to use StreamingPCollectionViewWriterFn");
    Coder<?> coder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(Structs.getObject(cloudUserFn, PropertyNames.ENCODING)));
    checkState(coder instanceof FullWindowedValueCoder, "Expected to received an instanceof an %s but got %s", FullWindowedValueCoder.class.getSimpleName(), coder);
    FullWindowedValueCoder<?> windowedValueCoder = (FullWindowedValueCoder<?>) coder;
    return new StreamingPCollectionViewWriterParDoFn((StreamingModeExecutionContext.StreamingModeStepContext) stepContext, new TupleTag<>(Structs.getString(cloudUserFn, WorkerPropertyNames.SIDE_INPUT_ID)), (Coder) windowedValueCoder.getValueCoder(), (Coder) windowedValueCoder.getWindowCoder());
}

Example 3 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class BeamFnMapTaskExecutorTest method generateDataflowStepContext.

/**
 * Generates bare minumum DataflowStepContext to use for testing.
 *
 * @param valuesPrefix prefix for all types of names that are specified in DataflowStepContext.
 * @return new instance of DataflowStepContext
 */
private DataflowStepContext generateDataflowStepContext(String valuesPrefix) {
    NameContext nc = new NameContext() {

        @Override
        @Nullable
        public String stageName() {
            return valuesPrefix + "Stage";
        }

        @Override
        @Nullable
        public String originalName() {
            return valuesPrefix + "OriginalName";
        }

        @Override
        @Nullable
        public String systemName() {
            return valuesPrefix + "SystemName";
        }

        @Override
        @Nullable
        public String userName() {
            return valuesPrefix + "UserName";
        }
    };
    DataflowStepContext dsc = new DataflowStepContext(nc) {

        @Override
        @Nullable
        public <W extends BoundedWindow> TimerData getNextFiredTimer(Coder<W> windowCoder) {
            return null;
        }

        @Override
        public <W extends BoundedWindow> void setStateCleanupTimer(String timerId, W window, Coder<W> windowCoder, Instant cleanupTime, Instant cleanupOutputTimestamp) {
        }

        @Override
        public DataflowStepContext namespacedToUser() {
            return this;
        }

        @Override
        public StateInternals stateInternals() {
            return null;
        }

        @Override
        public TimerInternals timerInternals() {
            return null;
        }
    };
    return dsc;
}

Example 4 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserMonitoringInfoToCounterUpdateTransformerTest method testTransformThrowsIfMonitoringInfoWithWrongUrnPrefixReceived.

@Test
public void testTransformThrowsIfMonitoringInfoWithWrongUrnPrefixReceived() {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.ELEMENT_COUNT).setType(TypeUrns.SUM_INT64_TYPE).build();
    UserMonitoringInfoToCounterUpdateTransformer testObject = new UserMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    exception.expect(RuntimeException.class);
    testObject.transform(monitoringInfo);
}

Example 5 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserMonitoringInfoToCounterUpdateTransformerTest method testTransformReturnsValidCounterUpdateWhenValidUserMonitoringInfoReceived.

@Test
public void testTransformReturnsValidCounterUpdateWhenValidUserMonitoringInfoReceived() throws Exception {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    NameContext nc = NameContext.create("anyStageName", "anyOriginalName", "anySystemName", "anyUserName");
    DataflowStepContext dsc = mock(DataflowStepContext.class);
    when(dsc.getNameContext()).thenReturn(nc);
    stepContextMapping.put("anyValue", dsc);
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.USER_SUM_INT64).setType(TypeUrns.SUM_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.NAME, "anyName").putLabels(MonitoringInfoConstants.Labels.NAMESPACE, "anyNamespace").putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "anyValue").setPayload(encodeInt64Counter(1L)).build();
    UserMonitoringInfoToCounterUpdateTransformer testObject = new UserMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    CounterUpdate result = testObject.transform(monitoringInfo);
    assertNotNull(result);
    assertEqualsAsJson("{cumulative:true, integer:{highBits:0, lowBits:1}, " + "structuredNameAndMetadata:{metadata:{kind:'SUM'}, " + "name:{name:'anyName', origin:'USER', originNamespace:'anyNamespace', " + "originalStepName:'anyOriginalName'}}}", result);
}