Java Code Examples for org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness#close()

private void closeHarness(KeyedOneInputStreamOperatorTestHarness<?, ?, ?>[] harnessArr) throws Exception {
	for (KeyedOneInputStreamOperatorTestHarness<?, ?, ?> harness : harnessArr) {
		if (harness != null) {
			harness.close();
		}
	}
}
 

/**
 * Ensure that we get some output from the given operator when pushing in an element and
 * setting watermark and processing time to {@code Long.MAX_VALUE}.
 */
private static <K, IN, OUT> void processElementAndEnsureOutput(
		OneInputStreamOperator<IN, OUT> operator,
		KeySelector<IN, K> keySelector,
		TypeInformation<K> keyType,
		IN element) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
			new KeyedOneInputStreamOperatorTestHarness<>(
					operator,
					keySelector,
					keyType);

	if (operator instanceof OutputTypeConfigurable) {
		// use a dummy type since window functions just need the ExecutionConfig
		// this is also only needed for Fold, which we're getting rid off soon.
		((OutputTypeConfigurable) operator).setOutputType(BasicTypeInfo.STRING_TYPE_INFO, new ExecutionConfig());
	}

	testHarness.open();

	testHarness.setProcessingTime(0);
	testHarness.processWatermark(Long.MIN_VALUE);

	testHarness.processElement(new StreamRecord<>(element, 0));

	// provoke any processing-time/event-time triggers
	testHarness.setProcessingTime(Long.MAX_VALUE);
	testHarness.processWatermark(Long.MAX_VALUE);

	// we at least get the two watermarks and should also see an output element
	assertTrue(testHarness.getOutput().size() >= 3);

	testHarness.close();
}
 

/**
 * Ensure that we get some output from the given operator when pushing in an element and
 * setting watermark and processing time to {@code Long.MAX_VALUE}.
 */
private static <K, IN, OUT> void processElementAndEnsureOutput(
		OneInputStreamOperator<IN, OUT> operator,
		KeySelector<IN, K> keySelector,
		TypeInformation<K> keyType,
		IN element) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
			new KeyedOneInputStreamOperatorTestHarness<>(
					operator,
					keySelector,
					keyType);

	testHarness.open();

	testHarness.setProcessingTime(0);
	testHarness.processWatermark(Long.MIN_VALUE);

	testHarness.processElement(new StreamRecord<>(element, 0));

	// provoke any processing-time/event-time triggers
	testHarness.setProcessingTime(Long.MAX_VALUE);
	testHarness.processWatermark(Long.MAX_VALUE);

	// we at least get the two watermarks and should also see an output element
	assertTrue(testHarness.getOutput().size() >= 3);

	testHarness.close();
}
 

private void closeHarness(KeyedOneInputStreamOperatorTestHarness<?, ?, ?>[] harnessArr) throws Exception {
	for (KeyedOneInputStreamOperatorTestHarness<?, ?, ?> harness : harnessArr) {
		if (harness != null) {
			harness.close();
		}
	}
}
 

/**
 * Ensure that we get some output from the given operator when pushing in an element and
 * setting watermark and processing time to {@code Long.MAX_VALUE}.
 */
private static <K, IN, OUT> void processElementAndEnsureOutput(
		OneInputStreamOperator<IN, OUT> operator,
		KeySelector<IN, K> keySelector,
		TypeInformation<K> keyType,
		IN element) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
			new KeyedOneInputStreamOperatorTestHarness<>(
					operator,
					keySelector,
					keyType);

	if (operator instanceof OutputTypeConfigurable) {
		// use a dummy type since window functions just need the ExecutionConfig
		// this is also only needed for Fold, which we're getting rid off soon.
		((OutputTypeConfigurable) operator).setOutputType(BasicTypeInfo.STRING_TYPE_INFO, new ExecutionConfig());
	}

	testHarness.open();

	testHarness.setProcessingTime(0);
	testHarness.processWatermark(Long.MIN_VALUE);

	testHarness.processElement(new StreamRecord<>(element, 0));

	// provoke any processing-time/event-time triggers
	testHarness.setProcessingTime(Long.MAX_VALUE);
	testHarness.processWatermark(Long.MAX_VALUE);

	// we at least get the two watermarks and should also see an output element
	assertTrue(testHarness.getOutput().size() >= 3);

	testHarness.close();
}
 

/**
 * Ensure that we get some output from the given operator when pushing in an element and
 * setting watermark and processing time to {@code Long.MAX_VALUE}.
 */
private static <K, IN, OUT> void processElementAndEnsureOutput(
		OneInputStreamOperator<IN, OUT> operator,
		KeySelector<IN, K> keySelector,
		TypeInformation<K> keyType,
		IN element) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
			new KeyedOneInputStreamOperatorTestHarness<>(
					operator,
					keySelector,
					keyType);

	testHarness.open();

	testHarness.setProcessingTime(0);
	testHarness.processWatermark(Long.MIN_VALUE);

	testHarness.processElement(new StreamRecord<>(element, 0));

	// provoke any processing-time/event-time triggers
	testHarness.setProcessingTime(Long.MAX_VALUE);
	testHarness.processWatermark(Long.MAX_VALUE);

	// we at least get the two watermarks and should also see an output element
	assertTrue(testHarness.getOutput().size() >= 3);

	testHarness.close();
}
 

private void closeHarness(KeyedOneInputStreamOperatorTestHarness<?, ?, ?>[] harnessArr) throws Exception {
	for (KeyedOneInputStreamOperatorTestHarness<?, ?, ?> harness : harnessArr) {
		if (harness != null) {
			harness.close();
		}
	}
}
 

/**
 * Ensure that we get some output from the given operator when pushing in an element and
 * setting watermark and processing time to {@code Long.MAX_VALUE}.
 */
private static <K, IN, OUT> void processElementAndEnsureOutput(
		OneInputStreamOperator<IN, OUT> operator,
		KeySelector<IN, K> keySelector,
		TypeInformation<K> keyType,
		IN element) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
			new KeyedOneInputStreamOperatorTestHarness<>(
					operator,
					keySelector,
					keyType);

	if (operator instanceof OutputTypeConfigurable) {
		// use a dummy type since window functions just need the ExecutionConfig
		// this is also only needed for Fold, which we're getting rid off soon.
		((OutputTypeConfigurable) operator).setOutputType(BasicTypeInfo.STRING_TYPE_INFO, new ExecutionConfig());
	}

	testHarness.open();

	testHarness.setProcessingTime(0);
	testHarness.processWatermark(Long.MIN_VALUE);

	testHarness.processElement(new StreamRecord<>(element, 0));

	// provoke any processing-time/event-time triggers
	testHarness.setProcessingTime(Long.MAX_VALUE);
	testHarness.processWatermark(Long.MAX_VALUE);

	// we at least get the two watermarks and should also see an output element
	assertTrue(testHarness.getOutput().size() >= 3);

	testHarness.close();
}
 

/**
 * Ensure that we get some output from the given operator when pushing in an element and
 * setting watermark and processing time to {@code Long.MAX_VALUE}.
 */
private static <K, IN, OUT> void processElementAndEnsureOutput(
		OneInputStreamOperator<IN, OUT> operator,
		KeySelector<IN, K> keySelector,
		TypeInformation<K> keyType,
		IN element) throws Exception {

	KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness =
			new KeyedOneInputStreamOperatorTestHarness<>(
					operator,
					keySelector,
					keyType);

	testHarness.open();

	testHarness.setProcessingTime(0);
	testHarness.processWatermark(Long.MIN_VALUE);

	testHarness.processElement(new StreamRecord<>(element, 0));

	// provoke any processing-time/event-time triggers
	testHarness.setProcessingTime(Long.MAX_VALUE);
	testHarness.processWatermark(Long.MAX_VALUE);

	// we at least get the two watermarks and should also see an output element
	assertTrue(testHarness.getOutput().size() >= 3);

	testHarness.close();
}
 

/**
 * Tests that the internal time of a CEP operator advances only given watermarks. See FLINK-5033
 */
@Test
public void testKeyedAdvancingTimeWithoutElements() throws Exception {
	final Event startEvent = new Event(42, "start", 1.0);
	final long watermarkTimestamp1 = 5L;
	final long watermarkTimestamp2 = 13L;

	final Map<String, List<Event>> expectedSequence = new HashMap<>(2);
	expectedSequence.put("start", Collections.<Event>singletonList(startEvent));

	final OutputTag<Tuple2<Map<String, List<Event>>, Long>> timedOut =
		new OutputTag<Tuple2<Map<String, List<Event>>, Long>>("timedOut") {};
	final KeyedOneInputStreamOperatorTestHarness<Integer, Event, Map<String, List<Event>>> harness =
		new KeyedOneInputStreamOperatorTestHarness<>(
			new CepOperator<>(
				Event.createTypeSerializer(),
				false,
				new NFAFactory(true),
				null,
				null,
				new TimedOutProcessFunction(timedOut),
				null), new KeySelector<Event, Integer>() {
			private static final long serialVersionUID = 7219185117566268366L;

			@Override
			public Integer getKey(Event value) throws Exception {
				return value.getId();
			}
		}, BasicTypeInfo.INT_TYPE_INFO);

	try {
		String rocksDbPath = tempFolder.newFolder().getAbsolutePath();
		RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(new MemoryStateBackend());
		rocksDBStateBackend.setDbStoragePath(rocksDbPath);

		harness.setStateBackend(rocksDBStateBackend);
		harness.setup(
			new KryoSerializer<>(
				(Class<Map<String, List<Event>>>) (Object) Map.class,
				new ExecutionConfig()));
		harness.open();

		harness.processElement(new StreamRecord<>(startEvent, 3L));
		harness.processWatermark(new Watermark(watermarkTimestamp1));
		harness.processWatermark(new Watermark(watermarkTimestamp2));

		Queue<Object> result = harness.getOutput();
		Queue<StreamRecord<Tuple2<Map<String, List<Event>>, Long>>> sideOutput = harness.getSideOutput(timedOut);

		assertEquals(2L, result.size());
		assertEquals(1L, sideOutput.size());

		Object watermark1 = result.poll();

		assertTrue(watermark1 instanceof Watermark);

		assertEquals(watermarkTimestamp1, ((Watermark) watermark1).getTimestamp());

		Tuple2<Map<String, List<Event>>, Long> leftResult = sideOutput.poll().getValue();

		assertEquals(watermarkTimestamp2, (long) leftResult.f1);
		assertEquals(expectedSequence, leftResult.f0);

		Object watermark2 = result.poll();

		assertTrue(watermark2 instanceof Watermark);

		assertEquals(watermarkTimestamp2, ((Watermark) watermark2).getTimestamp());
	} finally {
		harness.close();
	}
}
 

/**
 * Tests that the internal time of a CEP operator advances only given watermarks. See FLINK-5033
 */
@Test
public void testKeyedAdvancingTimeWithoutElements() throws Exception {
	final Event startEvent = new Event(42, "start", 1.0);
	final long watermarkTimestamp1 = 5L;
	final long watermarkTimestamp2 = 13L;

	final Map<String, List<Event>> expectedSequence = new HashMap<>(2);
	expectedSequence.put("start", Collections.<Event>singletonList(startEvent));

	final OutputTag<Tuple2<Map<String, List<Event>>, Long>> timedOut =
		new OutputTag<Tuple2<Map<String, List<Event>>, Long>>("timedOut") {};
	final KeyedOneInputStreamOperatorTestHarness<Integer, Event, Map<String, List<Event>>> harness =
		new KeyedOneInputStreamOperatorTestHarness<>(
			new CepOperator<>(
				Event.createTypeSerializer(),
				false,
				new NFAFactory(true),
				null,
				null,
				new TimedOutProcessFunction(timedOut),
				null), new KeySelector<Event, Integer>() {
			private static final long serialVersionUID = 7219185117566268366L;

			@Override
			public Integer getKey(Event value) throws Exception {
				return value.getId();
			}
		}, BasicTypeInfo.INT_TYPE_INFO);

	try {
		String rocksDbPath = tempFolder.newFolder().getAbsolutePath();
		RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(new MemoryStateBackend());
		rocksDBStateBackend.setDbStoragePath(rocksDbPath);

		harness.setStateBackend(rocksDBStateBackend);
		harness.setup(
			new KryoSerializer<>(
				(Class<Map<String, List<Event>>>) (Object) Map.class,
				new ExecutionConfig()));
		harness.open();

		harness.processElement(new StreamRecord<>(startEvent, 3L));
		harness.processWatermark(new Watermark(watermarkTimestamp1));
		harness.processWatermark(new Watermark(watermarkTimestamp2));

		Queue<Object> result = harness.getOutput();
		Queue<StreamRecord<Tuple2<Map<String, List<Event>>, Long>>> sideOutput = harness.getSideOutput(timedOut);

		assertEquals(2L, result.size());
		assertEquals(1L, sideOutput.size());

		Object watermark1 = result.poll();

		assertTrue(watermark1 instanceof Watermark);

		assertEquals(watermarkTimestamp1, ((Watermark) watermark1).getTimestamp());

		Tuple2<Map<String, List<Event>>, Long> leftResult = sideOutput.poll().getValue();

		assertEquals(watermarkTimestamp2, (long) leftResult.f1);
		assertEquals(expectedSequence, leftResult.f0);

		Object watermark2 = result.poll();

		assertTrue(watermark2 instanceof Watermark);

		assertEquals(watermarkTimestamp2, ((Watermark) watermark2).getTimestamp());
	} finally {
		harness.close();
	}
}
 

@Test
public void testTimerCleanupOfPendingTimerList() throws Exception {
  // test harness
  WindowDoFnOperator<Long, Long, Long> windowDoFnOperator = getWindowDoFnOperator();
  KeyedOneInputStreamOperatorTestHarness<
          ByteBuffer, WindowedValue<KeyedWorkItem<Long, Long>>, WindowedValue<KV<Long, Long>>>
      testHarness = createTestHarness(windowDoFnOperator);
  testHarness.open();

  DoFnOperator<KeyedWorkItem<Long, Long>, KV<Long, Long>>.FlinkTimerInternals timerInternals =
      windowDoFnOperator.timerInternals;

  // process elements
  IntervalWindow window = new IntervalWindow(new Instant(0), Duration.millis(100));
  IntervalWindow window2 = new IntervalWindow(new Instant(100), Duration.millis(100));
  testHarness.processWatermark(0L);

  // Use two different keys to check for correct watermark hold calculation
  testHarness.processElement(
      Item.builder().key(1L).timestamp(1L).value(100L).window(window).build().toStreamRecord());
  testHarness.processElement(
      Item.builder()
          .key(2L)
          .timestamp(150L)
          .value(150L)
          .window(window2)
          .build()
          .toStreamRecord());

  testHarness.processWatermark(1);

  // Note that the following is 1 because the state is key-partitioned
  assertThat(Iterables.size(timerInternals.pendingTimersById.keys()), is(1));

  assertThat(testHarness.numKeyedStateEntries(), is(6));
  assertThat(windowDoFnOperator.getCurrentOutputWatermark(), is(1L));
  assertThat(timerInternals.getMinOutputTimestampMs(), is(Long.MAX_VALUE));

  // close window
  testHarness.processWatermark(100L);

  // Note that the following is zero because we only the first key is active
  assertThat(Iterables.size(timerInternals.pendingTimersById.keys()), is(0));

  assertThat(testHarness.numKeyedStateEntries(), is(3));
  assertThat(windowDoFnOperator.getCurrentOutputWatermark(), is(100L));
  assertThat(timerInternals.getMinOutputTimestampMs(), is(Long.MAX_VALUE));

  testHarness.processWatermark(200L);

  // All the state has been cleaned up
  assertThat(testHarness.numKeyedStateEntries(), is(0));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      containsInAnyOrder(
          WindowedValue.of(
              KV.of(1L, 100L), new Instant(99), window, PaneInfo.createPane(true, true, ON_TIME)),
          WindowedValue.of(
              KV.of(2L, 150L),
              new Instant(199),
              window2,
              PaneInfo.createPane(true, true, ON_TIME))));

  // cleanup
  testHarness.close();
}
 

@Test
public void testDeduping() throws Exception {

  KeyedOneInputStreamOperatorTestHarness<
          ByteBuffer, WindowedValue<ValueWithRecordId<String>>, WindowedValue<String>>
      harness = getDebupingHarness();

  harness.open();

  String key1 = "key1";
  String key2 = "key2";

  harness.processElement(
      new StreamRecord<>(
          WindowedValue.valueInGlobalWindow(
              new ValueWithRecordId<>(key1, key1.getBytes(StandardCharsets.UTF_8)))));

  harness.processElement(
      new StreamRecord<>(
          WindowedValue.valueInGlobalWindow(
              new ValueWithRecordId<>(key2, key2.getBytes(StandardCharsets.UTF_8)))));

  harness.processElement(
      new StreamRecord<>(
          WindowedValue.valueInGlobalWindow(
              new ValueWithRecordId<>(key1, key1.getBytes(StandardCharsets.UTF_8)))));

  assertThat(
      stripStreamRecordFromWindowedValue(harness.getOutput()),
      contains(WindowedValue.valueInGlobalWindow(key1), WindowedValue.valueInGlobalWindow(key2)));

  OperatorSubtaskState snapshot = harness.snapshot(0L, 0L);

  harness.close();

  harness = getDebupingHarness();
  harness.setup();
  harness.initializeState(snapshot);
  harness.open();

  String key3 = "key3";

  harness.processElement(
      new StreamRecord<>(
          WindowedValue.valueInGlobalWindow(
              new ValueWithRecordId<>(key2, key2.getBytes(StandardCharsets.UTF_8)))));

  harness.processElement(
      new StreamRecord<>(
          WindowedValue.valueInGlobalWindow(
              new ValueWithRecordId<>(key3, key3.getBytes(StandardCharsets.UTF_8)))));

  assertThat(
      stripStreamRecordFromWindowedValue(harness.getOutput()),
      contains(WindowedValue.valueInGlobalWindow(key3)));

  harness.close();
}
 

/**
 * Tests that the internal time of a CEP operator advances only given watermarks. See FLINK-5033
 */
@Test
public void testKeyedAdvancingTimeWithoutElements() throws Exception {
	final Event startEvent = new Event(42, "start", 1.0);
	final long watermarkTimestamp1 = 5L;
	final long watermarkTimestamp2 = 13L;

	final Map<String, List<Event>> expectedSequence = new HashMap<>(2);
	expectedSequence.put("start", Collections.<Event>singletonList(startEvent));

	final OutputTag<Tuple2<Map<String, List<Event>>, Long>> timedOut =
		new OutputTag<Tuple2<Map<String, List<Event>>, Long>>("timedOut") {};
	final KeyedOneInputStreamOperatorTestHarness<Integer, Event, Map<String, List<Event>>> harness =
		new KeyedOneInputStreamOperatorTestHarness<>(
			new CepOperator<>(
				Event.createTypeSerializer(),
				false,
				new NFAFactory(true),
				null,
				null,
				new TimedOutProcessFunction(timedOut),
				null), new KeySelector<Event, Integer>() {
			private static final long serialVersionUID = 7219185117566268366L;

			@Override
			public Integer getKey(Event value) throws Exception {
				return value.getId();
			}
		}, BasicTypeInfo.INT_TYPE_INFO);

	try {
		String rocksDbPath = tempFolder.newFolder().getAbsolutePath();
		RocksDBStateBackend rocksDBStateBackend = new RocksDBStateBackend(new MemoryStateBackend());
		rocksDBStateBackend.setDbStoragePath(rocksDbPath);

		harness.setStateBackend(rocksDBStateBackend);
		harness.setup(
			new KryoSerializer<>(
				(Class<Map<String, List<Event>>>) (Object) Map.class,
				new ExecutionConfig()));
		harness.open();

		harness.processElement(new StreamRecord<>(startEvent, 3L));
		harness.processWatermark(new Watermark(watermarkTimestamp1));
		harness.processWatermark(new Watermark(watermarkTimestamp2));

		Queue<Object> result = harness.getOutput();
		Queue<StreamRecord<Tuple2<Map<String, List<Event>>, Long>>> sideOutput = harness.getSideOutput(timedOut);

		assertEquals(2L, result.size());
		assertEquals(1L, sideOutput.size());

		Object watermark1 = result.poll();

		assertTrue(watermark1 instanceof Watermark);

		assertEquals(watermarkTimestamp1, ((Watermark) watermark1).getTimestamp());

		Tuple2<Map<String, List<Event>>, Long> leftResult = sideOutput.poll().getValue();

		assertEquals(watermarkTimestamp2, (long) leftResult.f1);
		assertEquals(expectedSequence, leftResult.f0);

		Object watermark2 = result.poll();

		assertTrue(watermark2 instanceof Watermark);

		assertEquals(watermarkTimestamp2, ((Watermark) watermark2).getTimestamp());
	} finally {
		harness.close();
	}
}