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

Example 1 with TwoInputStreamOperatorTestHarness

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

the class CoProcessOperatorTest method testTimestampAndWatermarkQuerying.

@Test
public void testTimestampAndWatermarkQuerying() throws Exception {
    CoProcessOperator<Integer, String, String> operator = new CoProcessOperator<>(new WatermarkQueryingProcessFunction());
    TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness = new TwoInputStreamOperatorTestHarness<>(operator);
    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 2 with TwoInputStreamOperatorTestHarness

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

the class CoStreamMapTest method testCoMap.

@Test
public void testCoMap() throws Exception {
    CoStreamMap<Double, Integer, String> operator = new CoStreamMap<Double, Integer, String>(new MyCoMap());
    TwoInputStreamOperatorTestHarness<Double, Integer, String> testHarness = new TwoInputStreamOperatorTestHarness<Double, Integer, String>(operator);
    long initialTime = 0L;
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<Object>();
    testHarness.open();
    testHarness.processElement1(new StreamRecord<Double>(1.1d, initialTime + 1));
    testHarness.processElement1(new StreamRecord<Double>(1.2d, initialTime + 2));
    testHarness.processElement1(new StreamRecord<Double>(1.3d, initialTime + 3));
    testHarness.processWatermark1(new Watermark(initialTime + 3));
    testHarness.processElement1(new StreamRecord<Double>(1.4d, initialTime + 4));
    testHarness.processElement1(new StreamRecord<Double>(1.5d, initialTime + 5));
    testHarness.processElement2(new StreamRecord<Integer>(1, initialTime + 1));
    testHarness.processElement2(new StreamRecord<Integer>(2, initialTime + 2));
    testHarness.processWatermark2(new Watermark(initialTime + 2));
    testHarness.processElement2(new StreamRecord<Integer>(3, initialTime + 3));
    testHarness.processElement2(new StreamRecord<Integer>(4, initialTime + 4));
    testHarness.processElement2(new StreamRecord<Integer>(5, initialTime + 5));
    expectedOutput.add(new StreamRecord<String>("1.1", initialTime + 1));
    expectedOutput.add(new StreamRecord<String>("1.2", initialTime + 2));
    expectedOutput.add(new StreamRecord<String>("1.3", initialTime + 3));
    expectedOutput.add(new StreamRecord<String>("1.4", initialTime + 4));
    expectedOutput.add(new StreamRecord<String>("1.5", initialTime + 5));
    expectedOutput.add(new StreamRecord<String>("1", initialTime + 1));
    expectedOutput.add(new StreamRecord<String>("2", initialTime + 2));
    expectedOutput.add(new Watermark(initialTime + 2));
    expectedOutput.add(new StreamRecord<String>("3", initialTime + 3));
    expectedOutput.add(new StreamRecord<String>("4", initialTime + 4));
    expectedOutput.add(new StreamRecord<String>("5", initialTime + 5));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
}

Example 3 with TwoInputStreamOperatorTestHarness

use of org.apache.flink.streaming.util.TwoInputStreamOperatorTestHarness 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();
}

Example 4 with TwoInputStreamOperatorTestHarness

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

the class DoFnOperatorTest method testSideInputs.

void testSideInputs(boolean keyed) throws Exception {
    Coder<WindowedValue<String>> coder = 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 = StringUtf8Coder.of();
    KeySelector<WindowedValue<String>, ByteBuffer> keySelector = null;
    if (keyed) {
        keySelector = value -> FlinkKeyUtils.encodeKey(value.getValue(), keyCoder);
    }
    DoFnOperator<String, String> doFnOperator = new DoFnOperator<>(new IdentityDoFn<>(), "stepName", coder, Collections.emptyMap(), outputTag, Collections.emptyList(), new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, coder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())), WindowingStrategy.of(FixedWindows.of(Duration.millis(100))), sideInputMapping, /* side-input mapping */
    ImmutableList.of(view1, view2), /* side inputs */
    FlinkPipelineOptions.defaults(), keyed ? keyCoder : null, keyed ? keySelector : null, DoFnSchemaInformation.create(), Collections.emptyMap());
    TwoInputStreamOperatorTestHarness<WindowedValue<String>, RawUnionValue, WindowedValue<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, keySelector, null, new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of(), FlinkPipelineOptions.defaults()));
    }
    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(PCollectionViewTesting.materializeValuesFor(view1.getPipeline().getOptions(), View.asIterable(), "hello", "ciao"), new Instant(0), firstWindow))));
    testHarness.processElement2(new StreamRecord<>(new RawUnionValue(2, valuesInWindow(PCollectionViewTesting.materializeValuesFor(view2.getPipeline().getOptions(), View.asIterable(), "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(PCollectionViewTesting.materializeValuesFor(view1.getPipeline().getOptions(), View.asIterable(), "hello", "ciao"), new Instant(1000), firstWindow))));
    testHarness.processElement2(new StreamRecord<>(new RawUnionValue(2, valuesInWindow(PCollectionViewTesting.materializeValuesFor(view2.getPipeline().getOptions(), View.asIterable(), "foo", "bar"), new Instant(1000), secondWindow))));
    assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), contains(helloElement, worldElement));
    testHarness.close();
}

Example 5 with TwoInputStreamOperatorTestHarness

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

the class CoStreamFlatMapTest method testCoFlatMap.

@Test
public void testCoFlatMap() throws Exception {
    CoStreamFlatMap<String, Integer, String> operator = new CoStreamFlatMap<String, Integer, String>(new MyCoFlatMap());
    TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = new TwoInputStreamOperatorTestHarness<String, Integer, String>(operator);
    long initialTime = 0L;
    ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<Object>();
    testHarness.open();
    testHarness.processElement1(new StreamRecord<String>("abc", initialTime + 1));
    testHarness.processElement1(new StreamRecord<String>("def", initialTime + 2));
    testHarness.processWatermark1(new Watermark(initialTime + 2));
    testHarness.processElement1(new StreamRecord<String>("ghi", initialTime + 3));
    testHarness.processElement2(new StreamRecord<Integer>(1, initialTime + 1));
    testHarness.processElement2(new StreamRecord<Integer>(2, initialTime + 2));
    testHarness.processWatermark2(new Watermark(initialTime + 3));
    testHarness.processElement2(new StreamRecord<Integer>(3, initialTime + 3));
    testHarness.processElement2(new StreamRecord<Integer>(4, initialTime + 4));
    testHarness.processElement2(new StreamRecord<Integer>(5, initialTime + 5));
    expectedOutput.add(new StreamRecord<String>("a", initialTime + 1));
    expectedOutput.add(new StreamRecord<String>("b", initialTime + 1));
    expectedOutput.add(new StreamRecord<String>("c", initialTime + 1));
    expectedOutput.add(new StreamRecord<String>("d", initialTime + 2));
    expectedOutput.add(new StreamRecord<String>("e", initialTime + 2));
    expectedOutput.add(new StreamRecord<String>("f", initialTime + 2));
    expectedOutput.add(new StreamRecord<String>("g", initialTime + 3));
    expectedOutput.add(new StreamRecord<String>("h", initialTime + 3));
    expectedOutput.add(new StreamRecord<String>("i", initialTime + 3));
    expectedOutput.add(new StreamRecord<String>("1", initialTime + 1));
    expectedOutput.add(new StreamRecord<String>("2", initialTime + 2));
    expectedOutput.add(new Watermark(initialTime + 2));
    expectedOutput.add(new StreamRecord<String>("3", initialTime + 3));
    expectedOutput.add(new StreamRecord<String>("4", initialTime + 4));
    expectedOutput.add(new StreamRecord<String>("5", initialTime + 5));
    TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
}