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

private static OperatorSubtaskState repartitionInitState(
	final OperatorSubtaskState initState,
	final int numKeyGroups,
	final int oldParallelism,
	final int newParallelism,
	final int subtaskIndex
) {
	return AbstractStreamOperatorTestHarness.repartitionOperatorState(initState, numKeyGroups, oldParallelism, newParallelism, subtaskIndex);
}
 

/**
 * Test restoring from a non-empty state taken using a previous Flink version, when some partitions could be
 * found for topics.
 */
@Test
public void testRestore() throws Exception {
	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
			new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));

	testHarness.open();

	// assert that there are partitions and is identical to expected list
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());

	// on restore, subscribedPartitionsToStartOffsets should be identical to the restored state
	assertEquals(PARTITION_STATE, consumerFunction.getSubscribedPartitionsToStartOffsets());

	// assert that state is correctly restored from legacy checkpoint
	assertTrue(consumerFunction.getRestoredState() != null);
	assertEquals(PARTITION_STATE, consumerFunction.getRestoredState());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

/**
 * Test restoring from savepoints before version Flink 1.3 should fail if discovery is enabled.
 */
@Test
public void testRestoreFailsWithNonEmptyPreFlink13StatesIfDiscoveryEnabled() throws Exception {
	assumeTrue(testMigrateVersion == MigrationVersion.v1_3 || testMigrateVersion == MigrationVersion.v1_2);

	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, 1000L); // discovery enabled

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
		new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file; should fail since discovery is enabled
	try {
		testHarness.initializeState(
			OperatorSnapshotUtil.getResourceFilename(
				"kafka-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));

		fail("Restore from savepoints from version before Flink 1.3.x should have failed if discovery is enabled.");
	} catch (Exception e) {
		Assert.assertTrue(e instanceof IllegalArgumentException);
	}
}
 

private static <T> AbstractStreamOperatorTestHarness<T> createTestHarness(
	SourceFunction<T> source, int numSubtasks, int subtaskIndex) throws Exception {

	AbstractStreamOperatorTestHarness<T> testHarness =
		new AbstractStreamOperatorTestHarness<>(
			new StreamSource<>(source), maxParallelism, numSubtasks, subtaskIndex);

	testHarness.setTimeCharacteristic(TimeCharacteristic.EventTime);

	return testHarness;
}
 

@Test
public void testOpen() throws Exception {
	MockSerializationSchema<Object> serializationSchema = new MockSerializationSchema<>();

	Properties config = TestUtils.getStandardProperties();
	FlinkKinesisProducer<Object> producer = new FlinkKinesisProducer<>(
		serializationSchema,
		config);
	AbstractStreamOperatorTestHarness<Object> testHarness = new AbstractStreamOperatorTestHarness<>(
		new StreamSink<>(producer), 1, 1, 0
	);

	testHarness.open();
	assertThat("Open method was not called", serializationSchema.isOpenCalled(), is(true));
}
 

private static void testUDF(TestUserFunction userFunction) throws Exception {
	OperatorSubtaskState snapshot;
	try (AbstractStreamOperatorTestHarness<Integer> testHarness = createTestHarness(userFunction)) {
		testHarness.open();
		snapshot = testHarness.snapshot(0L, 0L);
		assertFalse(userFunction.isRestored());
	}
	try (AbstractStreamOperatorTestHarness<Integer> testHarness = createTestHarness(userFunction)) {
		testHarness.initializeState(snapshot);
		testHarness.open();
		assertTrue(userFunction.isRestored());
	}
}
 

@Test
public void testOpen() throws Exception {
	MockDeserializationSchema<String> deserializationSchema = new MockDeserializationSchema<>();

	RMQSource<String> consumer = new RMQTestSource(deserializationSchema);
	AbstractStreamOperatorTestHarness<String> testHarness = new AbstractStreamOperatorTestHarness<>(
		new StreamSource<>(consumer), 1, 1, 0
	);

	testHarness.open();
	assertThat("Open method was not called", deserializationSchema.isOpenCalled(), is(true));
}
 

/**
 * Test restoring from a non-empty state taken using a previous Flink version, when some partitions could be
 * found for topics.
 */
@Test
public void testRestore() throws Exception {
	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
			new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));

	testHarness.open();

	// assert that there are partitions and is identical to expected list
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());

	// on restore, subscribedPartitionsToStartOffsets should be identical to the restored state
	assertEquals(PARTITION_STATE, consumerFunction.getSubscribedPartitionsToStartOffsets());

	// assert that state is correctly restored from legacy checkpoint
	assertTrue(consumerFunction.getRestoredState() != null);
	assertEquals(PARTITION_STATE, consumerFunction.getRestoredState());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

private static OperatorSubtaskState repartitionInitState(
	final OperatorSubtaskState initState,
	final int numKeyGroups,
	final int oldParallelism,
	final int newParallelism,
	final int subtaskIndex
) {
	return AbstractStreamOperatorTestHarness.repartitionOperatorState(initState, numKeyGroups, oldParallelism, newParallelism, subtaskIndex);
}
 

private static <T> AbstractStreamOperatorTestHarness<T> createTestHarness(
	SourceFunction<T> source, int numSubtasks, int subtaskIndex) throws Exception {

	AbstractStreamOperatorTestHarness<T> testHarness =
		new AbstractStreamOperatorTestHarness<>(
			new StreamSource<>(source), maxParallelism, numSubtasks, subtaskIndex);

	testHarness.setTimeCharacteristic(TimeCharacteristic.EventTime);

	return testHarness;
}
 

/**
 * Test restoring from savepoints before version Flink 1.3 should fail if discovery is enabled.
 */
@Test
public void testRestoreFailsWithNonEmptyPreFlink13StatesIfDiscoveryEnabled() throws Exception {
	assumeTrue(testMigrateVersion == MigrationVersion.v1_3 || testMigrateVersion == MigrationVersion.v1_2);

	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, 1000L); // discovery enabled

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
		new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file; should fail since discovery is enabled
	try {
		testHarness.initializeState(
			OperatorSnapshotUtil.getResourceFilename(
				"kafka-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));

		fail("Restore from savepoints from version before Flink 1.3.x should have failed if discovery is enabled.");
	} catch (Exception e) {
		Assert.assertTrue(e instanceof IllegalArgumentException);
	}
}
 

/**
 * Test restoring from an legacy empty state, when no partitions could be found for topics.
 */
@Test
public void testRestoreFromEmptyStateNoPartitions() throws Exception {
	final DummyFlinkKafkaConsumer<String> consumerFunction =
			new DummyFlinkKafkaConsumer<>(
				Collections.singletonList("dummy-topic"),
				Collections.<KafkaTopicPartition>emptyList(),
				FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator = new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-empty-state-snapshot"));

	testHarness.open();

	// assert that no partitions were found and is empty
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());

	// assert that no state was restored
	assertTrue(consumerFunction.getRestoredState().isEmpty());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

private static void testUDF(TestUserFunction userFunction) throws Exception {
	OperatorSubtaskState snapshot;
	try (AbstractStreamOperatorTestHarness<Integer> testHarness = createTestHarness(userFunction)) {
		testHarness.open();
		snapshot = testHarness.snapshot(0L, 0L);
		assertFalse(userFunction.isRestored());
	}
	try (AbstractStreamOperatorTestHarness<Integer> testHarness = createTestHarness(userFunction)) {
		testHarness.initializeState(snapshot);
		testHarness.open();
		assertTrue(userFunction.isRestored());
	}
}
 

/**
 * Test restoring from an legacy empty state, when no partitions could be found for topics.
 */
@Test
public void testRestoreFromEmptyStateNoPartitions() throws Exception {
	final DummyFlinkKafkaConsumer<String> consumerFunction =
			new DummyFlinkKafkaConsumer<>(
				Collections.singletonList("dummy-topic"),
				Collections.<KafkaTopicPartition>emptyList(),
				FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator = new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-empty-state-snapshot"));

	testHarness.open();

	// assert that no partitions were found and is empty
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());

	// assert that no state was restored
	assertTrue(consumerFunction.getRestoredState().isEmpty());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

private static <T> AbstractStreamOperatorTestHarness<T> createTestHarness(
        SourceFunction<T> source, int numSubtasks, int subtaskIndex) throws Exception {

    AbstractStreamOperatorTestHarness<T> testHarness =
            new AbstractStreamOperatorTestHarness<>(
                    new StreamSource<>(source), maxParallelism, numSubtasks, subtaskIndex);

    testHarness.setTimeCharacteristic(TimeCharacteristic.EventTime);

    return testHarness;
}
 

@Test
public void testOpen() throws Exception {
	MockSerializationSchema<String> serializationSchema = new MockSerializationSchema<>();

	RMQSink<String> producer = new RMQSink<>(rmqConnectionConfig, serializationSchema, publishOptions);
	AbstractStreamOperatorTestHarness<Object> testHarness = new AbstractStreamOperatorTestHarness<>(
		new StreamSink<>(producer), 1, 1, 0
	);

	testHarness.open();
	assertThat("Open method was not called", serializationSchema.isOpenCalled(), is(true));
}
 

private static OperatorSubtaskState repartitionInitState(
		final OperatorSubtaskState initState,
		final int numKeyGroups,
		final int oldParallelism,
		final int newParallelism,
		final int subtaskIndex
) {
	return AbstractStreamOperatorTestHarness.repartitionOperatorState(initState, numKeyGroups, oldParallelism, newParallelism, subtaskIndex);
}
 

private static AbstractStreamOperatorTestHarness<Integer> createTestHarness(TestUserFunction userFunction) throws Exception {
	return new AbstractStreamOperatorTestHarness<>(
		new StreamMap<>(userFunction),
		1,
		1,
		0);
}
 

private static OperatorSubtaskState repartitionInitState(
	final OperatorSubtaskState initState,
	final int numKeyGroups,
	final int oldParallelism,
	final int newParallelism,
	final int subtaskIndex
) {
	return AbstractStreamOperatorTestHarness.repartitionOperatorState(initState, numKeyGroups, oldParallelism, newParallelism, subtaskIndex);
}
 

private static OperatorSubtaskState repartitionInitState(
		final OperatorSubtaskState initState,
		final int numKeyGroups,
		final int oldParallelism,
		final int newParallelism,
		final int subtaskIndex
) {
	return AbstractStreamOperatorTestHarness.repartitionOperatorState(initState, numKeyGroups, oldParallelism, newParallelism, subtaskIndex);
}
 

private static OperatorSubtaskState repartitionInitState(
		final OperatorSubtaskState initState,
		final int numKeyGroups,
		final int oldParallelism,
		final int newParallelism,
		final int subtaskIndex
) {
	return AbstractStreamOperatorTestHarness.repartitionOperatorState(initState, numKeyGroups, oldParallelism, newParallelism, subtaskIndex);
}
 

private static void testUDF(TestUserFunction userFunction) throws Exception {
	OperatorSubtaskState snapshot;
	try (AbstractStreamOperatorTestHarness<Integer> testHarness = createTestHarness(userFunction)) {
		testHarness.open();
		snapshot = testHarness.snapshot(0L, 0L);
		assertFalse(userFunction.isRestored());
	}
	try (AbstractStreamOperatorTestHarness<Integer> testHarness = createTestHarness(userFunction)) {
		testHarness.initializeState(snapshot);
		testHarness.open();
		assertTrue(userFunction.isRestored());
	}
}
 

/**
 * Test restoring from a non-empty state taken using a previous Flink version, when some partitions could be
 * found for topics.
 */
@Test
public void testRestore() throws Exception {
	final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

	final DummyFlinkKafkaConsumer<String> consumerFunction =
		new DummyFlinkKafkaConsumer<>(TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

	StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
			new StreamSource<>(consumerFunction);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();

	// restore state from binary snapshot file
	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"kafka-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));

	testHarness.open();

	// assert that there are partitions and is identical to expected list
	assertTrue(consumerFunction.getSubscribedPartitionsToStartOffsets() != null);
	assertTrue(!consumerFunction.getSubscribedPartitionsToStartOffsets().isEmpty());

	// on restore, subscribedPartitionsToStartOffsets should be identical to the restored state
	assertEquals(PARTITION_STATE, consumerFunction.getSubscribedPartitionsToStartOffsets());

	// assert that state is correctly restored from legacy checkpoint
	assertTrue(consumerFunction.getRestoredState() != null);
	assertEquals(PARTITION_STATE, consumerFunction.getRestoredState());

	consumerOperator.close();
	consumerOperator.cancel();
}
 

@Test
public void testCheckpointAndRestore() {
	try {
		final int numElements = 10000;

		List<Integer> data = new ArrayList<Integer>(numElements);
		List<Integer> result = new ArrayList<Integer>(numElements);

		for (int i = 0; i < numElements; i++) {
			data.add(i);
		}

		final FromElementsFunction<Integer> source = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> src = new StreamSource<>(source);
		AbstractStreamOperatorTestHarness<Integer> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
		testHarness.open();

		final SourceFunction.SourceContext<Integer> ctx = new ListSourceContext<Integer>(result, 2L);

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					source.run(ctx);
				}
				catch (Throwable t) {
					error[0] = t;
				}
			}
		};
		runner.start();

		// wait for a bit
		Thread.sleep(1000);

		// make a checkpoint
		List<Integer> checkpointData = new ArrayList<>(numElements);
		OperatorSubtaskState handles = null;
		synchronized (ctx.getCheckpointLock()) {
			handles = testHarness.snapshot(566, System.currentTimeMillis());
			checkpointData.addAll(result);
		}

		// cancel the source
		source.cancel();
		runner.join();

		// check for errors
		if (error[0] != null) {
			System.err.println("Error in asynchronous source runner");
			error[0].printStackTrace();
			fail("Error in asynchronous source runner");
		}

		final FromElementsFunction<Integer> sourceCopy = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<Integer> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
		testHarnessCopy.setup();
		testHarnessCopy.initializeState(handles);
		testHarnessCopy.open();

		// recovery run
		SourceFunction.SourceContext<Integer> newCtx = new ListSourceContext<>(checkpointData);

		sourceCopy.run(newCtx);

		assertEquals(data, checkpointData);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testCheckpointAndRestore() {
	try {
		final int numElements = 10000;

		List<Integer> data = new ArrayList<Integer>(numElements);
		List<Integer> result = new ArrayList<Integer>(numElements);

		for (int i = 0; i < numElements; i++) {
			data.add(i);
		}

		final FromElementsFunction<Integer> source = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> src = new StreamSource<>(source);
		AbstractStreamOperatorTestHarness<Integer> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
		testHarness.open();

		final SourceFunction.SourceContext<Integer> ctx = new ListSourceContext<Integer>(result, 2L);

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					source.run(ctx);
				}
				catch (Throwable t) {
					error[0] = t;
				}
			}
		};
		runner.start();

		// wait for a bit
		Thread.sleep(1000);

		// make a checkpoint
		List<Integer> checkpointData = new ArrayList<>(numElements);
		OperatorSubtaskState handles = null;
		synchronized (ctx.getCheckpointLock()) {
			handles = testHarness.snapshot(566, System.currentTimeMillis());
			checkpointData.addAll(result);
		}

		// cancel the source
		source.cancel();
		runner.join();

		// check for errors
		if (error[0] != null) {
			System.err.println("Error in asynchronous source runner");
			error[0].printStackTrace();
			fail("Error in asynchronous source runner");
		}

		final FromElementsFunction<Integer> sourceCopy = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<Integer> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
		testHarnessCopy.setup();
		testHarnessCopy.initializeState(handles);
		testHarnessCopy.open();

		// recovery run
		SourceFunction.SourceContext<Integer> newCtx = new ListSourceContext<>(checkpointData);

		sourceCopy.run(newCtx);

		assertEquals(data, checkpointData);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testReaderScalingDown() throws Exception {
	// simulates the scenario of scaling down from 2 to 1 instances

	final OneShotLatch waitingLatch = new OneShotLatch();

	// create the first instance and let it process the first split till element 5
	final OneShotLatch triggerLatch1 = new OneShotLatch();
	BlockingFileInputFormat format1 = new BlockingFileInputFormat(
		triggerLatch1, waitingLatch, new Path("test"), 20, 5);
	FileInputSplit[] splits = format1.createInputSplits(2);

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, String> testHarness1 = getTestHarness(format1, 2, 0);
	testHarness1.open();
	testHarness1.processElement(new StreamRecord<>(getTimestampedSplit(0, splits[0])));

	// wait until its arrives to element 5
	if (!triggerLatch1.isTriggered()) {
		triggerLatch1.await();
	}

	// create the second instance and let it process the second split till element 15
	final OneShotLatch triggerLatch2 = new OneShotLatch();
	BlockingFileInputFormat format2 = new BlockingFileInputFormat(
		triggerLatch2, waitingLatch, new Path("test"), 20, 15);

	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, String> testHarness2 = getTestHarness(format2, 2, 1);
	testHarness2.open();
	testHarness2.processElement(new StreamRecord<>(getTimestampedSplit(0, splits[1])));

	// wait until its arrives to element 15
	if (!triggerLatch2.isTriggered()) {
		triggerLatch2.await();
	}

	// 1) clear the outputs of the two previous instances so that
	// we can compare their newly produced outputs with the merged one
	testHarness1.getOutput().clear();
	testHarness2.getOutput().clear();

	// 2) take the snapshots from the previous instances and merge them
	// into a new one which will be then used to initialize a third instance
	OperatorSubtaskState mergedState = AbstractStreamOperatorTestHarness.
		repackageState(
			testHarness2.snapshot(0, 0),
			testHarness1.snapshot(0, 0)
		);

	// 3) and repartition to get the initialized state when scaling down.
	OperatorSubtaskState initState =
		AbstractStreamOperatorTestHarness.repartitionOperatorState(mergedState, maxParallelism, 2, 1, 0);

	// create the third instance
	final OneShotLatch wLatch = new OneShotLatch();
	final OneShotLatch tLatch = new OneShotLatch();

	BlockingFileInputFormat format = new BlockingFileInputFormat(wLatch, tLatch, new Path("test"), 20, 5);
	OneInputStreamOperatorTestHarness<TimestampedFileInputSplit, String> testHarness = getTestHarness(format, 1, 0);

	// initialize the state of the new operator with the constructed by
	// combining the partial states of the instances above.
	testHarness.initializeState(initState);
	testHarness.open();

	// now restart the waiting operators
	wLatch.trigger();
	tLatch.trigger();
	waitingLatch.trigger();

	// and wait for the processing to finish
	synchronized (testHarness1.getCheckpointLock()) {
		testHarness1.close();
	}
	synchronized (testHarness2.getCheckpointLock()) {
		testHarness2.close();
	}
	synchronized (testHarness.getCheckpointLock()) {
		testHarness.close();
	}

	Queue<Object> expectedResult = new ArrayDeque<>();
	putElementsInQ(expectedResult, testHarness1.getOutput());
	putElementsInQ(expectedResult, testHarness2.getOutput());

	Queue<Object> actualResult = new ArrayDeque<>();
	putElementsInQ(actualResult, testHarness.getOutput());

	Assert.assertEquals(20, actualResult.size());
	Assert.assertArrayEquals(expectedResult.toArray(), actualResult.toArray());
}
 

@Test
public void testCheckpointing() throws Exception {
	source.autoAck = false;

	StreamSource<String, RMQSource<String>> src = new StreamSource<>(source);
	AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	sourceThread.start();

	Thread.sleep(5);

	final Random random = new Random(System.currentTimeMillis());
	int numSnapshots = 50;
	long previousSnapshotId;
	long lastSnapshotId = 0;

	long totalNumberOfAcks = 0;

	for (int i = 0; i < numSnapshots; i++) {
		long snapshotId = random.nextLong();
		OperatorSubtaskState data;

		synchronized (DummySourceContext.lock) {
			data = testHarness.snapshot(snapshotId, System.currentTimeMillis());
			previousSnapshotId = lastSnapshotId;
			lastSnapshotId = messageId;
		}
		// let some time pass
		Thread.sleep(5);

		// check if the correct number of messages have been snapshotted
		final long numIds = lastSnapshotId - previousSnapshotId;

		RMQTestSource sourceCopy = new RMQTestSource();
		StreamSource<String, RMQTestSource> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<String> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);

		testHarnessCopy.setup();
		testHarnessCopy.initializeState(data);
		testHarnessCopy.open();

		ArrayDeque<Tuple2<Long, Set<String>>> deque = sourceCopy.getRestoredState();
		Set<String> messageIds = deque.getLast().f1;

		assertEquals(numIds, messageIds.size());
		if (messageIds.size() > 0) {
			assertTrue(messageIds.contains(Long.toString(lastSnapshotId)));
		}

		// check if the messages are being acknowledged and the transaction committed
		synchronized (DummySourceContext.lock) {
			source.notifyCheckpointComplete(snapshotId);
		}
		totalNumberOfAcks += numIds;

	}

	Mockito.verify(source.channel, Mockito.times((int) totalNumberOfAcks)).basicAck(Mockito.anyLong(), Mockito.eq(false));
	Mockito.verify(source.channel, Mockito.times(numSnapshots)).txCommit();

}
 

@Test
public void testFunctionRestore() throws Exception {
	String testBasePath = hdfsURI + "/" + UUID.randomUUID() + "/";

	org.apache.hadoop.fs.Path path = null;
	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<org.apache.hadoop.fs.Path, String> file = createFileAndFillWithData(testBasePath, "file", i, "This is test line.");
		path = file.f0;
		fileModTime = hdfs.getFileStatus(file.f0).getModificationTime();
	}

	TextInputFormat format = new TextInputFormat(new Path(testBasePath));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	final DummySourceContext sourceContext = new DummySourceContext() {
		@Override
		public void collect(TimestampedFileInputSplit element) {
			latch.trigger();
		}
	};

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(sourceContext);
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	// first condition for the source to have updated its state: emit at least one element
	if (!latch.isTriggered()) {
		latch.await();
	}

	// second condition for the source to have updated its state: it's not on the lock anymore,
	// this means it has processed all the splits and updated its state.
	synchronized (sourceContext.getCheckpointLock()) {}

	OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
	monitoringFunction.cancel();
	runner.join();

	testHarness.close();

	final ContinuousFileMonitoringFunction<String> monitoringFunctionCopy =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> srcCopy =
		new StreamSource<>(monitoringFunctionCopy);

	AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarnessCopy =
		new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
	testHarnessCopy.initializeState(snapshot);
	testHarnessCopy.open();

	Assert.assertNull(error[0]);
	Assert.assertEquals(fileModTime, monitoringFunctionCopy.getGlobalModificationTime());

	hdfs.delete(path, false);
}
 

@Test
public void testScalingDown() throws Exception {
	final File outDir = tempFolder.newFolder();

	OneInputStreamOperatorTestHarness<String, Object> testHarness1 = createRescalingTestSink(outDir, 3, 0, 100);
	testHarness1.setup();
	testHarness1.open();

	OneInputStreamOperatorTestHarness<String, Object> testHarness2 = createRescalingTestSink(outDir, 3, 1, 100);
	testHarness2.setup();
	testHarness2.open();

	OneInputStreamOperatorTestHarness<String, Object> testHarness3 = createRescalingTestSink(outDir, 3, 2, 100);
	testHarness3.setup();
	testHarness3.open();

	testHarness1.processElement(new StreamRecord<>("test1", 0L));
	checkLocalFs(outDir, 1, 0, 0, 0);

	testHarness2.processElement(new StreamRecord<>("test2", 0L));
	checkLocalFs(outDir, 2, 0, 0, 0);

	testHarness3.processElement(new StreamRecord<>("test3", 0L));
	testHarness3.processElement(new StreamRecord<>("test4", 0L));
	checkLocalFs(outDir, 4, 0, 0, 0);

	// intentionally we snapshot them in the reverse order so that the states are shuffled
	OperatorSubtaskState mergedSnapshot = AbstractStreamOperatorTestHarness.repackageState(
		testHarness3.snapshot(0, 0),
		testHarness1.snapshot(0, 0),
		testHarness2.snapshot(0, 0)
	);

	OperatorSubtaskState initState1 = AbstractStreamOperatorTestHarness.repartitionOperatorState(
		mergedSnapshot, maxParallelism, 3, 2, 0);

	testHarness1 = createRescalingTestSink(outDir, 2, 0, 100);
	testHarness1.setup();
	testHarness1.initializeState(initState1);
	testHarness1.open();

	checkLocalFs(outDir, 1, 0, 3, 3);

	OperatorSubtaskState initState2 = AbstractStreamOperatorTestHarness.repartitionOperatorState(
		mergedSnapshot, maxParallelism, 3, 2, 1);

	testHarness2 = createRescalingTestSink(outDir, 2, 1, 100);
	testHarness2.setup();
	testHarness2.initializeState(initState2);
	testHarness2.open();

	checkLocalFs(outDir, 0, 0, 4, 4);
}
 

@Test
public void testStateAndTimerStateShufflingScalingDown() throws Exception {
	final int maxParallelism = 10;

	// first get two keys that will fall into different key-group ranges that go
	// to different operator subtasks when we restore

	// get two sub key-ranges so that we can restore two ranges separately
	KeyGroupRange subKeyGroupRange1 = new KeyGroupRange(0, (maxParallelism / 2) - 1);
	KeyGroupRange subKeyGroupRange2 = new KeyGroupRange(subKeyGroupRange1.getEndKeyGroup() + 1, maxParallelism - 1);

	// get two different keys, one per sub range
	int key1 = getKeyInKeyGroupRange(subKeyGroupRange1, maxParallelism);
	int key2 = getKeyInKeyGroupRange(subKeyGroupRange2, maxParallelism);

	OperatorSubtaskState snapshot1, snapshot2;
	// register some state with both instances and scale down to parallelism 1
	try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness1 =
			createTestHarness(maxParallelism, 2, 0)) {

		testHarness1.setup();
		testHarness1.open();

		testHarness1.processWatermark(0L);
		testHarness1.setProcessingTime(0L);

		testHarness1.processElement(new Tuple2<>(key1, "SET_EVENT_TIME_TIMER:30"), 0);
		testHarness1.processElement(new Tuple2<>(key1, "SET_PROC_TIME_TIMER:30"), 0);
		testHarness1.processElement(new Tuple2<>(key1, "SET_STATE:HELLO"), 0);

		snapshot1 = testHarness1.snapshot(0, 0);
	}

	try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness2 =
			createTestHarness(maxParallelism, 2, 1)) {
		testHarness2.setup();
		testHarness2.open();

		testHarness2.processWatermark(0L);
		testHarness2.setProcessingTime(0L);

		testHarness2.processElement(new Tuple2<>(key2, "SET_EVENT_TIME_TIMER:40"), 0);
		testHarness2.processElement(new Tuple2<>(key2, "SET_PROC_TIME_TIMER:40"), 0);
		testHarness2.processElement(new Tuple2<>(key2, "SET_STATE:CIAO"), 0);

		snapshot2 = testHarness2.snapshot(0, 0);
	}
	// take a snapshot from each one of the "parallel" instances of the operator
	// and combine them into one so that we can scale down

	OperatorSubtaskState repackagedState =
		AbstractStreamOperatorTestHarness.repackageState(snapshot1, snapshot2);

	OperatorSubtaskState initSubTaskState =
		AbstractStreamOperatorTestHarness.repartitionOperatorState(repackagedState, maxParallelism, 2, 1, 0);

	try (KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, String>, String> testHarness3 =
			createTestHarness(maxParallelism, 1, 0)) {
		testHarness3.setup();
		testHarness3.initializeState(initSubTaskState);
		testHarness3.open();

		testHarness3.processWatermark(30L);
		assertThat(extractResult(testHarness3), contains("ON_EVENT_TIME:HELLO"));
		assertTrue(extractResult(testHarness3).isEmpty());

		testHarness3.processWatermark(40L);
		assertThat(extractResult(testHarness3), contains("ON_EVENT_TIME:CIAO"));
		assertTrue(extractResult(testHarness3).isEmpty());

		testHarness3.setProcessingTime(30L);
		assertThat(extractResult(testHarness3), contains("ON_PROC_TIME:HELLO"));
		assertTrue(extractResult(testHarness3).isEmpty());

		testHarness3.setProcessingTime(40L);
		assertThat(extractResult(testHarness3), contains("ON_PROC_TIME:CIAO"));
		assertTrue(extractResult(testHarness3).isEmpty());
	}
}