Examples with MockEnvironment org.apache.flink.runtime.operators.testutils.MockEnvironment used on opensource projects

Example 1 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class SourceFunctionUtil method runSourceFunction.

public static <T extends Serializable> List<T> runSourceFunction(SourceFunction<T> sourceFunction) throws Exception {
    final List<T> outputs = new ArrayList<T>();
    if (sourceFunction instanceof RichFunction) {
        AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class);
        when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig());
        RuntimeContext runtimeContext = new StreamingRuntimeContext(operator, new MockEnvironment("MockTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024), new HashMap<String, Accumulator<?, ?>>());
        ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext);
        ((RichFunction) sourceFunction).open(new Configuration());
    }
    try {
        SourceFunction.SourceContext<T> ctx = new CollectingSourceContext<T>(new Object(), outputs);
        sourceFunction.run(ctx);
    } catch (Exception e) {
        throw new RuntimeException("Cannot invoke source.", e);
    }
    return outputs;
}

Example 2 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class LargeSortingDataInputITCase method multiInputKeySorting.

@Test
@SuppressWarnings({ "rawtypes", "unchecked" })
public void multiInputKeySorting() throws Exception {
    int numberOfRecords = 500_000;
    GeneratedRecordsDataInput input1 = new GeneratedRecordsDataInput(numberOfRecords, 0);
    GeneratedRecordsDataInput input2 = new GeneratedRecordsDataInput(numberOfRecords, 1);
    KeySelector<Tuple3<Integer, String, byte[]>, String> keySelector = value -> value.f1;
    try (MockEnvironment environment = MockEnvironment.builder().build()) {
        SelectableSortingInputs selectableSortingInputs = MultiInputSortingDataInput.wrapInputs(new DummyInvokable(), new StreamTaskInput[] { input1, input2 }, new KeySelector[] { keySelector, keySelector }, new TypeSerializer[] { GeneratedRecordsDataInput.SERIALIZER, GeneratedRecordsDataInput.SERIALIZER }, new StringSerializer(), new StreamTaskInput[0], environment.getMemoryManager(), environment.getIOManager(), true, 1.0, new Configuration(), new ExecutionConfig());
        StreamTaskInput<?>[] sortingDataInputs = selectableSortingInputs.getSortedInputs();
        try (StreamTaskInput<Tuple3<Integer, String, byte[]>> sortedInput1 = (StreamTaskInput<Tuple3<Integer, String, byte[]>>) sortingDataInputs[0];
            StreamTaskInput<Tuple3<Integer, String, byte[]>> sortedInput2 = (StreamTaskInput<Tuple3<Integer, String, byte[]>>) sortingDataInputs[1]) {
            VerifyingOutput<String> output = new VerifyingOutput<>(keySelector);
            StreamMultipleInputProcessor multiSortedProcessor = new StreamMultipleInputProcessor(new MultipleInputSelectionHandler(selectableSortingInputs.getInputSelectable(), 2), new StreamOneInputProcessor[] { new StreamOneInputProcessor(sortedInput1, output, new DummyOperatorChain()), new StreamOneInputProcessor(sortedInput2, output, new DummyOperatorChain()) });
            DataInputStatus inputStatus;
            do {
                inputStatus = multiSortedProcessor.processInput();
            } while (inputStatus != DataInputStatus.END_OF_INPUT);
            assertThat(output.getSeenRecords(), equalTo(numberOfRecords * 2));
        }
    }
}

Example 3 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class BatchExecutionInternalTimeServiceTest method testBatchExecutionManagerCanBeInstantiatedWithBatchStateBackend.

@Test
public void testBatchExecutionManagerCanBeInstantiatedWithBatchStateBackend() throws Exception {
    expectedException.expect(IllegalStateException.class);
    expectedException.expectMessage("Batch execution specific time service can work only with BatchExecutionKeyedStateBackend");
    MockEnvironment mockEnvironment = MockEnvironment.builder().build();
    AbstractKeyedStateBackend<Integer> stateBackend = new MemoryStateBackend().createKeyedStateBackend(mockEnvironment, new JobID(), "dummy", KEY_SERIALIZER, 2, new KeyGroupRange(0, 1), mockEnvironment.getTaskKvStateRegistry(), TtlTimeProvider.DEFAULT, new UnregisteredMetricsGroup(), Collections.emptyList(), new CloseableRegistry());
    BatchExecutionInternalTimeServiceManager.create(stateBackend, this.getClass().getClassLoader(), new DummyKeyContext(), new TestProcessingTimeService(), Collections.emptyList());
}

Example 4 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class MultiInputSortingDataInputsTest method twoInputOrderTest.

@SuppressWarnings("unchecked")
public void twoInputOrderTest(int preferredIndex, int sortedIndex) throws Exception {
    CollectingDataOutput<Object> collectingDataOutput = new CollectingDataOutput<>();
    List<StreamElement> sortedInputElements = Arrays.asList(new StreamRecord<>(1, 3), new StreamRecord<>(1, 1), new StreamRecord<>(2, 1), new StreamRecord<>(2, 3), new StreamRecord<>(1, 2), new StreamRecord<>(2, 2), Watermark.MAX_WATERMARK);
    CollectionDataInput<Integer> sortedInput = new CollectionDataInput<>(sortedInputElements, sortedIndex);
    List<StreamElement> preferredInputElements = Arrays.asList(new StreamRecord<>(99, 3), new StreamRecord<>(99, 1), new Watermark(99L));
    CollectionDataInput<Integer> preferredInput = new CollectionDataInput<>(preferredInputElements, preferredIndex);
    KeySelector<Integer, Integer> keySelector = value -> value;
    try (MockEnvironment environment = MockEnvironment.builder().build()) {
        SelectableSortingInputs selectableSortingInputs = MultiInputSortingDataInput.wrapInputs(new DummyInvokable(), new StreamTaskInput[] { sortedInput }, new KeySelector[] { keySelector }, new TypeSerializer[] { new IntSerializer() }, new IntSerializer(), new StreamTaskInput[] { preferredInput }, environment.getMemoryManager(), environment.getIOManager(), true, 1.0, new Configuration(), new ExecutionConfig());
        StreamTaskInput<?>[] sortingDataInputs = selectableSortingInputs.getSortedInputs();
        StreamTaskInput<?>[] preferredDataInputs = selectableSortingInputs.getPassThroughInputs();
        try (StreamTaskInput<Object> preferredTaskInput = (StreamTaskInput<Object>) preferredDataInputs[0];
            StreamTaskInput<Object> sortedTaskInput = (StreamTaskInput<Object>) sortingDataInputs[0]) {
            MultipleInputSelectionHandler selectionHandler = new MultipleInputSelectionHandler(selectableSortingInputs.getInputSelectable(), 2);
            @SuppressWarnings("rawtypes") StreamOneInputProcessor[] inputProcessors = new StreamOneInputProcessor[2];
            inputProcessors[preferredIndex] = new StreamOneInputProcessor<>(preferredTaskInput, collectingDataOutput, new DummyOperatorChain());
            inputProcessors[sortedIndex] = new StreamOneInputProcessor<>(sortedTaskInput, collectingDataOutput, new DummyOperatorChain());
            StreamMultipleInputProcessor processor = new StreamMultipleInputProcessor(selectionHandler, inputProcessors);
            DataInputStatus inputStatus;
            do {
                inputStatus = processor.processInput();
            } while (inputStatus != DataInputStatus.END_OF_INPUT);
        }
    }
    assertThat(collectingDataOutput.events, equalTo(Arrays.asList(new StreamRecord<>(99, 3), new StreamRecord<>(99, 1), // max watermark from the preferred input
    new Watermark(99L), new StreamRecord<>(1, 1), new StreamRecord<>(1, 2), new StreamRecord<>(1, 3), new StreamRecord<>(2, 1), new StreamRecord<>(2, 2), new StreamRecord<>(2, 3), // max watermark from the sorted input
    Watermark.MAX_WATERMARK)));
}

Example 5 with MockEnvironment

use of org.apache.flink.runtime.operators.testutils.MockEnvironment in project flink by apache.

the class SortingDataInputTest method simpleVariableLengthKeySorting.

@Test
public void simpleVariableLengthKeySorting() throws Exception {
    CollectingDataOutput<Integer> collectingDataOutput = new CollectingDataOutput<>();
    CollectionDataInput<Integer> input = new CollectionDataInput<>(Arrays.asList(new StreamRecord<>(1, 3), new StreamRecord<>(1, 1), new StreamRecord<>(2, 1), new StreamRecord<>(2, 3), new StreamRecord<>(1, 2), new StreamRecord<>(2, 2)));
    MockEnvironment environment = MockEnvironment.builder().build();
    SortingDataInput<Integer, String> sortingDataInput = new SortingDataInput<>(input, new IntSerializer(), new StringSerializer(), (KeySelector<Integer, String>) value -> "" + value, environment.getMemoryManager(), environment.getIOManager(), true, 1.0, new Configuration(), new DummyInvokable(), new ExecutionConfig());
    DataInputStatus inputStatus;
    do {
        inputStatus = sortingDataInput.emitNext(collectingDataOutput);
    } while (inputStatus != DataInputStatus.END_OF_INPUT);
    assertThat(collectingDataOutput.events, equalTo(Arrays.asList(new StreamRecord<>(1, 1), new StreamRecord<>(1, 2), new StreamRecord<>(1, 3), new StreamRecord<>(2, 1), new StreamRecord<>(2, 2), new StreamRecord<>(2, 3))));
}