Examples with KeyedTwoInputStreamOperatorTestHarness org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness used on opensource projects

Example 1 with KeyedTwoInputStreamOperatorTestHarness

use of org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness in project flink by apache.

the class KeyedCoProcessOperatorTest method testEventTimeTimerWithState.

/**
	 * Verifies that we don't have leakage between different keys.
	 */
@Test
public void testEventTimeTimerWithState() throws Exception {
    KeyedCoProcessOperator<String, Integer, String, String> operator = new KeyedCoProcessOperator<>(new EventTimeTriggeringStatefulProcessFunction());
    TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness = new KeyedTwoInputStreamOperatorTestHarness<>(operator, new IntToStringKeySelector<>(), new IdentityKeySelector<String>(), BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setup();
    testHarness.open();
    testHarness.processWatermark1(new Watermark(1));
    testHarness.processWatermark2(new Watermark(1));
    // should set timer for 6
    testHarness.processElement1(new StreamRecord<>(17, 0L));
    testHarness.processWatermark1(new Watermark(2));
    testHarness.processWatermark2(new Watermark(2));
    // should set timer for 7
    testHarness.processElement2(new StreamRecord<>("42", 1L));
    testHarness.processWatermark1(new Watermark(6));
    testHarness.processWatermark2(new Watermark(6));
    testHarness.processWatermark1(new Watermark(7));
    testHarness.processWatermark2(new Watermark(7));
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
    expectedOutput.add(new Watermark(1L));
    expectedOutput.add(new StreamRecord<>("INPUT1:17", 0L));
    expectedOutput.add(new Watermark(2L));
    expectedOutput.add(new StreamRecord<>("INPUT2:42", 1L));
    expectedOutput.add(new StreamRecord<>("STATE:17", 6L));
    expectedOutput.add(new Watermark(6L));
    expectedOutput.add(new StreamRecord<>("STATE:42", 7L));
    expectedOutput.add(new Watermark(7L));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
    testHarness.close();
}

Example 2 with KeyedTwoInputStreamOperatorTestHarness

use of org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness in project flink by apache.

the class KeyedCoProcessOperatorTest method testTimestampAndWatermarkQuerying.

@Test
public void testTimestampAndWatermarkQuerying() throws Exception {
    KeyedCoProcessOperator<String, Integer, String, String> operator = new KeyedCoProcessOperator<>(new WatermarkQueryingProcessFunction());
    TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness = new KeyedTwoInputStreamOperatorTestHarness<>(operator, new IntToStringKeySelector<>(), new IdentityKeySelector<String>(), BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setup();
    testHarness.open();
    testHarness.processWatermark1(new Watermark(17));
    testHarness.processWatermark2(new Watermark(17));
    testHarness.processElement1(new StreamRecord<>(5, 12L));
    testHarness.processWatermark1(new Watermark(42));
    testHarness.processWatermark2(new Watermark(42));
    testHarness.processElement2(new StreamRecord<>("6", 13L));
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
    expectedOutput.add(new Watermark(17L));
    expectedOutput.add(new StreamRecord<>("5WM:17 TS:12", 12L));
    expectedOutput.add(new Watermark(42L));
    expectedOutput.add(new StreamRecord<>("6WM:42 TS:13", 13L));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
    testHarness.close();
}

Example 3 with KeyedTwoInputStreamOperatorTestHarness

use of org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness in project flink by apache.

the class KeyedCoProcessOperatorTest method testEventTimeTimers.

@Test
public void testEventTimeTimers() throws Exception {
    KeyedCoProcessOperator<String, Integer, String, String> operator = new KeyedCoProcessOperator<>(new EventTimeTriggeringProcessFunction());
    TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness = new KeyedTwoInputStreamOperatorTestHarness<>(operator, new IntToStringKeySelector<>(), new IdentityKeySelector<String>(), BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setup();
    testHarness.open();
    testHarness.processElement1(new StreamRecord<>(17, 42L));
    testHarness.processElement2(new StreamRecord<>("18", 42L));
    testHarness.processWatermark1(new Watermark(5));
    testHarness.processWatermark2(new Watermark(5));
    testHarness.processWatermark1(new Watermark(6));
    testHarness.processWatermark2(new Watermark(6));
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
    expectedOutput.add(new StreamRecord<>("INPUT1:17", 42L));
    expectedOutput.add(new StreamRecord<>("INPUT2:18", 42L));
    expectedOutput.add(new StreamRecord<>("1777", 5L));
    expectedOutput.add(new Watermark(5L));
    expectedOutput.add(new StreamRecord<>("1777", 6L));
    expectedOutput.add(new Watermark(6L));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
    testHarness.close();
}

Example 4 with KeyedTwoInputStreamOperatorTestHarness

use of org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness in project flink by apache.

the class KeyedCoProcessOperatorTest method testSnapshotAndRestore.

@Test
public void testSnapshotAndRestore() throws Exception {
    KeyedCoProcessOperator<String, Integer, String, String> operator = new KeyedCoProcessOperator<>(new BothTriggeringProcessFunction());
    TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness = new KeyedTwoInputStreamOperatorTestHarness<>(operator, new IntToStringKeySelector<>(), new IdentityKeySelector<String>(), BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setup();
    testHarness.open();
    testHarness.processElement1(new StreamRecord<>(5, 12L));
    testHarness.processElement2(new StreamRecord<>("5", 12L));
    // snapshot and restore from scratch
    OperatorStateHandles snapshot = testHarness.snapshot(0, 0);
    testHarness.close();
    operator = new KeyedCoProcessOperator<>(new BothTriggeringProcessFunction());
    testHarness = new KeyedTwoInputStreamOperatorTestHarness<>(operator, new IntToStringKeySelector<>(), new IdentityKeySelector<String>(), BasicTypeInfo.STRING_TYPE_INFO);
    testHarness.setup();
    testHarness.initializeState(snapshot);
    testHarness.open();
    testHarness.setProcessingTime(5);
    testHarness.processWatermark1(new Watermark(6));
    testHarness.processWatermark2(new Watermark(6));
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
    expectedOutput.add(new StreamRecord<>("PROC:1777", 5L));
    expectedOutput.add(new StreamRecord<>("EVENT:1777", 6L));
    expectedOutput.add(new Watermark(6));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
    testHarness.close();
}

Example 5 with KeyedTwoInputStreamOperatorTestHarness

use of org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness in project beam by apache.

the class DoFnOperatorTest method testSideInputs.

public void testSideInputs(boolean keyed) throws Exception {
    WindowedValue.ValueOnlyWindowedValueCoder<String> windowedValueCoder = WindowedValue.getValueOnlyCoder(StringUtf8Coder.of());
    TupleTag<String> outputTag = new TupleTag<>("main-output");
    ImmutableMap<Integer, PCollectionView<?>> sideInputMapping = ImmutableMap.<Integer, PCollectionView<?>>builder().put(1, view1).put(2, view2).build();
    Coder<String> keyCoder = null;
    if (keyed) {
        keyCoder = StringUtf8Coder.of();
    }
    DoFnOperator<String, String, String> doFnOperator = new DoFnOperator<>(new IdentityDoFn<String>(), "stepName", windowedValueCoder, outputTag, Collections.<TupleTag<?>>emptyList(), new DoFnOperator.DefaultOutputManagerFactory<String>(), WindowingStrategy.globalDefault(), sideInputMapping, /* side-input mapping */
    ImmutableList.<PCollectionView<?>>of(view1, view2), /* side inputs */
    PipelineOptionsFactory.as(FlinkPipelineOptions.class), keyCoder);
    TwoInputStreamOperatorTestHarness<WindowedValue<String>, RawUnionValue, String> testHarness = new TwoInputStreamOperatorTestHarness<>(doFnOperator);
    if (keyed) {
        // we use a dummy key for the second input since it is considered to be broadcast
        testHarness = new KeyedTwoInputStreamOperatorTestHarness<>(doFnOperator, new StringKeySelector(), new DummyKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
    }
    testHarness.open();
    IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(100));
    IntervalWindow secondWindow = new IntervalWindow(new Instant(0), new Instant(500));
    // test the keep of sideInputs events
    testHarness.processElement2(new StreamRecord<>(new RawUnionValue(1, valuesInWindow(ImmutableList.of("hello", "ciao"), new Instant(0), firstWindow))));
    testHarness.processElement2(new StreamRecord<>(new RawUnionValue(2, valuesInWindow(ImmutableList.of("foo", "bar"), new Instant(0), secondWindow))));
    // push in a regular elements
    WindowedValue<String> helloElement = valueInWindow("Hello", new Instant(0), firstWindow);
    WindowedValue<String> worldElement = valueInWindow("World", new Instant(1000), firstWindow);
    testHarness.processElement1(new StreamRecord<>(helloElement));
    testHarness.processElement1(new StreamRecord<>(worldElement));
    // test the keep of pushed-back events
    testHarness.processElement2(new StreamRecord<>(new RawUnionValue(1, valuesInWindow(ImmutableList.of("hello", "ciao"), new Instant(1000), firstWindow))));
    testHarness.processElement2(new StreamRecord<>(new RawUnionValue(2, valuesInWindow(ImmutableList.of("foo", "bar"), new Instant(1000), secondWindow))));
    assertThat(this.<String>stripStreamRecordFromWindowedValue(testHarness.getOutput()), contains(helloElement, worldElement));
    testHarness.close();
}