Examples with FlinkExecutableStageContextFactory org.apache.beam.runners.flink.translation.functions.FlinkExecutableStageContextFactory used on opensource projects

Example 1 with FlinkExecutableStageContextFactory

use of org.apache.beam.runners.flink.translation.functions.FlinkExecutableStageContextFactory in project beam by apache.

the class ExecutableStageDoFnOperatorTest method getOperator.

@SuppressWarnings("rawtypes")
private ExecutableStageDoFnOperator getOperator(TupleTag<Integer> mainOutput, List<TupleTag<?>> additionalOutputs, DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory, WindowingStrategy windowingStrategy, @Nullable Coder keyCoder, Coder windowedInputCoder) {
    FlinkExecutableStageContextFactory contextFactory = Mockito.mock(FlinkExecutableStageContextFactory.class);
    when(contextFactory.get(any())).thenReturn(stageContext);
    final ExecutableStagePayload stagePayload;
    if (keyCoder != null) {
        stagePayload = this.stagePayloadWithUserState;
    } else {
        stagePayload = this.stagePayload;
    }
    ExecutableStageDoFnOperator<Integer, Integer> operator = new ExecutableStageDoFnOperator<>("transform", windowedInputCoder, Collections.emptyMap(), mainOutput, additionalOutputs, outputManagerFactory, Collections.emptyMap(), /* sideInputTagMapping */
    Collections.emptyList(), /* sideInputs */
    Collections.emptyMap(), /* sideInputId mapping */
    FlinkPipelineOptions.defaults(), stagePayload, jobInfo, contextFactory, createOutputMap(mainOutput, additionalOutputs), windowingStrategy, keyCoder, keyCoder != null ? new KvToByteBufferKeySelector<>(keyCoder, null) : null);
    Whitebox.setInternalState(operator, "stateRequestHandler", stateRequestHandler);
    return operator;
}