org.apache.flink.streaming.util.TwoInputStreamOperatorTestHarness Java Examples

private static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final int maxParallelism,
		final int numTasks,
		final int taskIdx,
		final OperatorSubtaskState initState,
		final MapStateDescriptor<?, ?>... descriptors) throws Exception {

	TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> testHarness = new TwoInputStreamOperatorTestHarness<>(
			new CoBroadcastWithNonKeyedOperator<>(
					Preconditions.checkNotNull(function),
					Arrays.asList(descriptors)),
			maxParallelism, numTasks, taskIdx
	);
	testHarness.setup();
	testHarness.initializeState(initState);
	testHarness.open();

	return testHarness;
}
 

private static <KEY, IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final TypeInformation<KEY> keyTypeInfo,
		final KeySelector<IN1, KEY> keyKeySelector,
		final KeyedBroadcastProcessFunction<KEY, IN1, IN2, OUT> function,
		final int maxParallelism,
		final int numTasks,
		final int taskIdx) throws Exception {

	return getInitializedTestHarness(
			keyTypeInfo,
			keyKeySelector,
			function,
			maxParallelism,
			numTasks,
			taskIdx,
			null);
}
 

private TwoInputStreamOperatorTestHarness<Trade, Customer, EnrichedTrade> setupHarness() throws Exception {
	// instantiate operator
	KeyedCoProcessOperator<Long, Trade, Customer, EnrichedTrade> operator =
			new KeyedCoProcessOperator<>(new EventTimeJoinExercise.EventTimeJoinFunction());

	// setup test harness
	TwoInputStreamOperatorTestHarness<Trade, Customer, EnrichedTrade> testHarness =
			new KeyedTwoInputStreamOperatorTestHarness<>(operator,
					(Trade t) -> t.customerId,
					(Customer c) -> c.customerId,
					BasicTypeInfo.LONG_TYPE_INFO);

	testHarness.setup();
	testHarness.open();

	return testHarness;
}
 

@Test
public void testMultiStateSupport() throws Exception {
	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness =
					getInitializedTestHarness(new FunctionWithMultipleStates(), STATE_DESCRIPTOR, STATE_DESCRIPTOR_A)
	) {
		testHarness.processElement2(new StreamRecord<>(5, 12L));
		testHarness.processElement2(new StreamRecord<>(6, 13L));

		testHarness.processElement1(new StreamRecord<>("9", 15L));

		Queue<Object> expectedBr = new ConcurrentLinkedQueue<>();
		expectedBr.add(new StreamRecord<>("9:key.6->6", 15L));
		expectedBr.add(new StreamRecord<>("9:key.5->5", 15L));
		expectedBr.add(new StreamRecord<>("9:5->value.5", 15L));
		expectedBr.add(new StreamRecord<>("9:6->value.6", 15L));

		TestHarnessUtil.assertOutputEquals("Wrong Side Output", expectedBr, testHarness.getOutput());
	}
}
 

@Test
public void testMultiStateSupport() throws Exception {
	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness =
					getInitializedTestHarness(new FunctionWithMultipleStates(), STATE_DESCRIPTOR, STATE_DESCRIPTOR_A)
	) {
		testHarness.processElement2(new StreamRecord<>(5, 12L));
		testHarness.processElement2(new StreamRecord<>(6, 13L));

		testHarness.processElement1(new StreamRecord<>("9", 15L));

		Queue<Object> expectedBr = new ConcurrentLinkedQueue<>();
		expectedBr.add(new StreamRecord<>("9:key.6->6", 15L));
		expectedBr.add(new StreamRecord<>("9:key.5->5", 15L));
		expectedBr.add(new StreamRecord<>("9:5->value.5", 15L));
		expectedBr.add(new StreamRecord<>("9:6->value.6", 15L));

		TestHarnessUtil.assertOutputEquals("Wrong Side Output", expectedBr, testHarness.getOutput());
	}
}
 

private static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final int maxParallelism,
		final int numTasks,
		final int taskIdx,
		final OperatorSubtaskState initState,
		final MapStateDescriptor<?, ?>... descriptors) throws Exception {

	TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> testHarness = new TwoInputStreamOperatorTestHarness<>(
			new CoBroadcastWithNonKeyedOperator<>(
					Preconditions.checkNotNull(function),
					Arrays.asList(descriptors)),
			maxParallelism, numTasks, taskIdx
	);
	testHarness.setup();
	testHarness.initializeState(initState);
	testHarness.open();

	return testHarness;
}
 

private static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final int maxParallelism,
		final int numTasks,
		final int taskIdx,
		final OperatorSubtaskState initState,
		final MapStateDescriptor<?, ?>... descriptors) throws Exception {

	TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> testHarness = new TwoInputStreamOperatorTestHarness<>(
			new CoBroadcastWithNonKeyedOperator<>(
					Preconditions.checkNotNull(function),
					Arrays.asList(descriptors)),
			maxParallelism, numTasks, taskIdx
	);
	testHarness.setup();
	testHarness.initializeState(initState);
	testHarness.open();

	return testHarness;
}
 

@Test
public void testBroadcastState() throws Exception {

	final Set<String> keysToRegister = new HashSet<>();
	keysToRegister.add("test1");
	keysToRegister.add("test2");
	keysToRegister.add("test3");

	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					new TestFunction(keysToRegister), STATE_DESCRIPTOR)
	) {
		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>("6", 13L));
		testHarness.processElement1(new StreamRecord<>("6", 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		Queue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

		expectedOutput.add(new Watermark(10L));
		expectedOutput.add(new StreamRecord<>("5WM:10 TS:12", 12L));
		expectedOutput.add(new Watermark(40L));
		expectedOutput.add(new StreamRecord<>("6WM:40 TS:13", 13L));
		expectedOutput.add(new StreamRecord<>("6WM:40 TS:15", 15L));
		expectedOutput.add(new Watermark(50L));

		TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
	}
}
 

@Test
public void testBroadcastState() throws Exception {

	final Set<String> keysToRegister = new HashSet<>();
	keysToRegister.add("test1");
	keysToRegister.add("test2");
	keysToRegister.add("test3");

	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					new TestFunction(keysToRegister), STATE_DESCRIPTOR)
	) {
		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>("6", 13L));
		testHarness.processElement1(new StreamRecord<>("6", 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		Queue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

		expectedOutput.add(new Watermark(10L));
		expectedOutput.add(new StreamRecord<>("5WM:10 TS:12", 12L));
		expectedOutput.add(new Watermark(40L));
		expectedOutput.add(new StreamRecord<>("6WM:40 TS:13", 13L));
		expectedOutput.add(new StreamRecord<>("6WM:40 TS:15", 15L));
		expectedOutput.add(new Watermark(50L));

		TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
	}
}
 

@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();
}
 

@Test
public void testSideOutput() throws Exception {
	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					new FunctionWithSideOutput(), STATE_DESCRIPTOR)
	) {

		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>("6", 13L));
		testHarness.processElement1(new StreamRecord<>("6", 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		ConcurrentLinkedQueue<StreamRecord<String>> expectedBr = new ConcurrentLinkedQueue<>();
		expectedBr.add(new StreamRecord<>("BR:5 WM:10 TS:12", 12L));

		ConcurrentLinkedQueue<StreamRecord<String>> expectedNonBr = new ConcurrentLinkedQueue<>();
		expectedNonBr.add(new StreamRecord<>("NON-BR:6 WM:40 TS:13", 13L));
		expectedNonBr.add(new StreamRecord<>("NON-BR:6 WM:40 TS:15", 15L));

		ConcurrentLinkedQueue<StreamRecord<String>> brSideOutput = testHarness.getSideOutput(FunctionWithSideOutput.BROADCAST_TAG);
		ConcurrentLinkedQueue<StreamRecord<String>> nonBrSideOutput = testHarness.getSideOutput(FunctionWithSideOutput.NON_BROADCAST_TAG);

		TestHarnessUtil.assertOutputEquals("Wrong Side Output", expectedBr, brSideOutput);
		TestHarnessUtil.assertOutputEquals("Wrong Side Output", expectedNonBr, nonBrSideOutput);
	}
}
 

@Test
public void testProcessingTimeTimers() throws Exception {

	KeyedCoProcessOperator<String, Integer, String, String> operator =
			new KeyedCoProcessOperator<>(new ProcessingTimeTriggeringProcessFunction());

	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));
	testHarness.processElement2(new StreamRecord<>("18"));

	testHarness.setProcessingTime(5);
	testHarness.setProcessingTime(6);

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("INPUT1:17"));
	expectedOutput.add(new StreamRecord<>("INPUT2:18"));
	expectedOutput.add(new StreamRecord<>("1777"));
	expectedOutput.add(new StreamRecord<>("1777"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

private static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> getInitializedTestHarness(
		final BroadcastProcessFunction<IN1, IN2, OUT> function,
		final MapStateDescriptor<?, ?>... descriptors) throws Exception {

	return getInitializedTestHarness(
			function,
			1,
			1,
			0,
			descriptors);
}
 

@Test
public void nonKeyedParDoSideInputCheckpointing() throws Exception {
  sideInputCheckpointing(
      () -> {
        Coder<WindowedValue<String>> coder =
            WindowedValue.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder());
        TupleTag<String> outputTag = new TupleTag<>("main-output");

        ImmutableMap<Integer, PCollectionView<?>> sideInputMapping =
            ImmutableMap.<Integer, PCollectionView<?>>builder()
                .put(1, view1)
                .put(2, view2)
                .build();

        DoFnOperator<String, String> doFnOperator =
            new DoFnOperator<>(
                new IdentityDoFn<>(),
                "stepName",
                coder,
                Collections.emptyMap(),
                outputTag,
                Collections.emptyList(),
                new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, coder),
                WindowingStrategy.globalDefault(),
                sideInputMapping, /* side-input mapping */
                ImmutableList.of(view1, view2), /* side inputs */
                PipelineOptionsFactory.as(FlinkPipelineOptions.class),
                null,
                null,
                DoFnSchemaInformation.create(),
                Collections.emptyMap());

        return new TwoInputStreamOperatorTestHarness<>(doFnOperator);
      });
}
 

@Test
public void testGetCurrentKeyFromContext() throws Exception {
	KeyedCoProcessOperator<String, Integer, String, String> operator =
			new KeyedCoProcessOperator<>(new AppendCurrentKeyProcessFunction());

	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));
	testHarness.processElement1(new StreamRecord<>(6));
	testHarness.processElement2(new StreamRecord<>("hello"));
	testHarness.processElement2(new StreamRecord<>("world"));

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("5,5"));
	expectedOutput.add(new StreamRecord<>("6,6"));
	expectedOutput.add(new StreamRecord<>("hello,hello"));
	expectedOutput.add(new StreamRecord<>("world,world"));

	TestHarnessUtil.assertOutputEquals(
			"Output was not correct.",
			expectedOutput,
			testHarness.getOutput());

	testHarness.close();
}
 

@Test
public void testTimestampAndWatermarkQuerying() throws Exception {

	LegacyKeyedCoProcessOperator<String, Integer, String, String> operator =
			new LegacyKeyedCoProcessOperator<>(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();
}
 

@Test
public void testTimestampAndProcessingTimeQuerying() throws Exception {

	LegacyKeyedCoProcessOperator<String, Integer, String, String> operator =
			new LegacyKeyedCoProcessOperator<>(new ProcessingTimeQueryingProcessFunction());

	TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness =
			new KeyedTwoInputStreamOperatorTestHarness<>(
					operator,
					new IntToStringKeySelector<>(),
					new IdentityKeySelector<String>(),
					BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.setup();
	testHarness.open();

	testHarness.setProcessingTime(17);
	testHarness.processElement1(new StreamRecord<>(5));

	testHarness.setProcessingTime(42);
	testHarness.processElement2(new StreamRecord<>("6"));

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("5PT:17 TS:null"));
	expectedOutput.add(new StreamRecord<>("6PT:42 TS:null"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@Test
public void testReadBroadcastState() throws Exception {
	try (TwoInputStreamOperatorTestHarness<Void, Integer, Void> testHarness = getTestHarness()) {
		testHarness.open();

		testHarness.processElement2(new StreamRecord<>(1));
		testHarness.processElement2(new StreamRecord<>(2));
		testHarness.processElement2(new StreamRecord<>(3));

		OperatorSubtaskState subtaskState = testHarness.snapshot(0, 0);
		OperatorState state = new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4);
		state.putState(0, subtaskState);

		OperatorStateInputSplit split = new OperatorStateInputSplit(subtaskState.getManagedOperatorState(), 0);

		BroadcastStateInputFormat<Integer, Integer> format = new BroadcastStateInputFormat<>(state, descriptor);

		format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
		format.open(split);

		Map<Integer, Integer> results = new HashMap<>(3);

		while (!format.reachedEnd()) {
			Tuple2<Integer, Integer> entry = format.nextRecord(null);
			results.put(entry.f0, entry.f1);
		}

		Map<Integer, Integer> expected = new HashMap<>(3);
		expected.put(1, 1);
		expected.put(2, 2);
		expected.put(3, 3);

		Assert.assertEquals("Failed to read correct list state from state backend", expected, results);
	}
}
 

@Test
public void testProcessingTimeTimers() throws Exception {

	KeyedCoProcessOperator<String, Integer, String, String> operator =
			new KeyedCoProcessOperator<>(new ProcessingTimeTriggeringProcessFunction());

	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));
	testHarness.processElement2(new StreamRecord<>("18"));

	testHarness.setProcessingTime(5);
	testHarness.setProcessingTime(6);

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("INPUT1:17"));
	expectedOutput.add(new StreamRecord<>("INPUT2:18"));
	expectedOutput.add(new StreamRecord<>("1777"));
	expectedOutput.add(new StreamRecord<>("1777"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@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<>("17:1777", 5L));
	expectedOutput.add(new Watermark(5L));
	expectedOutput.add(new StreamRecord<>("18:1777", 6L));
	expectedOutput.add(new Watermark(6L));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@Test
public void testTimestampAndProcessingTimeQuerying() throws Exception {

	KeyedCoProcessOperator<String, Integer, String, String> operator =
			new KeyedCoProcessOperator<>(new ProcessingTimeQueryingProcessFunction());

	TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness =
			new KeyedTwoInputStreamOperatorTestHarness<>(
					operator,
					new IntToStringKeySelector<>(),
					new IdentityKeySelector<String>(),
					BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.setup();
	testHarness.open();

	testHarness.setProcessingTime(17);
	testHarness.processElement1(new StreamRecord<>(5));

	testHarness.setProcessingTime(42);
	testHarness.processElement2(new StreamRecord<>("6"));

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("5PT:17 TS:null"));
	expectedOutput.add(new StreamRecord<>("6PT:42 TS:null"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@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();
}
 

@Test
public void testProcessingTimeTimers() throws Exception {

	LegacyKeyedCoProcessOperator<String, Integer, String, String> operator =
			new LegacyKeyedCoProcessOperator<>(new ProcessingTimeTriggeringProcessFunction());

	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));
	testHarness.processElement2(new StreamRecord<>("18"));

	testHarness.setProcessingTime(5);
	testHarness.setProcessingTime(6);

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("INPUT1:17"));
	expectedOutput.add(new StreamRecord<>("INPUT2:18"));
	expectedOutput.add(new StreamRecord<>("1777"));
	expectedOutput.add(new StreamRecord<>("1777"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@Test
public void testSideOutput() throws Exception {
	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					BasicTypeInfo.STRING_TYPE_INFO,
					new IdentityKeySelector<>(),
					new FunctionWithSideOutput())
	) {

		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>("6", 13L));
		testHarness.processElement1(new StreamRecord<>("6", 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		Queue<StreamRecord<String>> expectedBr = new ConcurrentLinkedQueue<>();
		expectedBr.add(new StreamRecord<>("BR:5 WM:10 TS:12", 12L));

		Queue<StreamRecord<String>> expectedNonBr = new ConcurrentLinkedQueue<>();
		expectedNonBr.add(new StreamRecord<>("NON-BR:6 WM:40 TS:13", 13L));
		expectedNonBr.add(new StreamRecord<>("NON-BR:6 WM:40 TS:15", 15L));

		TestHarnessUtil.assertOutputEquals(
				"Wrong Side Output",
				expectedBr,
				testHarness.getSideOutput(FunctionWithSideOutput.BROADCAST_TAG));

		TestHarnessUtil.assertOutputEquals(
				"Wrong Side Output",
				expectedNonBr,
				testHarness.getSideOutput(FunctionWithSideOutput.NON_BROADCAST_TAG));
	}
}
 

@Test
public void testFunctionWithTimer() throws Exception {
	final String expectedKey = "6";

	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					BasicTypeInfo.STRING_TYPE_INFO,
					new IdentityKeySelector<>(),
					new FunctionWithTimerOnKeyed(41L, expectedKey))
	) {
		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>(expectedKey, 13L));
		testHarness.processElement1(new StreamRecord<>(expectedKey, 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		Queue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

		expectedOutput.add(new Watermark(10L));
		expectedOutput.add(new StreamRecord<>("BR:5 WM:10 TS:12", 12L));
		expectedOutput.add(new Watermark(40L));
		expectedOutput.add(new StreamRecord<>("NON-BR:6 WM:40 TS:13", 13L));
		expectedOutput.add(new StreamRecord<>("NON-BR:6 WM:40 TS:15", 15L));
		expectedOutput.add(new StreamRecord<>("TIMER:41", 41L));
		expectedOutput.add(new Watermark(50L));

		TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
	}
}
 

/**
 * Verifies that we don't have leakage between different keys.
 */
@Test
public void testProcessingTimeTimerWithState() throws Exception {

	LegacyKeyedCoProcessOperator<String, Integer, String, String> operator =
			new LegacyKeyedCoProcessOperator<>(new ProcessingTimeTriggeringStatefulProcessFunction());

	TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness =
			new KeyedTwoInputStreamOperatorTestHarness<>(
					operator,
					new IntToStringKeySelector<>(),
					new IdentityKeySelector<String>(),
					BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.setup();
	testHarness.open();

	testHarness.setProcessingTime(1);
	testHarness.processElement1(new StreamRecord<>(17)); // should set timer for 6
	testHarness.processElement1(new StreamRecord<>(13)); // should set timer for 6

	testHarness.setProcessingTime(2);
	testHarness.processElement1(new StreamRecord<>(13)); // should delete timer again
	testHarness.processElement2(new StreamRecord<>("42")); // should set timer for 7

	testHarness.setProcessingTime(6);
	testHarness.setProcessingTime(7);

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("INPUT1:17"));
	expectedOutput.add(new StreamRecord<>("INPUT1:13"));
	expectedOutput.add(new StreamRecord<>("INPUT2:42"));
	expectedOutput.add(new StreamRecord<>("STATE:17"));
	expectedOutput.add(new StreamRecord<>("STATE:42"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@Test
public void testTimestampAndProcessingTimeQuerying() throws Exception {

	KeyedCoProcessOperator<String, Integer, String, String> operator =
			new KeyedCoProcessOperator<>(new ProcessingTimeQueryingProcessFunction());

	TwoInputStreamOperatorTestHarness<Integer, String, String> testHarness =
			new KeyedTwoInputStreamOperatorTestHarness<>(
					operator,
					new IntToStringKeySelector<>(),
					new IdentityKeySelector<String>(),
					BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.setup();
	testHarness.open();

	testHarness.setProcessingTime(17);
	testHarness.processElement1(new StreamRecord<>(5));

	testHarness.setProcessingTime(42);
	testHarness.processElement2(new StreamRecord<>("6"));

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

	expectedOutput.add(new StreamRecord<>("5PT:17 TS:null"));
	expectedOutput.add(new StreamRecord<>("6PT:42 TS:null"));

	TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());

	testHarness.close();
}
 

@Test
public void testEventTimeTimers() throws Exception {

	LegacyKeyedCoProcessOperator<String, Integer, String, String> operator =
			new LegacyKeyedCoProcessOperator<>(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();
}
 

@Test
public void testFunctionWithTimer() throws Exception {
	final String expectedKey = "6";

	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					BasicTypeInfo.STRING_TYPE_INFO,
					new IdentityKeySelector<>(),
					new FunctionWithTimerOnKeyed(41L, expectedKey))
	) {
		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>(expectedKey, 13L));
		testHarness.processElement1(new StreamRecord<>(expectedKey, 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		Queue<Object> expectedOutput = new ConcurrentLinkedQueue<>();

		expectedOutput.add(new Watermark(10L));
		expectedOutput.add(new StreamRecord<>("BR:5 WM:10 TS:12", 12L));
		expectedOutput.add(new Watermark(40L));
		expectedOutput.add(new StreamRecord<>("NON-BR:6 WM:40 TS:13", 13L));
		expectedOutput.add(new StreamRecord<>("NON-BR:6 WM:40 TS:15", 15L));
		expectedOutput.add(new StreamRecord<>("TIMER:41", 41L));
		expectedOutput.add(new Watermark(50L));

		TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
	}
}
 

@Test
public void testSideOutput() throws Exception {
	try (
			TwoInputStreamOperatorTestHarness<String, Integer, String> testHarness = getInitializedTestHarness(
					BasicTypeInfo.STRING_TYPE_INFO,
					new IdentityKeySelector<>(),
					new FunctionWithSideOutput())
	) {

		testHarness.processWatermark1(new Watermark(10L));
		testHarness.processWatermark2(new Watermark(10L));
		testHarness.processElement2(new StreamRecord<>(5, 12L));

		testHarness.processWatermark1(new Watermark(40L));
		testHarness.processWatermark2(new Watermark(40L));
		testHarness.processElement1(new StreamRecord<>("6", 13L));
		testHarness.processElement1(new StreamRecord<>("6", 15L));

		testHarness.processWatermark1(new Watermark(50L));
		testHarness.processWatermark2(new Watermark(50L));

		Queue<StreamRecord<String>> expectedBr = new ConcurrentLinkedQueue<>();
		expectedBr.add(new StreamRecord<>("BR:5 WM:10 TS:12", 12L));

		Queue<StreamRecord<String>> expectedNonBr = new ConcurrentLinkedQueue<>();
		expectedNonBr.add(new StreamRecord<>("NON-BR:6 WM:40 TS:13", 13L));
		expectedNonBr.add(new StreamRecord<>("NON-BR:6 WM:40 TS:15", 15L));

		TestHarnessUtil.assertOutputEquals(
				"Wrong Side Output",
				expectedBr,
				testHarness.getSideOutput(FunctionWithSideOutput.BROADCAST_TAG));

		TestHarnessUtil.assertOutputEquals(
				"Wrong Side Output",
				expectedNonBr,
				testHarness.getSideOutput(FunctionWithSideOutput.NON_BROADCAST_TAG));
	}
}