org.apache.flink.streaming.runtime.tasks.StreamTaskTestHarness Java Examples

private StreamTaskTestHarness<?> startTestHarness() throws Exception {
	final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(
			OneInputStreamTask::new,
			BasicTypeInfo.STRING_TYPE_INFO,
			BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.setupOutputForSingletonOperatorChain();

	StreamConfig streamConfig = testHarness.getStreamConfig();
	streamConfig.setChainIndex(0);
	streamConfig.setStreamOperator(new StreamMap<String, String>(new DummyMapFunction<>()));

	testHarness.invoke();
	testHarness.waitForTaskRunning();

	return testHarness;
}
 

@Test
public void testMaxWatermarkIsForwardedLastForFiniteSource() throws Exception {
	StreamSource<String, ?> sourceOperator = new StreamSource<>(new FiniteSource(true));
	StreamTaskTestHarness<String> testHarness = setupSourceStreamTask(sourceOperator, BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.invoke();
	testHarness.waitForTaskCompletion();

	ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
	expectedOutput.add(new StreamRecord<>("Hello"));
	expectedOutput.add(Watermark.MAX_WATERMARK);

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

@Test
public void testNoMaxWatermarkOnImmediateCancel() throws Exception {
	StreamSource<String, ?> sourceOperator = new StreamSource<>(new InfiniteSource<>());
	StreamTaskTestHarness<String> testHarness = setupSourceStreamTask(
		sourceOperator, BasicTypeInfo.STRING_TYPE_INFO, true);

	testHarness.invoke();
	try {
		testHarness.waitForTaskCompletion();
		fail("should throw an exception");
	} catch (Throwable t) {
		if (!ExceptionUtils.findThrowable(t, CancelTaskException.class).isPresent()) {
			throw t;
		}
	}
	assertTrue(testHarness.getOutput().isEmpty());
}
 

@Test
public void testNoMaxWatermarkOnAsyncCancel() throws Exception {
	StreamSource<String, ?> sourceOperator = new StreamSource<>(new InfiniteSource<>());
	StreamTaskTestHarness<String> testHarness = setupSourceStreamTask(sourceOperator, BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.invoke();
	testHarness.waitForTaskRunning();
	Thread.sleep(200);
	testHarness.getTask().cancel();
	try {
		testHarness.waitForTaskCompletion();
	} catch (Throwable t) {
		if (!ExceptionUtils.findThrowable(t, CancelTaskException.class).isPresent()) {
			throw t;
		}
	}
	assertTrue(testHarness.getOutput().isEmpty());
}
 

@Test
public void testIterationHeadWatermarkEmission() throws Exception {
	StreamTaskTestHarness<Integer> harness = new StreamTaskTestHarness<>(
			StreamIterationHead::new,
			BasicTypeInfo.INT_TYPE_INFO);
	harness.setupOutputForSingletonOperatorChain();
	harness.getStreamConfig().setIterationId("1");
	harness.getStreamConfig().setIterationWaitTime(1);

	harness.invoke();
	harness.waitForTaskCompletion();

	assertEquals(1, harness.getOutput().size());
	assertEquals(new Watermark(Long.MAX_VALUE), harness.getOutput().peek());
}
 

@Test
public void testIterationHeadWatermarkEmission() throws Exception {
	StreamTaskTestHarness<Integer> harness = new StreamTaskTestHarness<>(
			StreamIterationHead::new,
			BasicTypeInfo.INT_TYPE_INFO);
	harness.setupOutputForSingletonOperatorChain();
	harness.getStreamConfig().setIterationId("1");
	harness.getStreamConfig().setIterationWaitTime(1);

	harness.invoke();
	harness.waitForTaskCompletion();

	assertEquals(1, harness.getOutput().size());
	assertEquals(new Watermark(Long.MAX_VALUE), harness.getOutput().peek());
}
 

@Test
public void testIterationHeadWatermarkEmission() throws Exception {
	StreamTaskTestHarness<Integer> harness = new StreamTaskTestHarness<>(
			StreamIterationHead::new,
			BasicTypeInfo.INT_TYPE_INFO);
	harness.setupOutputForSingletonOperatorChain();
	harness.getStreamConfig().setIterationId("1");
	harness.getStreamConfig().setIterationWaitTime(1);

	harness.invoke();
	harness.waitForTaskCompletion();

	assertEquals(1, harness.getOutput().size());
	assertEquals(new Watermark(Long.MAX_VALUE), harness.getOutput().peek());
}
 

private void stopTestHarness(StreamTaskTestHarness<?> testHarness, long timeout) throws Exception {
	testHarness.endInput();
	testHarness.waitForTaskCompletion();

	// thread needs to die in time
	long deadline = System.currentTimeMillis() + timeout;
	while (StreamTask.TRIGGER_THREAD_GROUP.activeCount() > 0 && System.currentTimeMillis() < deadline) {
		Thread.sleep(10);
	}

	assertEquals("Trigger timer thread did not properly shut down",
			0, StreamTask.TRIGGER_THREAD_GROUP.activeCount());
}
 

@Test
public void testEmitMaxWatermarkForFiniteSource() throws Exception {
	StreamSource<String, ?> sourceOperator = new StreamSource<>(new FiniteSource());
	StreamTaskTestHarness<String> testHarness = setupSourceStreamTask(sourceOperator, BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.invoke();
	testHarness.waitForTaskCompletion();

	assertEquals(1, testHarness.getOutput().size());
	assertEquals(Watermark.MAX_WATERMARK, testHarness.getOutput().peek());
}
 

private static <T> StreamTaskTestHarness<T> setupSourceStreamTask(
	StreamSource<T, ?> sourceOperator,
	TypeInformation<T> outputType,
	final boolean cancelImmediatelyAfterCreation) {

	final StreamTaskTestHarness<T> testHarness = new StreamTaskTestHarness<>(
		(env) -> {
			SourceStreamTask<T, ?, ?> sourceTask = new SourceStreamTask<>(env);
			if (cancelImmediatelyAfterCreation) {
				try {
					sourceTask.cancel();
				} catch (Exception e) {
					throw new RuntimeException(e);
				}
			}
			return sourceTask;
		},
		outputType);
	testHarness.setupOutputForSingletonOperatorChain();

	StreamConfig streamConfig = testHarness.getStreamConfig();
	streamConfig.setStreamOperator(sourceOperator);
	streamConfig.setOperatorID(new OperatorID());
	streamConfig.setTimeCharacteristic(TimeCharacteristic.EventTime);

	return testHarness;
}
 

private static <T> StreamTaskTestHarness<T> setupSourceStreamTask(
	StreamSource<T, ?> sourceOperator,
	TypeInformation<T> outputType) {

	return setupSourceStreamTask(sourceOperator, outputType, false);
}