/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.streaming.connectors.kafka;

import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.serialization.TypeInformationSerializationSchema;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.client.program.ProgramInvocationException;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.memory.DataInputView;
import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.client.JobCancellationException;
import org.apache.flink.runtime.client.JobExecutionException;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobStatus;
import org.apache.flink.runtime.state.CheckpointListener;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.checkpoint.ListCheckpointed;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.graph.StreamingJobGraphGenerator;
import org.apache.flink.streaming.connectors.kafka.config.StartupMode;
import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
import org.apache.flink.streaming.connectors.kafka.testutils.DataGenerators;
import org.apache.flink.streaming.connectors.kafka.testutils.FailingIdentityMapper;
import org.apache.flink.streaming.connectors.kafka.testutils.PartitionValidatingMapper;
import org.apache.flink.streaming.connectors.kafka.testutils.ThrottledMapper;
import org.apache.flink.streaming.connectors.kafka.testutils.Tuple2FlinkPartitioner;
import org.apache.flink.streaming.connectors.kafka.testutils.ValidatingExactlyOnceSink;
import org.apache.flink.streaming.util.serialization.KeyedDeserializationSchema;
import org.apache.flink.streaming.util.serialization.KeyedDeserializationSchemaWrapper;
import org.apache.flink.streaming.util.serialization.KeyedSerializationSchema;
import org.apache.flink.streaming.util.serialization.KeyedSerializationSchemaWrapper;
import org.apache.flink.streaming.util.serialization.TypeInformationKeyValueSerializationSchema;
import org.apache.flink.test.util.SuccessException;
import org.apache.flink.testutils.junit.RetryOnException;
import org.apache.flink.testutils.junit.RetryRule;
import org.apache.flink.util.Collector;
import org.apache.flink.util.ExceptionUtils;

import org.apache.flink.shaded.guava18.com.google.common.collect.Iterables;

import kafka.consumer.Consumer;
import kafka.consumer.ConsumerConfig;
import kafka.consumer.ConsumerIterator;
import kafka.consumer.KafkaStream;
import kafka.javaapi.consumer.ConsumerConnector;
import kafka.message.MessageAndMetadata;
import kafka.server.KafkaServer;
import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.errors.TimeoutException;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;

import javax.management.MBeanServer;
import javax.management.ObjectName;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumerBase.KEY_PARTITION_DISCOVERY_INTERVAL_MILLIS;
import static org.apache.flink.streaming.connectors.kafka.testutils.ClusterCommunicationUtils.getRunningJobs;
import static org.apache.flink.streaming.connectors.kafka.testutils.ClusterCommunicationUtils.waitUntilJobIsRunning;
import static org.apache.flink.streaming.connectors.kafka.testutils.ClusterCommunicationUtils.waitUntilNoJobIsRunning;
import static org.apache.flink.test.util.TestUtils.tryExecute;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

/**
 * Abstract test base for all Kafka consumer tests.
 */
@SuppressWarnings("serial")
public abstract class KafkaConsumerTestBase extends KafkaTestBase {

	@Rule
	public RetryRule retryRule = new RetryRule();

	private ClusterClient<?> client;

	// ------------------------------------------------------------------------
	//  Common Test Preparation
	// ------------------------------------------------------------------------

	/**
	 * Makes sure that no job is on the JobManager any more from any previous tests that use
	 * the same mini cluster. Otherwise, missing slots may happen.
	 */
	@Before
	public void setClientAndEnsureNoJobIsLingering() throws Exception {
		client = flink.getClusterClient();
		waitUntilNoJobIsRunning(client);
	}

	// ------------------------------------------------------------------------
	//  Suite of Tests
	//
	//  The tests here are all not activated (by an @Test tag), but need
	//  to be invoked from the extending classes. That way, the classes can
	//  select which tests to run.
	// ------------------------------------------------------------------------

	/**
	 * Test that ensures the KafkaConsumer is properly failing if the topic doesnt exist
	 * and a wrong broker was specified.
	 *
	 * @throws Exception
	 */
	public void runFailOnNoBrokerTest() throws Exception {
		try {
			Properties properties = new Properties();

			StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
			see.getConfig().disableSysoutLogging();
			see.setRestartStrategy(RestartStrategies.noRestart());
			see.setParallelism(1);

			// use wrong ports for the consumers
			properties.setProperty("bootstrap.servers", "localhost:80");
			properties.setProperty("zookeeper.connect", "localhost:80");
			properties.setProperty("group.id", "test");
			properties.setProperty("request.timeout.ms", "3000"); // let the test fail fast
			properties.setProperty("socket.timeout.ms", "3000");
			properties.setProperty("session.timeout.ms", "2000");
			properties.setProperty("fetch.max.wait.ms", "2000");
			properties.setProperty("heartbeat.interval.ms", "1000");
			properties.putAll(secureProps);
			FlinkKafkaConsumerBase<String> source = kafkaServer.getConsumer("doesntexist", new SimpleStringSchema(), properties);
			DataStream<String> stream = see.addSource(source);
			stream.print();
			see.execute("No broker test");
		} catch (JobExecutionException jee) {
			if (kafkaServer.getVersion().equals("0.9") || kafkaServer.getVersion().equals("0.10") || kafkaServer.getVersion().equals("0.11")) {
				assertTrue(jee.getCause() instanceof TimeoutException);

				TimeoutException te = (TimeoutException) jee.getCause();

				assertEquals("Timeout expired while fetching topic metadata", te.getMessage());
			} else {
				assertTrue(jee.getCause() instanceof RuntimeException);

				RuntimeException re = (RuntimeException) jee.getCause();

				assertTrue(re.getMessage().contains("Unable to retrieve any partitions"));
			}
		}
	}

	/**
	 * Ensures that the committed offsets to Kafka are the offsets of "the next record to process".
	 */
	public void runCommitOffsetsToKafka() throws Exception {
		// 3 partitions with 50 records each (0-49, so the expected commit offset of each partition should be 50)
		final int parallelism = 3;
		final int recordsInEachPartition = 50;

		final String topicName = writeSequence("testCommitOffsetsToKafkaTopic", recordsInEachPartition, parallelism, 1);

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		env.enableCheckpointing(200, CheckpointingMode.EXACTLY_ONCE);

		Properties props = new Properties();
		props.putAll(standardProps);
		// We need to turn on the partition dynamic discovery so that the source task won't exit after reaching the
		// log end.
		props.setProperty(KEY_PARTITION_DISCOVERY_INTERVAL_MILLIS, "10000");
		DataStream<String> stream = env.addSource(kafkaServer.getConsumer(topicName, new SimpleStringSchema(), props));
		stream.addSink(new DiscardingSink<String>());

		final AtomicReference<Throwable> errorRef = new AtomicReference<>();
		final Thread runner = new Thread("runner") {
			@Override
			public void run() {
				try {
					env.execute();
				}
				catch (Throwable t) {
					if (!(t instanceof JobCancellationException)) {
						errorRef.set(t);
					}
				}
			}
		};
		runner.start();

		final Long l50 = 50L; // the final committed offset in Kafka should be 50
		final long deadline = 30_000_000_000L + System.nanoTime();

		KafkaTestEnvironment.KafkaOffsetHandler kafkaOffsetHandler = kafkaServer.createOffsetHandler();

		do {
			Long o1 = kafkaOffsetHandler.getCommittedOffset(topicName, 0);
			Long o2 = kafkaOffsetHandler.getCommittedOffset(topicName, 1);
			Long o3 = kafkaOffsetHandler.getCommittedOffset(topicName, 2);

			if (l50.equals(o1) && l50.equals(o2) && l50.equals(o3)) {
				break;
			}

			Thread.sleep(100);
		}
		while (System.nanoTime() < deadline);

		// cancel the job & wait for the job to finish
		client.cancel(Iterables.getOnlyElement(getRunningJobs(client)));
		runner.join();

		final Throwable t = errorRef.get();
		if (t != null) {
			throw new RuntimeException("Job failed with an exception", t);
		}

		// final check to see if offsets are correctly in Kafka
		Long o1 = kafkaOffsetHandler.getCommittedOffset(topicName, 0);
		Long o2 = kafkaOffsetHandler.getCommittedOffset(topicName, 1);
		Long o3 = kafkaOffsetHandler.getCommittedOffset(topicName, 2);
		Assert.assertEquals(Long.valueOf(50L), o1);
		Assert.assertEquals(Long.valueOf(50L), o2);
		Assert.assertEquals(Long.valueOf(50L), o3);

		kafkaOffsetHandler.close();
		deleteTestTopic(topicName);
	}

	/**
	 * This test ensures that when the consumers retrieve some start offset from kafka (earliest, latest), that this offset
	 * is committed to Kafka, even if some partitions are not read.
	 *
	 * <p>Test:
	 * - Create 3 partitions
	 * - write 50 messages into each.
	 * - Start three consumers with auto.offset.reset='latest' and wait until they committed into Kafka.
	 * - Check if the offsets in Kafka are set to 50 for the three partitions
	 *
	 * <p>See FLINK-3440 as well
	 */
	public void runAutoOffsetRetrievalAndCommitToKafka() throws Exception {
		// 3 partitions with 50 records each (0-49, so the expected commit offset of each partition should be 50)
		final int parallelism = 3;
		final int recordsInEachPartition = 50;

		final String topicName = writeSequence("testAutoOffsetRetrievalAndCommitToKafkaTopic", recordsInEachPartition, parallelism, 1);

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		env.setParallelism(parallelism);
		env.enableCheckpointing(200);

		Properties readProps = new Properties();
		readProps.putAll(standardProps);
		readProps.setProperty("auto.offset.reset", "latest"); // set to reset to latest, so that partitions are initially not read

		DataStream<String> stream = env.addSource(kafkaServer.getConsumer(topicName, new SimpleStringSchema(), readProps));
		stream.addSink(new DiscardingSink<String>());

		final AtomicReference<Throwable> errorRef = new AtomicReference<>();
		final Thread runner = new Thread("runner") {
			@Override
			public void run() {
				try {
					env.execute();
				}
				catch (Throwable t) {
					if (!(t instanceof JobCancellationException)) {
						errorRef.set(t);
					}
				}
			}
		};
		runner.start();

		KafkaTestEnvironment.KafkaOffsetHandler kafkaOffsetHandler = kafkaServer.createOffsetHandler();

		final Long l50 = 50L; // the final committed offset in Kafka should be 50
		final long deadline = 30_000_000_000L + System.nanoTime();
		do {
			Long o1 = kafkaOffsetHandler.getCommittedOffset(topicName, 0);
			Long o2 = kafkaOffsetHandler.getCommittedOffset(topicName, 1);
			Long o3 = kafkaOffsetHandler.getCommittedOffset(topicName, 2);

			if (l50.equals(o1) && l50.equals(o2) && l50.equals(o3)) {
				break;
			}

			Thread.sleep(100);
		}
		while (System.nanoTime() < deadline);

		// cancel the job & wait for the job to finish
		client.cancel(Iterables.getOnlyElement(getRunningJobs(client)));
		runner.join();

		final Throwable t = errorRef.get();
		if (t != null) {
			throw new RuntimeException("Job failed with an exception", t);
		}

		// final check to see if offsets are correctly in Kafka
		Long o1 = kafkaOffsetHandler.getCommittedOffset(topicName, 0);
		Long o2 = kafkaOffsetHandler.getCommittedOffset(topicName, 1);
		Long o3 = kafkaOffsetHandler.getCommittedOffset(topicName, 2);
		Assert.assertEquals(Long.valueOf(50L), o1);
		Assert.assertEquals(Long.valueOf(50L), o2);
		Assert.assertEquals(Long.valueOf(50L), o3);

		kafkaOffsetHandler.close();
		deleteTestTopic(topicName);
	}

	/**
	 * This test ensures that when explicitly set to start from earliest record, the consumer
	 * ignores the "auto.offset.reset" behaviour as well as any committed group offsets in Kafka.
	 */
	public void runStartFromEarliestOffsets() throws Exception {
		// 3 partitions with 50 records each (0-49, so the expected commit offset of each partition should be 50)
		final int parallelism = 3;
		final int recordsInEachPartition = 50;

		final String topicName = writeSequence("testStartFromEarliestOffsetsTopic", recordsInEachPartition, parallelism, 1);

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.setParallelism(parallelism);

		Properties readProps = new Properties();
		readProps.putAll(standardProps);
		readProps.setProperty("auto.offset.reset", "latest"); // this should be ignored

		// the committed offsets should be ignored
		KafkaTestEnvironment.KafkaOffsetHandler kafkaOffsetHandler = kafkaServer.createOffsetHandler();
		kafkaOffsetHandler.setCommittedOffset(topicName, 0, 23);
		kafkaOffsetHandler.setCommittedOffset(topicName, 1, 31);
		kafkaOffsetHandler.setCommittedOffset(topicName, 2, 43);

		readSequence(env, StartupMode.EARLIEST, null, null, readProps, parallelism, topicName, recordsInEachPartition, 0);

		kafkaOffsetHandler.close();
		deleteTestTopic(topicName);
	}

	/**
	 * This test ensures that when explicitly set to start from latest record, the consumer
	 * ignores the "auto.offset.reset" behaviour as well as any committed group offsets in Kafka.
	 */
	public void runStartFromLatestOffsets() throws Exception {
		// 50 records written to each of 3 partitions before launching a latest-starting consuming job
		final int parallelism = 3;
		final int recordsInEachPartition = 50;

		// each partition will be written an extra 200 records
		final int extraRecordsInEachPartition = 200;

		// all already existing data in the topic, before the consuming topology has started, should be ignored
		final String topicName = writeSequence("testStartFromLatestOffsetsTopic", recordsInEachPartition, parallelism, 1);

		// the committed offsets should be ignored
		KafkaTestEnvironment.KafkaOffsetHandler kafkaOffsetHandler = kafkaServer.createOffsetHandler();
		kafkaOffsetHandler.setCommittedOffset(topicName, 0, 23);
		kafkaOffsetHandler.setCommittedOffset(topicName, 1, 31);
		kafkaOffsetHandler.setCommittedOffset(topicName, 2, 43);

		// job names for the topologies for writing and consuming the extra records
		final String consumeExtraRecordsJobName = "Consume Extra Records Job";
		final String writeExtraRecordsJobName = "Write Extra Records Job";

		// serialization / deserialization schemas for writing and consuming the extra records
		final TypeInformation<Tuple2<Integer, Integer>> resultType =
			TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {});

		final KeyedSerializationSchema<Tuple2<Integer, Integer>> serSchema =
			new KeyedSerializationSchemaWrapper<>(
				new TypeInformationSerializationSchema<>(resultType, new ExecutionConfig()));

		final KeyedDeserializationSchema<Tuple2<Integer, Integer>> deserSchema =
			new KeyedDeserializationSchemaWrapper<>(
				new TypeInformationSerializationSchema<>(resultType, new ExecutionConfig()));

		// setup and run the latest-consuming job
		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.setParallelism(parallelism);

		final Properties readProps = new Properties();
		readProps.putAll(standardProps);
		readProps.setProperty("auto.offset.reset", "earliest"); // this should be ignored

		FlinkKafkaConsumerBase<Tuple2<Integer, Integer>> latestReadingConsumer =
			kafkaServer.getConsumer(topicName, deserSchema, readProps);
		latestReadingConsumer.setStartFromLatest();

		env
			.addSource(latestReadingConsumer).setParallelism(parallelism)
			.flatMap(new FlatMapFunction<Tuple2<Integer, Integer>, Object>() {
				@Override
				public void flatMap(Tuple2<Integer, Integer> value, Collector<Object> out) throws Exception {
					if (value.f1 - recordsInEachPartition < 0) {
						throw new RuntimeException("test failed; consumed a record that was previously written: " + value);
					}
				}
			}).setParallelism(1)
			.addSink(new DiscardingSink<>());

		JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());
		final JobID consumeJobId = jobGraph.getJobID();

		final AtomicReference<Throwable> error = new AtomicReference<>();
		Thread consumeThread = new Thread(new Runnable() {
			@Override
			public void run() {
				try {
					client.setDetached(false);
					client.submitJob(jobGraph, KafkaConsumerTestBase.class.getClassLoader(), false);
				} catch (Throwable t) {
					if (!ExceptionUtils.findThrowable(t, JobCancellationException.class).isPresent()) {
						error.set(t);
					}
				}
			}
		});
		consumeThread.start();

		// wait until the consuming job has started, to be extra safe
		waitUntilJobIsRunning(client);

		// setup the extra records writing job
		final StreamExecutionEnvironment env2 = StreamExecutionEnvironment.getExecutionEnvironment();

		env2.setParallelism(parallelism);

		DataStream<Tuple2<Integer, Integer>> extraRecordsStream = env2
			.addSource(new RichParallelSourceFunction<Tuple2<Integer, Integer>>() {

				private boolean running = true;

				@Override
				public void run(SourceContext<Tuple2<Integer, Integer>> ctx) throws Exception {
					int count = recordsInEachPartition; // the extra records should start from the last written value
					int partition = getRuntimeContext().getIndexOfThisSubtask();

					while (running && count < recordsInEachPartition + extraRecordsInEachPartition) {
						ctx.collect(new Tuple2<>(partition, count));
						count++;
					}
				}

				@Override
				public void cancel() {
					running = false;
				}
			});

		kafkaServer.produceIntoKafka(extraRecordsStream, topicName, serSchema, readProps, null);

		try {
			env2.execute(writeExtraRecordsJobName);
		}
		catch (Exception e) {
			throw new RuntimeException("Writing extra records failed", e);
		}

		// cancel the consume job after all extra records are written
		client.cancel(consumeJobId);
		consumeThread.join();

		kafkaOffsetHandler.close();
		deleteTestTopic(topicName);

		// check whether the consuming thread threw any test errors;
		// test will fail here if the consume job had incorrectly read any records other than the extra records
		final Throwable consumerError = error.get();
		if (consumerError != null) {
			throw new Exception("Exception in the consuming thread", consumerError);
		}
	}

	/**
	 * This test ensures that the consumer correctly uses group offsets in Kafka, and defaults to "auto.offset.reset"
	 * behaviour when necessary, when explicitly configured to start from group offsets.
	 *
	 * <p>The partitions and their committed group offsets are setup as:
	 * 	partition 0 --> committed offset 23
	 * 	partition 1 --> no commit offset
	 * 	partition 2 --> committed offset 43
	 *
	 * <p>When configured to start from group offsets, each partition should read:
	 * 	partition 0 --> start from offset 23, read to offset 49 (27 records)
	 * 	partition 1 --> default to "auto.offset.reset" (set to earliest), so start from offset 0, read to offset 49 (50 records)
	 * 	partition 2 --> start from offset 43, read to offset 49 (7 records)
	 */
	public void runStartFromGroupOffsets() throws Exception {
		// 3 partitions with 50 records each (offsets 0-49)
		final int parallelism = 3;
		final int recordsInEachPartition = 50;

		final String topicName = writeSequence("testStartFromGroupOffsetsTopic", recordsInEachPartition, parallelism, 1);

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.setParallelism(parallelism);

		Properties readProps = new Properties();
		readProps.putAll(standardProps);
		readProps.setProperty("auto.offset.reset", "earliest");

		// the committed group offsets should be used as starting points
		KafkaTestEnvironment.KafkaOffsetHandler kafkaOffsetHandler = kafkaServer.createOffsetHandler();

		// only partitions 0 and 2 have group offsets committed
		kafkaOffsetHandler.setCommittedOffset(topicName, 0, 23);
		kafkaOffsetHandler.setCommittedOffset(topicName, 2, 43);

		Map<Integer, Tuple2<Integer, Integer>> partitionsToValueCountAndStartOffsets = new HashMap<>();
		partitionsToValueCountAndStartOffsets.put(0, new Tuple2<>(27, 23)); // partition 0 should read offset 23-49
		partitionsToValueCountAndStartOffsets.put(1, new Tuple2<>(50, 0)); // partition 1 should read offset 0-49
		partitionsToValueCountAndStartOffsets.put(2, new Tuple2<>(7, 43));	// partition 2 should read offset 43-49

		readSequence(env, StartupMode.GROUP_OFFSETS, null, null, readProps, topicName, partitionsToValueCountAndStartOffsets);

		kafkaOffsetHandler.close();
		deleteTestTopic(topicName);
	}

	/**
	 * This test ensures that the consumer correctly uses user-supplied specific offsets when explicitly configured to
	 * start from specific offsets. For partitions which a specific offset can not be found for, the starting position
	 * for them should fallback to the group offsets behaviour.
	 *
	 * <p>4 partitions will have 50 records with offsets 0 to 49. The supplied specific offsets map is:
	 * 	partition 0 --> start from offset 19
	 * 	partition 1 --> not set
	 * 	partition 2 --> start from offset 22
	 * 	partition 3 --> not set
	 * 	partition 4 --> start from offset 26 (this should be ignored because the partition does not exist)
	 *
	 * <p>The partitions and their committed group offsets are setup as:
	 * 	partition 0 --> committed offset 23
	 * 	partition 1 --> committed offset 31
	 * 	partition 2 --> committed offset 43
	 * 	partition 3 --> no commit offset
	 *
	 * <p>When configured to start from these specific offsets, each partition should read:
	 * 	partition 0 --> start from offset 19, read to offset 49 (31 records)
	 * 	partition 1 --> fallback to group offsets, so start from offset 31, read to offset 49 (19 records)
	 * 	partition 2 --> start from offset 22, read to offset 49 (28 records)
	 * 	partition 3 --> fallback to group offsets, but since there is no group offset for this partition,
	 * 	                will default to "auto.offset.reset" (set to "earliest"),
	 * 	                so start from offset 0, read to offset 49 (50 records)
	 */
	public void runStartFromSpecificOffsets() throws Exception {
		// 4 partitions with 50 records each (offsets 0-49)
		final int parallelism = 4;
		final int recordsInEachPartition = 50;

		final String topicName = writeSequence("testStartFromSpecificOffsetsTopic", recordsInEachPartition, parallelism, 1);

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.setParallelism(parallelism);

		Properties readProps = new Properties();
		readProps.putAll(standardProps);
		readProps.setProperty("auto.offset.reset", "earliest"); // partition 3 should default back to this behaviour

		Map<KafkaTopicPartition, Long> specificStartupOffsets = new HashMap<>();
		specificStartupOffsets.put(new KafkaTopicPartition(topicName, 0), 19L);
		specificStartupOffsets.put(new KafkaTopicPartition(topicName, 2), 22L);
		specificStartupOffsets.put(new KafkaTopicPartition(topicName, 4), 26L); // non-existing partition, should be ignored

		// only the committed offset for partition 1 should be used, because partition 1 has no entry in specific offset map
		KafkaTestEnvironment.KafkaOffsetHandler kafkaOffsetHandler = kafkaServer.createOffsetHandler();
		kafkaOffsetHandler.setCommittedOffset(topicName, 0, 23);
		kafkaOffsetHandler.setCommittedOffset(topicName, 1, 31);
		kafkaOffsetHandler.setCommittedOffset(topicName, 2, 43);

		Map<Integer, Tuple2<Integer, Integer>> partitionsToValueCountAndStartOffsets = new HashMap<>();
		partitionsToValueCountAndStartOffsets.put(0, new Tuple2<>(31, 19)); // partition 0 should read offset 19-49
		partitionsToValueCountAndStartOffsets.put(1, new Tuple2<>(19, 31)); // partition 1 should read offset 31-49
		partitionsToValueCountAndStartOffsets.put(2, new Tuple2<>(28, 22));	// partition 2 should read offset 22-49
		partitionsToValueCountAndStartOffsets.put(3, new Tuple2<>(50, 0));	// partition 3 should read offset 0-49

		readSequence(env, StartupMode.SPECIFIC_OFFSETS, specificStartupOffsets, null, readProps, topicName, partitionsToValueCountAndStartOffsets);

		kafkaOffsetHandler.close();
		deleteTestTopic(topicName);
	}

	/**
	 * This test ensures that the consumer correctly uses user-supplied timestamp when explicitly configured to
	 * start from timestamp.
	 *
	 * <p>The validated Kafka data is written in 2 steps: first, an initial 50 records is written to each partition.
	 * After that, another 30 records is appended to each partition. Before each step, a timestamp is recorded.
	 * For the validation, when the read job is configured to start from the first timestamp, each partition should start
	 * from offset 0 and read a total of 80 records. When configured to start from the second timestamp,
	 * each partition should start from offset 50 and read on the remaining 30 appended records.
	 */
	public void runStartFromTimestamp() throws Exception {
		// 4 partitions with 50 records each
		final int parallelism = 4;
		final int initialRecordsInEachPartition = 50;
		final int appendRecordsInEachPartition = 30;

		// attempt to create an appended test sequence, where the timestamp of writing the appended sequence
		// is assured to be larger than the timestamp of the original sequence.
		long firstTimestamp = System.currentTimeMillis();
		String topic = writeSequence("runStartFromTimestamp", initialRecordsInEachPartition, parallelism, 1);

		long secondTimestamp = 0;
		while (secondTimestamp <= firstTimestamp) {
			Thread.sleep(1000);
			secondTimestamp = System.currentTimeMillis();
		}
		writeAppendSequence(topic, initialRecordsInEachPartition, appendRecordsInEachPartition, parallelism);

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();
		env.setParallelism(parallelism);

		Properties readProps = new Properties();
		readProps.putAll(standardProps);

		readSequence(env, StartupMode.TIMESTAMP, null, firstTimestamp,
			readProps, parallelism, topic, initialRecordsInEachPartition + appendRecordsInEachPartition, 0);
		readSequence(env, StartupMode.TIMESTAMP, null, secondTimestamp,
			readProps, parallelism, topic, appendRecordsInEachPartition, initialRecordsInEachPartition);

		deleteTestTopic(topic);
	}

	/**
	 * Ensure Kafka is working on both producer and consumer side.
	 * This executes a job that contains two Flink pipelines.
	 *
	 * <pre>
	 * (generator source) --> (kafka sink)-[KAFKA-TOPIC]-(kafka source) --> (validating sink)
	 * </pre>
	 *
	 * <p>We need to externally retry this test. We cannot let Flink's retry mechanism do it, because the Kafka producer
	 * does not guarantee exactly-once output. Hence a recovery would introduce duplicates that
	 * cause the test to fail.
	 *
	 * <p>This test also ensures that FLINK-3156 doesn't happen again:
	 *
	 * <p>The following situation caused a NPE in the FlinkKafkaConsumer
	 *
	 * <p>topic-1 <-- elements are only produced into topic1.
	 * topic-2
	 *
	 * <p>Therefore, this test is consuming as well from an empty topic.
	 *
	 */
	@RetryOnException(times = 2, exception = kafka.common.NotLeaderForPartitionException.class)
	public void runSimpleConcurrentProducerConsumerTopology() throws Exception {
		final String topic = "concurrentProducerConsumerTopic_" + UUID.randomUUID().toString();
		final String additionalEmptyTopic = "additionalEmptyTopic_" + UUID.randomUUID().toString();

		final int parallelism = 3;
		final int elementsPerPartition = 100;
		final int totalElements = parallelism * elementsPerPartition;

		createTestTopic(topic, parallelism, 2);
		createTestTopic(additionalEmptyTopic, parallelism, 1); // create an empty topic which will remain empty all the time

		final StreamExecutionEnvironment env =
				StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(parallelism);
//		env.enableCheckpointing(500);
		env.setRestartStrategy(RestartStrategies.noRestart()); // fail immediately
		env.getConfig().disableSysoutLogging();

		TypeInformation<Tuple2<Long, String>> longStringType = TypeInformation.of(new TypeHint<Tuple2<Long, String>>(){});

		TypeInformationSerializationSchema<Tuple2<Long, String>> sourceSchema =
				new TypeInformationSerializationSchema<>(longStringType, env.getConfig());

		TypeInformationSerializationSchema<Tuple2<Long, String>> sinkSchema =
				new TypeInformationSerializationSchema<>(longStringType, env.getConfig());

		// ----------- add producer dataflow ----------

		DataStream<Tuple2<Long, String>> stream = env.addSource(new RichParallelSourceFunction<Tuple2<Long, String>>() {

			private boolean running = true;

			@Override
			public void run(SourceContext<Tuple2<Long, String>> ctx) throws InterruptedException {
				int cnt = getRuntimeContext().getIndexOfThisSubtask() * elementsPerPartition;
				int limit = cnt + elementsPerPartition;

				while (running && cnt < limit) {
					ctx.collect(new Tuple2<>(1000L + cnt, "kafka-" + cnt));
					cnt++;
					// we delay data generation a bit so that we are sure that some checkpoints are
					// triggered (for FLINK-3156)
					Thread.sleep(50);
				}
			}

			@Override
			public void cancel() {
				running = false;
			}
		});
		Properties producerProperties = FlinkKafkaProducerBase.getPropertiesFromBrokerList(brokerConnectionStrings);
		producerProperties.setProperty("retries", "3");
		producerProperties.putAll(secureProps);
		kafkaServer.produceIntoKafka(stream, topic, new KeyedSerializationSchemaWrapper<>(sinkSchema), producerProperties, null);

		// ----------- add consumer dataflow ----------

		List<String> topics = new ArrayList<>();
		topics.add(topic);
		topics.add(additionalEmptyTopic);

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<Tuple2<Long, String>> source = kafkaServer.getConsumer(topics, sourceSchema, props);

		DataStreamSource<Tuple2<Long, String>> consuming = env.addSource(source).setParallelism(parallelism);

		consuming.addSink(new RichSinkFunction<Tuple2<Long, String>>() {

			private int elCnt = 0;
			private BitSet validator = new BitSet(totalElements);

			@Override
			public void invoke(Tuple2<Long, String> value) throws Exception {
				String[] sp = value.f1.split("-");
				int v = Integer.parseInt(sp[1]);

				assertEquals(value.f0 - 1000, (long) v);

				assertFalse("Received tuple twice", validator.get(v));
				validator.set(v);
				elCnt++;

				if (elCnt == totalElements) {
					// check if everything in the bitset is set to true
					int nc;
					if ((nc = validator.nextClearBit(0)) != totalElements) {
						fail("The bitset was not set to 1 on all elements. Next clear:"
								+ nc + " Set: " + validator);
					}
					throw new SuccessException();
				}
			}

			@Override
			public void close() throws Exception {
				super.close();
			}
		}).setParallelism(1);

		try {
			tryExecutePropagateExceptions(env, "runSimpleConcurrentProducerConsumerTopology");
		}
		catch (ProgramInvocationException | JobExecutionException e) {
			// look for NotLeaderForPartitionException
			Throwable cause = e.getCause();

			// search for nested SuccessExceptions
			int depth = 0;
			while (cause != null && depth++ < 20) {
				if (cause instanceof kafka.common.NotLeaderForPartitionException) {
					throw (Exception) cause;
				}
				cause = cause.getCause();
			}
			throw e;
		}

		deleteTestTopic(topic);
	}

	/**
	 * Tests the proper consumption when having a 1:1 correspondence between kafka partitions and
	 * Flink sources.
	 */
	public void runOneToOneExactlyOnceTest() throws Exception {

		final String topic = "oneToOneTopic";
		final int parallelism = 5;
		final int numElementsPerPartition = 1000;
		final int totalElements = parallelism * numElementsPerPartition;
		final int failAfterElements = numElementsPerPartition / 3;

		createTestTopic(topic, parallelism, 1);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic,
				parallelism,
				numElementsPerPartition,
				true);

		// run the topology that fails and recovers

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.enableCheckpointing(500);
		env.setParallelism(parallelism);
		env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);

		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);

		env
				.addSource(kafkaSource)
				.map(new PartitionValidatingMapper(parallelism, 1))
				.map(new FailingIdentityMapper<Integer>(failAfterElements))
				.addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

		FailingIdentityMapper.failedBefore = false;
		tryExecute(env, "One-to-one exactly once test");

		deleteTestTopic(topic);
	}

	/**
	 * Tests the proper consumption when having a 1:1 correspondence between kafka partitions and
	 * Flink sources.
	 */
	public void runOneToOneTest(boolean isFinite) throws Exception {

		final String topic = "oneToOneTopic";
		final int parallelism = 5;
		final int numElementsPerPartition = 1000;
		final int totalElements = parallelism * numElementsPerPartition;

		createTestTopic(topic, parallelism, 1);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic,
				parallelism,
				numElementsPerPartition,
				true);

		// run the topology that fails and recovers

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.enableCheckpointing(5000);
		env.setParallelism(parallelism);
		env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);

		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);

		kafkaSource.setStopAtLatest(isFinite);

		env
				.addSource(kafkaSource)
				.addSink(new ValidatingExactlyOnceSink(totalElements, false)).setParallelism(1);

		env.execute();

		deleteTestTopic(topic);
	}

	/**
	 * Tests the proper consumption when having fewer Flink sources than Kafka partitions, so
	 * one Flink source will read multiple Kafka partitions.
	 */
	public void runOneSourceMultiplePartitionsExactlyOnceTest() throws Exception {
		final String topic = "oneToManyTopic";
		final int numPartitions = 5;
		final int numElementsPerPartition = 1000;
		final int totalElements = numPartitions * numElementsPerPartition;
		final int failAfterElements = numElementsPerPartition / 3;

		final int parallelism = 2;

		createTestTopic(topic, numPartitions, 1);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic,
				numPartitions,
				numElementsPerPartition,
				true);

		// run the topology that fails and recovers

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.enableCheckpointing(500);
		env.setParallelism(parallelism);
		env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);

		env
				.addSource(kafkaSource)
				.map(new PartitionValidatingMapper(numPartitions, 3))
				.map(new FailingIdentityMapper<Integer>(failAfterElements))
				.addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

		FailingIdentityMapper.failedBefore = false;
		tryExecute(env, "One-source-multi-partitions exactly once test");

		deleteTestTopic(topic);
	}

	/**
	 * Tests the proper consumption when having fewer Flink sources than Kafka partitions, so
	 * one Flink source will read multiple Kafka partitions.
	 */
	public void runOneSourceMultiplePartitionsFiniteTest() throws Exception {
		final String topic = "oneToManyTopic";
		final int numPartitions = 5;
		final int numElementsPerPartition = 1000;
		final int totalElements = numPartitions * numElementsPerPartition;

		final int parallelism = 2;

		createTestTopic(topic, numPartitions, 1);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic,
				numPartitions,
				numElementsPerPartition,
				true);

		// run the topology that fails and recovers

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.enableCheckpointing(500);
		env.setParallelism(parallelism);
		env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);
		kafkaSource.setStopAtLatest(true);
		env
				.addSource(kafkaSource)
				.addSink(new ValidatingExactlyOnceSink(totalElements, false)).setParallelism(1);

		env.execute("One-source-multi-partitions exactly once test");

		deleteTestTopic(topic);
	}

	/**
	 * Tests the proper consumption when having more Flink sources than Kafka partitions, which means
	 * that some Flink sources will read no partitions.
	 */
	public void runMultipleSourcesOnePartitionExactlyOnceTest() throws Exception {
		final String topic = "manyToOneTopic";
		final int numPartitions = 5;
		final int numElementsPerPartition = 1000;
		final int totalElements = numPartitions * numElementsPerPartition;
		final int failAfterElements = numElementsPerPartition / 3;

		final int parallelism = 8;

		createTestTopic(topic, numPartitions, 1);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic,
				numPartitions,
				numElementsPerPartition,
				true);

		// run the topology that fails and recovers

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.enableCheckpointing(500);
		env.setParallelism(parallelism);
		// set the number of restarts to one. The failing mapper will fail once, then it's only success exceptions.
		env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
		env.getConfig().disableSysoutLogging();
		env.setBufferTimeout(0);

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);

		env
			.addSource(kafkaSource)
			.map(new PartitionValidatingMapper(numPartitions, 1))
			.map(new FailingIdentityMapper<Integer>(failAfterElements))
			.addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

		FailingIdentityMapper.failedBefore = false;
		tryExecute(env, "multi-source-one-partitions exactly once test");

		deleteTestTopic(topic);
	}

	/**
	 * Tests the proper consumption when having more Flink sources than Kafka partitions, which means
	 * that some Flink sources will read no partitions.
	 */
	public void runMultipleSourcesOnePartitionFiniteTest() throws Exception {
		final String topic = "manyToOneTopic";
		final int numPartitions = 5;
		final int numElementsPerPartition = 1000;
		final int totalElements = numPartitions * numElementsPerPartition;

		final int parallelism = 8;

		createTestTopic(topic, numPartitions, 1);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic,
				numPartitions,
				numElementsPerPartition,
				true);

		// run the topology that fails and recovers

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.enableCheckpointing(500);
		env.setParallelism(parallelism);
		// set the number of restarts to one. The failing mapper will fail once, then it's only success exceptions.
		env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
		env.getConfig().disableSysoutLogging();
		env.setBufferTimeout(0);

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);
		kafkaSource.setStopAtLatest(true);
		env
			.addSource(kafkaSource)
			.addSink(new ValidatingExactlyOnceSink(totalElements, false)).setParallelism(1);

		env.execute();

		deleteTestTopic(topic);
	}

	/**
	 * Tests that the source can be properly canceled when reading full partitions.
	 */
	public void runCancelingOnFullInputTest() throws Exception {
		final String topic = "cancelingOnFullTopic";

		final int parallelism = 3;
		createTestTopic(topic, parallelism, 1);

		// launch a producer thread
		DataGenerators.InfiniteStringsGenerator generator =
				new DataGenerators.InfiniteStringsGenerator(kafkaServer, topic);
		generator.start();

		// launch a consumer asynchronously

		final AtomicReference<Throwable> jobError = new AtomicReference<>();

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(parallelism);
		env.enableCheckpointing(100);
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<String> source = kafkaServer.getConsumer(topic, new SimpleStringSchema(), props);

		env.addSource(source).addSink(new DiscardingSink<String>());

		JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());
		final JobID jobId = jobGraph.getJobID();

		final Runnable jobRunner = new Runnable() {
			@Override
			public void run() {
				try {
					client.setDetached(false);
					client.submitJob(jobGraph, KafkaConsumerTestBase.class.getClassLoader(), false);
				}
				catch (Throwable t) {
					jobError.set(t);
				}
			}
		};

		Thread runnerThread = new Thread(jobRunner, "program runner thread");
		runnerThread.start();

		// wait a bit before canceling
		Thread.sleep(2000);

		Throwable failueCause = jobError.get();
		if (failueCause != null) {
			failueCause.printStackTrace();
			Assert.fail("Test failed prematurely with: " + failueCause.getMessage());
		}

		// cancel
		client.cancel(jobId);

		// wait for the program to be done and validate that we failed with the right exception
		runnerThread.join();

		assertEquals(JobStatus.CANCELED, client.getJobStatus(jobId).get());

		if (generator.isAlive()) {
			generator.shutdown();
			generator.join();
		}
		else {
			Throwable t = generator.getError();
			if (t != null) {
				t.printStackTrace();
				fail("Generator failed: " + t.getMessage());
			} else {
				fail("Generator failed with no exception");
			}
		}

		deleteTestTopic(topic);
	}

	/**
	 * Tests that the source can be properly canceled when reading empty partitions.
	 */
	public void runCancelingOnEmptyInputTest() throws Exception {
		final String topic = "cancelingOnEmptyInputTopic";

		final int parallelism = 3;
		createTestTopic(topic, parallelism, 1);

		final AtomicReference<Throwable> error = new AtomicReference<>();

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(parallelism);
		env.enableCheckpointing(100);
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<String> source = kafkaServer.getConsumer(topic, new SimpleStringSchema(), props);

		env.addSource(source).addSink(new DiscardingSink<String>());

		JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env.getStreamGraph());
		final JobID jobId = jobGraph.getJobID();

		final Runnable jobRunner = new Runnable() {
			@Override
			public void run() {
				try {
					client.setDetached(false);
					client.submitJob(jobGraph, KafkaConsumerTestBase.class.getClassLoader(), false);
				}
				catch (Throwable t) {
					LOG.error("Job Runner failed with exception", t);
					error.set(t);
				}
			}
		};

		Thread runnerThread = new Thread(jobRunner, "program runner thread");
		runnerThread.start();

		// wait a bit before canceling
		Thread.sleep(2000);

		Throwable failueCause = error.get();
		if (failueCause != null) {
			failueCause.printStackTrace();
			Assert.fail("Test failed prematurely with: " + failueCause.getMessage());
		}
		// cancel
		client.cancel(jobId);

		// wait for the program to be done and validate that we failed with the right exception
		runnerThread.join();

		assertEquals(JobStatus.CANCELED, client.getJobStatus(jobId).get());

		deleteTestTopic(topic);
	}

	/**
	 * Test producing and consuming into multiple topics.
	 * @throws Exception
	 */
	public void runProduceConsumeMultipleTopics() throws Exception {
		final int numTopics = 5;
		final int numElements = 20;

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();

		// create topics with content
		final List<String> topics = new ArrayList<>();
		for (int i = 0; i < numTopics; i++) {
			final String topic = "topic-" + i;
			topics.add(topic);
			// create topic
			createTestTopic(topic, i + 1 /*partitions*/, 1);
		}

		// before FLINK-6078 the RemoteExecutionEnvironment set the parallelism to 1 as well
		env.setParallelism(1);

		// run first job, producing into all topics
		DataStream<Tuple3<Integer, Integer, String>> stream = env.addSource(new RichParallelSourceFunction<Tuple3<Integer, Integer, String>>() {

			@Override
			public void run(SourceContext<Tuple3<Integer, Integer, String>> ctx) throws Exception {
				int partition = getRuntimeContext().getIndexOfThisSubtask();

				for (int topicId = 0; topicId < numTopics; topicId++) {
					for (int i = 0; i < numElements; i++) {
						ctx.collect(new Tuple3<>(partition, i, "topic-" + topicId));
					}
				}
			}

			@Override
			public void cancel() {
			}
		});

		Tuple2WithTopicSchema schema = new Tuple2WithTopicSchema(env.getConfig());

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		kafkaServer.produceIntoKafka(stream, "dummy", schema, props, null);

		env.execute("Write to topics");

		// run second job consuming from multiple topics
		env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.getConfig().disableSysoutLogging();

		stream = env.addSource(kafkaServer.getConsumer(topics, schema, props));

		stream.flatMap(new FlatMapFunction<Tuple3<Integer, Integer, String>, Integer>() {
			Map<String, Integer> countPerTopic = new HashMap<>(numTopics);
			@Override
			public void flatMap(Tuple3<Integer, Integer, String> value, Collector<Integer> out) throws Exception {
				Integer count = countPerTopic.get(value.f2);
				if (count == null) {
					count = 1;
				} else {
					count++;
				}
				countPerTopic.put(value.f2, count);

				// check map:
				for (Map.Entry<String, Integer> el: countPerTopic.entrySet()) {
					if (el.getValue() < numElements) {
						break; // not enough yet
					}
					if (el.getValue() > numElements) {
						throw new RuntimeException("There is a failure in the test. I've read " +
								el.getValue() + " from topic " + el.getKey());
					}
				}
				// we've seen messages from all topics
				throw new SuccessException();
			}
		}).setParallelism(1);

		tryExecute(env, "Count elements from the topics");

		// delete all topics again
		for (int i = 0; i < numTopics; i++) {
			final String topic = "topic-" + i;
			deleteTestTopic(topic);
		}
	}

	/**
	 * Test Flink's Kafka integration also with very big records (30MB).
	 *
	 * <p>see http://stackoverflow.com/questions/21020347/kafka-sending-a-15mb-message
	 *
	 */
	public void runBigRecordTestTopology() throws Exception {

		final String topic = "bigRecordTestTopic";
		final int parallelism = 1; // otherwise, the kafka mini clusters may run out of heap space

		createTestTopic(topic, parallelism, 1);

		final TypeInformation<Tuple2<Long, byte[]>> longBytesInfo =
				TypeInformation.of(new TypeHint<Tuple2<Long, byte[]>>(){});

		final TypeInformationSerializationSchema<Tuple2<Long, byte[]>> serSchema =
				new TypeInformationSerializationSchema<>(longBytesInfo, new ExecutionConfig());

		final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setRestartStrategy(RestartStrategies.noRestart());
		env.getConfig().disableSysoutLogging();
		env.enableCheckpointing(100);
		env.setParallelism(parallelism);

		// add consuming topology:
		Properties consumerProps = new Properties();
		consumerProps.putAll(standardProps);
		consumerProps.setProperty("fetch.message.max.bytes", Integer.toString(1024 * 1024 * 14));
		consumerProps.setProperty("max.partition.fetch.bytes", Integer.toString(1024 * 1024 * 14)); // for the new fetcher
		consumerProps.setProperty("queued.max.message.chunks", "1");
		consumerProps.putAll(secureProps);

		FlinkKafkaConsumerBase<Tuple2<Long, byte[]>> source = kafkaServer.getConsumer(topic, serSchema, consumerProps);
		DataStreamSource<Tuple2<Long, byte[]>> consuming = env.addSource(source);

		consuming.addSink(new SinkFunction<Tuple2<Long, byte[]>>() {

			private int elCnt = 0;

			@Override
			public void invoke(Tuple2<Long, byte[]> value) throws Exception {
				elCnt++;
				if (value.f0 == -1) {
					// we should have seen 11 elements now.
					if (elCnt == 11) {
						throw new SuccessException();
					} else {
						throw new RuntimeException("There have been " + elCnt + " elements");
					}
				}
				if (elCnt > 10) {
					throw new RuntimeException("More than 10 elements seen: " + elCnt);
				}
			}
		});

		// add producing topology
		Properties producerProps = new Properties();
		producerProps.setProperty("max.request.size", Integer.toString(1024 * 1024 * 15));
		producerProps.setProperty("retries", "3");
		producerProps.putAll(secureProps);
		producerProps.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerConnectionStrings);

		DataStream<Tuple2<Long, byte[]>> stream = env.addSource(new RichSourceFunction<Tuple2<Long, byte[]>>() {

			private boolean running;

			@Override
			public void open(Configuration parameters) throws Exception {
				super.open(parameters);
				running = true;
			}

			@Override
			public void run(SourceContext<Tuple2<Long, byte[]>> ctx) throws Exception {
				Random rnd = new Random();
				long cnt = 0;
				int sevenMb = 1024 * 1024 * 7;

				while (running) {
					byte[] wl = new byte[sevenMb + rnd.nextInt(sevenMb)];
					ctx.collect(new Tuple2<>(cnt++, wl));

					Thread.sleep(100);

					if (cnt == 10) {
						// signal end
						ctx.collect(new Tuple2<>(-1L, new byte[]{1}));
						break;
					}
				}
			}

			@Override
			public void cancel() {
				running = false;
			}
		});

		kafkaServer.produceIntoKafka(stream, topic, new KeyedSerializationSchemaWrapper<>(serSchema), producerProps, null);

		tryExecute(env, "big topology test");
		deleteTestTopic(topic);
	}

	public void runBrokerFailureTest() throws Exception {
		final String topic = "brokerFailureTestTopic";

		final int parallelism = 2;
		final int numElementsPerPartition = 1000;
		final int totalElements = parallelism * numElementsPerPartition;
		final int failAfterElements = numElementsPerPartition / 3;

		createTestTopic(topic, parallelism, 2);

		DataGenerators.generateRandomizedIntegerSequence(
				StreamExecutionEnvironment.getExecutionEnvironment(),
				kafkaServer,
				topic, parallelism, numElementsPerPartition, true);

		// find leader to shut down
		int leaderId = kafkaServer.getLeaderToShutDown(topic);

		LOG.info("Leader to shutdown {}", leaderId);

		// run the topology (the consumers must handle the failures)

		DeserializationSchema<Integer> schema =
				new TypeInformationSerializationSchema<>(BasicTypeInfo.INT_TYPE_INFO, new ExecutionConfig());

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(parallelism);
		env.enableCheckpointing(500);
		env.setRestartStrategy(RestartStrategies.noRestart());
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		FlinkKafkaConsumerBase<Integer> kafkaSource = kafkaServer.getConsumer(topic, schema, props);

		env
				.addSource(kafkaSource)
				.map(new PartitionValidatingMapper(parallelism, 1))
				.map(new BrokerKillingMapper<Integer>(leaderId, failAfterElements))
				.addSink(new ValidatingExactlyOnceSink(totalElements)).setParallelism(1);

		BrokerKillingMapper.killedLeaderBefore = false;
		tryExecute(env, "Broker failure once test");

		// start a new broker:
		kafkaServer.restartBroker(leaderId);
	}

	public void runKeyValueTest() throws Exception {
		final String topic = "keyvaluetest";
		createTestTopic(topic, 1, 1);
		final int elementCount = 5000;

		// ----------- Write some data into Kafka -------------------

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(1);
		env.setRestartStrategy(RestartStrategies.noRestart());
		env.getConfig().disableSysoutLogging();

		DataStream<Tuple2<Long, PojoValue>> kvStream = env.addSource(new SourceFunction<Tuple2<Long, PojoValue>>() {
			@Override
			public void run(SourceContext<Tuple2<Long, PojoValue>> ctx) throws Exception {
				Random rnd = new Random(1337);
				for (long i = 0; i < elementCount; i++) {
					PojoValue pojo = new PojoValue();
					pojo.when = new Date(rnd.nextLong());
					pojo.lon = rnd.nextLong();
					pojo.lat = i;
					// make every second key null to ensure proper "null" serialization
					Long key = (i % 2 == 0) ? null : i;
					ctx.collect(new Tuple2<>(key, pojo));
				}
			}

			@Override
			public void cancel() {
			}
		});

		KeyedSerializationSchema<Tuple2<Long, PojoValue>> schema = new TypeInformationKeyValueSerializationSchema<>(Long.class, PojoValue.class, env.getConfig());
		Properties producerProperties = FlinkKafkaProducerBase.getPropertiesFromBrokerList(brokerConnectionStrings);
		producerProperties.setProperty("retries", "3");
		kafkaServer.produceIntoKafka(kvStream, topic, schema, producerProperties, null);
		env.execute("Write KV to Kafka");

		// ----------- Read the data again -------------------

		env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(1);
		env.setRestartStrategy(RestartStrategies.noRestart());
		env.getConfig().disableSysoutLogging();

		KeyedDeserializationSchema<Tuple2<Long, PojoValue>> readSchema = new TypeInformationKeyValueSerializationSchema<>(Long.class, PojoValue.class, env.getConfig());

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		DataStream<Tuple2<Long, PojoValue>> fromKafka = env.addSource(kafkaServer.getConsumer(topic, readSchema, props));
		fromKafka.flatMap(new RichFlatMapFunction<Tuple2<Long, PojoValue>, Object>() {
			long counter = 0;
			@Override
			public void flatMap(Tuple2<Long, PojoValue> value, Collector<Object> out) throws Exception {
				// the elements should be in order.
				Assert.assertTrue("Wrong value " + value.f1.lat, value.f1.lat == counter);
				if (value.f1.lat % 2 == 0) {
					assertNull("key was not null", value.f0);
				} else {
					Assert.assertTrue("Wrong value " + value.f0, value.f0 == counter);
				}
				counter++;
				if (counter == elementCount) {
					// we got the right number of elements
					throw new SuccessException();
				}
			}
		});

		tryExecute(env, "Read KV from Kafka");

		deleteTestTopic(topic);
	}

	private static class PojoValue {
		public Date when;
		public long lon;
		public long lat;
		public PojoValue() {}
	}

	/**
	 * Test delete behavior and metrics for producer.
	 * @throws Exception
	 */
	public void runAllDeletesTest() throws Exception {
		final String topic = "alldeletestest";
		createTestTopic(topic, 1, 1);
		final int elementCount = 300;

		// ----------- Write some data into Kafka -------------------

		StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(1);
		env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		env.getConfig().disableSysoutLogging();

		DataStream<Tuple2<byte[], PojoValue>> kvStream = env.addSource(new SourceFunction<Tuple2<byte[], PojoValue>>() {
			@Override
			public void run(SourceContext<Tuple2<byte[], PojoValue>> ctx) throws Exception {
				Random rnd = new Random(1337);
				for (long i = 0; i < elementCount; i++) {
					final byte[] key = new byte[200];
					rnd.nextBytes(key);
					ctx.collect(new Tuple2<>(key, (PojoValue) null));
				}
			}

			@Override
			public void cancel() {
			}
		});

		TypeInformationKeyValueSerializationSchema<byte[], PojoValue> schema = new TypeInformationKeyValueSerializationSchema<>(byte[].class, PojoValue.class, env.getConfig());

		Properties producerProperties = FlinkKafkaProducerBase.getPropertiesFromBrokerList(brokerConnectionStrings);
		producerProperties.setProperty("retries", "3");
		producerProperties.putAll(secureProps);
		kafkaServer.produceIntoKafka(kvStream, topic, schema, producerProperties, null);

		env.execute("Write deletes to Kafka");

		// ----------- Read the data again -------------------

		env = StreamExecutionEnvironment.getExecutionEnvironment();
		env.setParallelism(1);
		env.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		env.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);
		DataStream<Tuple2<byte[], PojoValue>> fromKafka = env.addSource(kafkaServer.getConsumer(topic, schema, props));

		fromKafka.flatMap(new RichFlatMapFunction<Tuple2<byte[], PojoValue>, Object>() {
			long counter = 0;
			@Override
			public void flatMap(Tuple2<byte[], PojoValue> value, Collector<Object> out) throws Exception {
				// ensure that deleted messages are passed as nulls
				assertNull(value.f1);
				counter++;
				if (counter == elementCount) {
					// we got the right number of elements
					throw new SuccessException();
				}
			}
		});

		tryExecute(env, "Read deletes from Kafka");

		deleteTestTopic(topic);
	}

	/**
	 * Test that ensures that DeserializationSchema.isEndOfStream() is properly evaluated.
	 *
	 * @throws Exception
	 */
	public void runEndOfStreamTest() throws Exception {

		final int elementCount = 300;
		final String topic = writeSequence("testEndOfStream", elementCount, 1, 1);

		// read using custom schema
		final StreamExecutionEnvironment env1 = StreamExecutionEnvironment.getExecutionEnvironment();
		env1.setParallelism(1);
		env1.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		env1.getConfig().disableSysoutLogging();

		Properties props = new Properties();
		props.putAll(standardProps);
		props.putAll(secureProps);

		DataStream<Tuple2<Integer, Integer>> fromKafka = env1.addSource(kafkaServer.getConsumer(topic, new FixedNumberDeserializationSchema(elementCount), props));
		fromKafka.flatMap(new FlatMapFunction<Tuple2<Integer, Integer>, Void>() {
			@Override
			public void flatMap(Tuple2<Integer, Integer> value, Collector<Void> out) throws Exception {
				// noop ;)
			}
		});

		tryExecute(env1, "Consume " + elementCount + " elements from Kafka");

		deleteTestTopic(topic);
	}

	/**
	 * Test metrics reporting for consumer.
	 *
	 * @throws Exception
	 */
	public void runMetricsTest() throws Throwable {

		// create a stream with 5 topics
		final String topic = "metricsStream";
		createTestTopic(topic, 5, 1);

		final Tuple1<Throwable> error = new Tuple1<>(null);

		// start job writing & reading data.
		final StreamExecutionEnvironment env1 = StreamExecutionEnvironment.getExecutionEnvironment();
		env1.setParallelism(1);
		env1.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		env1.getConfig().disableSysoutLogging();
		env1.disableOperatorChaining(); // let the source read everything into the network buffers

		TypeInformationSerializationSchema<Tuple2<Integer, Integer>> schema = new TypeInformationSerializationSchema<>(
				TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>(){}), env1.getConfig());

		DataStream<Tuple2<Integer, Integer>> fromKafka = env1.addSource(kafkaServer.getConsumer(topic, schema, standardProps));
		fromKafka.flatMap(new FlatMapFunction<Tuple2<Integer, Integer>, Void>() {
			@Override
			public void flatMap(Tuple2<Integer, Integer> value, Collector<Void> out) throws Exception {// no op
			}
		});

		DataStream<Tuple2<Integer, Integer>> fromGen = env1.addSource(new RichSourceFunction<Tuple2<Integer, Integer>>() {
			boolean running = true;

			@Override
			public void run(SourceContext<Tuple2<Integer, Integer>> ctx) throws Exception {
				int i = 0;
				while (running) {
					ctx.collect(Tuple2.of(i++, getRuntimeContext().getIndexOfThisSubtask()));
					Thread.sleep(1);
				}
			}

			@Override
			public void cancel() {
				running = false;
			}
		});

		kafkaServer.produceIntoKafka(fromGen, topic, new KeyedSerializationSchemaWrapper<>(schema), standardProps, null);

		JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(env1.getStreamGraph());
		final JobID jobId = jobGraph.getJobID();

		Runnable job = new Runnable() {
			@Override
			public void run() {
				try {
					client.setDetached(false);
					client.submitJob(jobGraph, KafkaConsumerTestBase.class.getClassLoader(), false);
				} catch (Throwable t) {
					if (!ExceptionUtils.findThrowable(t, JobCancellationException.class).isPresent()) {
						LOG.warn("Got exception during execution", t);
						error.f0 = t;
					}
				}
			}
		};
		Thread jobThread = new Thread(job);
		jobThread.start();

		try {
			// connect to JMX
			MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
			// wait until we've found all 5 offset metrics
			Set<ObjectName> offsetMetrics = mBeanServer.queryNames(new ObjectName("*current-offsets*:*"), null);
			while (offsetMetrics.size() < 5) { // test will time out if metrics are not properly working
				if (error.f0 != null) {
					// fail test early
					throw error.f0;
				}
				offsetMetrics = mBeanServer.queryNames(new ObjectName("*current-offsets*:*"), null);
				Thread.sleep(50);
			}
			Assert.assertEquals(5, offsetMetrics.size());
			// we can't rely on the consumer to have touched all the partitions already
			// that's why we'll wait until all five partitions have a positive offset.
			// The test will fail if we never meet the condition
			while (true) {
				int numPosOffsets = 0;
				// check that offsets are correctly reported
				for (ObjectName object : offsetMetrics) {
					Object offset = mBeanServer.getAttribute(object, "Value");
					if ((long) offset >= 0) {
						numPosOffsets++;
					}
				}
				if (numPosOffsets == 5) {
					break;
				}
				// wait for the consumer to consume on all partitions
				Thread.sleep(50);
			}

			// check if producer metrics are also available.
			Set<ObjectName> producerMetrics = mBeanServer.queryNames(new ObjectName("*KafkaProducer*:*"), null);
			Assert.assertTrue("No producer metrics found", producerMetrics.size() > 30);

			LOG.info("Found all JMX metrics. Cancelling job.");
		} finally {
			// cancel
			client.cancel(jobId);
			// wait for the job to finish (it should due to the cancel command above)
			jobThread.join();
		}

		if (error.f0 != null) {
			throw error.f0;
		}

		deleteTestTopic(topic);
	}

	private static class FixedNumberDeserializationSchema implements DeserializationSchema<Tuple2<Integer, Integer>> {

		final int finalCount;
		int count = 0;

		TypeInformation<Tuple2<Integer, Integer>> ti = TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>(){});
		TypeSerializer<Tuple2<Integer, Integer>> ser = ti.createSerializer(new ExecutionConfig());

		public FixedNumberDeserializationSchema(int finalCount) {
			this.finalCount = finalCount;
		}

		@Override
		public Tuple2<Integer, Integer> deserialize(byte[] message) throws IOException {
			DataInputView in = new DataInputViewStreamWrapper(new ByteArrayInputStream(message));
			return ser.deserialize(in);
		}

		@Override
		public boolean isEndOfStream(Tuple2<Integer, Integer> nextElement) {
			return ++count >= finalCount;
		}

		@Override
		public TypeInformation<Tuple2<Integer, Integer>> getProducedType() {
			return ti;
		}
	}

	// ------------------------------------------------------------------------
	//  Reading writing test data sets
	// ------------------------------------------------------------------------

	/**
	 * Runs a job using the provided environment to read a sequence of records from a single Kafka topic.
	 * The method allows to individually specify the expected starting offset and total read value count of each partition.
	 * The job will be considered successful only if all partition read results match the start offset and value count criteria.
	 */
	protected void readSequence(final StreamExecutionEnvironment env,
								final StartupMode startupMode,
								final Map<KafkaTopicPartition, Long> specificStartupOffsets,
								final Long startupTimestamp,
								final Properties cc,
								final String topicName,
								final Map<Integer, Tuple2<Integer, Integer>> partitionsToValuesCountAndStartOffset) throws Exception {
		final int sourceParallelism = partitionsToValuesCountAndStartOffset.keySet().size();

		int finalCountTmp = 0;
		for (Map.Entry<Integer, Tuple2<Integer, Integer>> valuesCountAndStartOffset : partitionsToValuesCountAndStartOffset.entrySet()) {
			finalCountTmp += valuesCountAndStartOffset.getValue().f0;
		}
		final int finalCount = finalCountTmp;

		final TypeInformation<Tuple2<Integer, Integer>> intIntTupleType =
				TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>(){});

		final TypeInformationSerializationSchema<Tuple2<Integer, Integer>> deser =
			new TypeInformationSerializationSchema<>(intIntTupleType, env.getConfig());

		// create the consumer
		cc.putAll(secureProps);
		FlinkKafkaConsumerBase<Tuple2<Integer, Integer>> consumer = kafkaServer.getConsumer(topicName, deser, cc);
		setKafkaConsumerOffset(startupMode, consumer, specificStartupOffsets, startupTimestamp);

		DataStream<Tuple2<Integer, Integer>> source = env
			.addSource(consumer).setParallelism(sourceParallelism)
			.map(new ThrottledMapper<Tuple2<Integer, Integer>>(20)).setParallelism(sourceParallelism);

		// verify data
		source.flatMap(new RichFlatMapFunction<Tuple2<Integer, Integer>, Integer>() {

			private HashMap<Integer, BitSet> partitionsToValueCheck;
			private int count = 0;

			@Override
			public void open(Configuration parameters) throws Exception {
				partitionsToValueCheck = new HashMap<>();
				for (Integer partition : partitionsToValuesCountAndStartOffset.keySet()) {
					partitionsToValueCheck.put(partition, new BitSet());
				}
			}

			@Override
			public void flatMap(Tuple2<Integer, Integer> value, Collector<Integer> out) throws Exception {
				int partition = value.f0;
				int val = value.f1;

				BitSet bitSet = partitionsToValueCheck.get(partition);
				if (bitSet == null) {
					throw new RuntimeException("Got a record from an unknown partition");
				} else {
					bitSet.set(val - partitionsToValuesCountAndStartOffset.get(partition).f1);
				}

				count++;

				LOG.info("Received message {}, total {} messages", value, count);

				// verify if we've seen everything
				if (count == finalCount) {
					for (Map.Entry<Integer, BitSet> partitionsToValueCheck : this.partitionsToValueCheck.entrySet()) {
						BitSet check = partitionsToValueCheck.getValue();
						int expectedValueCount = partitionsToValuesCountAndStartOffset.get(partitionsToValueCheck.getKey()).f0;

						if (check.cardinality() != expectedValueCount) {
							throw new RuntimeException("Expected cardinality to be " + expectedValueCount +
								", but was " + check.cardinality());
						} else if (check.nextClearBit(0) != expectedValueCount) {
							throw new RuntimeException("Expected next clear bit to be " + expectedValueCount +
								", but was " + check.cardinality());
						}
					}

					// test has passed
					throw new SuccessException();
				}
			}

		}).setParallelism(1);

		tryExecute(env, "Read data from Kafka");

		LOG.info("Successfully read sequence for verification");
	}

	/**
	 * Variant of {@link KafkaConsumerTestBase#readSequence(StreamExecutionEnvironment, StartupMode, Map, Long, Properties, String, Map)} to
	 * expect reading from the same start offset and the same value count for all partitions of a single Kafka topic.
	 */
	protected void readSequence(final StreamExecutionEnvironment env,
								final StartupMode startupMode,
								final Map<KafkaTopicPartition, Long> specificStartupOffsets,
								final Long startupTimestamp,
								final Properties cc,
								final int sourceParallelism,
								final String topicName,
								final int valuesCount,
								final int startFrom) throws Exception {
		HashMap<Integer, Tuple2<Integer, Integer>> partitionsToValuesCountAndStartOffset = new HashMap<>();
		for (int i = 0; i < sourceParallelism; i++) {
			partitionsToValuesCountAndStartOffset.put(i, new Tuple2<>(valuesCount, startFrom));
		}
		readSequence(env, startupMode, specificStartupOffsets, startupTimestamp, cc, topicName, partitionsToValuesCountAndStartOffset);
	}

	protected void setKafkaConsumerOffset(final StartupMode startupMode,
										final FlinkKafkaConsumerBase<Tuple2<Integer, Integer>> consumer,
										final Map<KafkaTopicPartition, Long> specificStartupOffsets,
										final Long startupTimestamp) {
		switch (startupMode) {
			case EARLIEST:
				consumer.setStartFromEarliest();
				break;
			case LATEST:
				consumer.setStartFromLatest();
				break;
			case SPECIFIC_OFFSETS:
				consumer.setStartFromSpecificOffsets(specificStartupOffsets);
				break;
			case GROUP_OFFSETS:
				consumer.setStartFromGroupOffsets();
				break;
			case TIMESTAMP:
				consumer.setStartFromTimestamp(startupTimestamp);
				break;
		}
	}

	protected String writeSequence(
			String baseTopicName,
			final int numElements,
			final int parallelism,
			final int replicationFactor) throws Exception {
		LOG.info("\n===================================\n" +
				"== Writing sequence of " + numElements + " into " + baseTopicName + " with p=" + parallelism + "\n" +
				"===================================");

		final TypeInformation<Tuple2<Integer, Integer>> resultType =
				TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {});

		final KeyedSerializationSchema<Tuple2<Integer, Integer>> serSchema =
				new KeyedSerializationSchemaWrapper<>(
						new TypeInformationSerializationSchema<>(resultType, new ExecutionConfig()));

		final KeyedDeserializationSchema<Tuple2<Integer, Integer>> deserSchema =
				new KeyedDeserializationSchemaWrapper<>(
						new TypeInformationSerializationSchema<>(resultType, new ExecutionConfig()));

		final int maxNumAttempts = 10;

		for (int attempt = 1; attempt <= maxNumAttempts; attempt++) {

			final String topicName = baseTopicName + '-' + attempt;

			LOG.info("Writing attempt #" + attempt);

			// -------- Write the Sequence --------

			createTestTopic(topicName, parallelism, replicationFactor);

			StreamExecutionEnvironment writeEnv = StreamExecutionEnvironment.getExecutionEnvironment();
			writeEnv.getConfig().setRestartStrategy(RestartStrategies.noRestart());
			writeEnv.getConfig().disableSysoutLogging();

			DataStream<Tuple2<Integer, Integer>> stream = writeEnv.addSource(new RichParallelSourceFunction<Tuple2<Integer, Integer>>() {

				private boolean running = true;

				@Override
				public void run(SourceContext<Tuple2<Integer, Integer>> ctx) throws Exception {
					int cnt = 0;
					int partition = getRuntimeContext().getIndexOfThisSubtask();

					while (running && cnt < numElements) {
						ctx.collect(new Tuple2<>(partition, cnt));
						cnt++;
					}
				}

				@Override
				public void cancel() {
					running = false;
				}
			}).setParallelism(parallelism);

			// the producer must not produce duplicates
			Properties producerProperties = FlinkKafkaProducerBase.getPropertiesFromBrokerList(brokerConnectionStrings);
			producerProperties.setProperty("retries", "0");
			producerProperties.putAll(secureProps);

			kafkaServer.produceIntoKafka(stream, topicName, serSchema, producerProperties, new Tuple2FlinkPartitioner(parallelism))
					.setParallelism(parallelism);

			try {
				writeEnv.execute("Write sequence");
			}
			catch (Exception e) {
				LOG.error("Write attempt failed, trying again", e);
				deleteTestTopic(topicName);
				waitUntilNoJobIsRunning(client);
				continue;
			}

			LOG.info("Finished writing sequence");

			// -------- Validate the Sequence --------

			// we need to validate the sequence, because kafka's producers are not exactly once
			LOG.info("Validating sequence");

			waitUntilNoJobIsRunning(client);

			if (validateSequence(topicName, parallelism, deserSchema, numElements)) {
				// everything is good!
				return topicName;
			}
			else {
				deleteTestTopic(topicName);
				// fall through the loop
			}
		}

		throw new Exception("Could not write a valid sequence to Kafka after " + maxNumAttempts + " attempts");
	}

	protected void writeAppendSequence(
			String topicName,
			final int originalNumElements,
			final int numElementsToAppend,
			final int parallelism) throws Exception {

		LOG.info("\n===================================\n" +
			"== Appending sequence of " + numElementsToAppend + " into " + topicName +
			"===================================");

		final TypeInformation<Tuple2<Integer, Integer>> resultType =
			TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {});

		final KeyedSerializationSchema<Tuple2<Integer, Integer>> serSchema =
			new KeyedSerializationSchemaWrapper<>(
				new TypeInformationSerializationSchema<>(resultType, new ExecutionConfig()));

		final KeyedDeserializationSchema<Tuple2<Integer, Integer>> deserSchema =
			new KeyedDeserializationSchemaWrapper<>(
				new TypeInformationSerializationSchema<>(resultType, new ExecutionConfig()));

		// -------- Write the append sequence --------

		StreamExecutionEnvironment writeEnv = StreamExecutionEnvironment.getExecutionEnvironment();
		writeEnv.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		writeEnv.getConfig().disableSysoutLogging();

		DataStream<Tuple2<Integer, Integer>> stream = writeEnv.addSource(new RichParallelSourceFunction<Tuple2<Integer, Integer>>() {

			private boolean running = true;

			@Override
			public void run(SourceContext<Tuple2<Integer, Integer>> ctx) throws Exception {
				int cnt = originalNumElements;
				int partition = getRuntimeContext().getIndexOfThisSubtask();

				while (running && cnt < numElementsToAppend + originalNumElements) {
					ctx.collect(new Tuple2<>(partition, cnt));
					cnt++;
				}
			}

			@Override
			public void cancel() {
				running = false;
			}
		}).setParallelism(parallelism);

		// the producer must not produce duplicates
		Properties producerProperties = FlinkKafkaProducerBase.getPropertiesFromBrokerList(brokerConnectionStrings);
		producerProperties.setProperty("retries", "0");
		producerProperties.putAll(secureProps);

		kafkaServer.produceIntoKafka(stream, topicName, serSchema, producerProperties, new Tuple2FlinkPartitioner(parallelism))
			.setParallelism(parallelism);

		try {
			writeEnv.execute("Write sequence");
		}
		catch (Exception e) {
			throw new Exception("Failed to append sequence to Kafka; append job failed.", e);
		}

		LOG.info("Finished writing append sequence");

		// we need to validate the sequence, because kafka's producers are not exactly once
		LOG.info("Validating sequence");
		while (!getRunningJobs(client).isEmpty()){
			Thread.sleep(50);
		}

		if (!validateSequence(topicName, parallelism, deserSchema, originalNumElements + numElementsToAppend)) {
			throw new Exception("Could not append a valid sequence to Kafka.");
		}
	}

	private boolean validateSequence(
			final String topic,
			final int parallelism,
			KeyedDeserializationSchema<Tuple2<Integer, Integer>> deserSchema,
			final int totalNumElements) throws Exception {

		final StreamExecutionEnvironment readEnv = StreamExecutionEnvironment.getExecutionEnvironment();
		readEnv.getConfig().setRestartStrategy(RestartStrategies.noRestart());
		readEnv.getConfig().disableSysoutLogging();
		readEnv.setParallelism(parallelism);

		Properties readProps = (Properties) standardProps.clone();
		readProps.setProperty("group.id", "flink-tests-validator");
		readProps.putAll(secureProps);
		FlinkKafkaConsumerBase<Tuple2<Integer, Integer>> consumer = kafkaServer.getConsumer(topic, deserSchema, readProps);
		consumer.setStartFromEarliest();

		readEnv
			.addSource(consumer)
			.map(new RichMapFunction<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>>() {

				private final int totalCount = parallelism * totalNumElements;
				private int count = 0;

				@Override
				public Tuple2<Integer, Integer> map(Tuple2<Integer, Integer> value) throws Exception {
					if (++count == totalCount) {
						throw new SuccessException();
					} else {
						return value;
					}
				}
			}).setParallelism(1)
			.addSink(new DiscardingSink<>()).setParallelism(1);

		final AtomicReference<Throwable> errorRef = new AtomicReference<>();

		JobGraph jobGraph = StreamingJobGraphGenerator.createJobGraph(readEnv.getStreamGraph());
		final JobID jobId = jobGraph.getJobID();

		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					client.setDetached(false);
					client.submitJob(jobGraph, KafkaConsumerTestBase.class.getClassLoader(), false);
					tryExecute(readEnv, "sequence validation");
				} catch (Throwable t) {
					if (!ExceptionUtils.findThrowable(t, SuccessException.class).isPresent()) {
						errorRef.set(t);
					}
				}
			}
		};
		runner.start();

		final long deadline = System.nanoTime() + 10_000_000_000L;
		long delay;
		while (runner.isAlive() && (delay = deadline - System.nanoTime()) > 0) {
			runner.join(delay / 1_000_000L);
		}

		boolean success;

		if (runner.isAlive()) {
			// did not finish in time, maybe the producer dropped one or more records and
			// the validation did not reach the exit point
			success = false;
			client.cancel(jobId);
		}
		else {
			Throwable error = errorRef.get();
			if (error != null) {
				success = false;
				LOG.info("Sequence validation job failed with exception", error);
			}
			else {
				success = true;
			}
		}

		waitUntilNoJobIsRunning(client);

		return success;
	}

	// ------------------------------------------------------------------------
	//  Debugging utilities
	// ------------------------------------------------------------------------

	/**
	 * Read topic to list, only using Kafka code.
	 */
	private static List<MessageAndMetadata<byte[], byte[]>> readTopicToList(String topicName, ConsumerConfig config, final int stopAfter) {
		ConsumerConnector consumerConnector = Consumer.createJavaConsumerConnector(config);
		// we request only one stream per consumer instance. Kafka will make sure that each consumer group
		// will see each message only once.
		Map<String, Integer> topicCountMap = Collections.singletonMap(topicName, 1);
		Map<String, List<KafkaStream<byte[], byte[]>>> streams = consumerConnector.createMessageStreams(topicCountMap);
		if (streams.size() != 1) {
			throw new RuntimeException("Expected only one message stream but got " + streams.size());
		}
		List<KafkaStream<byte[], byte[]>> kafkaStreams = streams.get(topicName);
		if (kafkaStreams == null) {
			throw new RuntimeException("Requested stream not available. Available streams: " + streams.toString());
		}
		if (kafkaStreams.size() != 1) {
			throw new RuntimeException("Requested 1 stream from Kafka, bot got " + kafkaStreams.size() + " streams");
		}
		LOG.info("Opening Consumer instance for topic '{}' on group '{}'", topicName, config.groupId());
		ConsumerIterator<byte[], byte[]> iteratorToRead = kafkaStreams.get(0).iterator();

		List<MessageAndMetadata<byte[], byte[]>> result = new ArrayList<>();
		int read = 0;
		while (iteratorToRead.hasNext()) {
			read++;
			result.add(iteratorToRead.next());
			if (read == stopAfter) {
				LOG.info("Read " + read + " elements");
				return result;
			}
		}
		return result;
	}

	private static class BrokerKillingMapper<T> extends RichMapFunction<T, T>
			implements ListCheckpointed<Integer>, CheckpointListener {

		private static final long serialVersionUID = 6334389850158707313L;

		public static volatile boolean killedLeaderBefore;
		public static volatile boolean hasBeenCheckpointedBeforeFailure;

		private final int shutdownBrokerId;
		private final int failCount;
		private int numElementsTotal;

		private boolean failer;
		private boolean hasBeenCheckpointed;

		public BrokerKillingMapper(int shutdownBrokerId, int failCount) {
			this.shutdownBrokerId = shutdownBrokerId;
			this.failCount = failCount;
		}

		@Override
		public void open(Configuration parameters) {
			failer = getRuntimeContext().getIndexOfThisSubtask() == 0;
		}

		@Override
		public T map(T value) throws Exception {
			numElementsTotal++;

			if (!killedLeaderBefore) {
				Thread.sleep(10);

				if (failer && numElementsTotal >= failCount) {
					// shut down a Kafka broker
					KafkaServer toShutDown = null;
					for (KafkaServer server : kafkaServer.getBrokers()) {

						if (kafkaServer.getBrokerId(server) == shutdownBrokerId) {
							toShutDown = server;
							break;
						}
					}

					if (toShutDown == null) {
						StringBuilder listOfBrokers = new StringBuilder();
						for (KafkaServer server : kafkaServer.getBrokers()) {
							listOfBrokers.append(kafkaServer.getBrokerId(server));
							listOfBrokers.append(" ; ");
						}

						throw new Exception("Cannot find broker to shut down: " + shutdownBrokerId
								+ " ; available brokers: " + listOfBrokers.toString());
					}
					else {
						hasBeenCheckpointedBeforeFailure = hasBeenCheckpointed;
						killedLeaderBefore = true;
						toShutDown.shutdown();
					}
				}
			}
			return value;
		}

		@Override
		public void notifyCheckpointComplete(long checkpointId) {
			hasBeenCheckpointed = true;
		}

		@Override
		public List<Integer> snapshotState(long checkpointId, long timestamp) throws Exception {
			return Collections.singletonList(this.numElementsTotal);
		}

		@Override
		public void restoreState(List<Integer> state) throws Exception {
			if (state.isEmpty() || state.size() > 1) {
				throw new RuntimeException("Test failed due to unexpected recovered state size " + state.size());
			}
			this.numElementsTotal = state.get(0);
		}
	}

	private static class Tuple2WithTopicSchema implements KeyedDeserializationSchema<Tuple3<Integer, Integer, String>>,
		KeyedSerializationSchema<Tuple3<Integer, Integer, String>> {

		private final TypeSerializer<Tuple2<Integer, Integer>> ts;

		public Tuple2WithTopicSchema(ExecutionConfig ec) {
			ts = TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>(){}).createSerializer(ec);
		}

		@Override
		public Tuple3<Integer, Integer, String> deserialize(byte[] messageKey, byte[] message, String topic, int partition, long offset) throws IOException {
			DataInputView in = new DataInputViewStreamWrapper(new ByteArrayInputStream(message));
			Tuple2<Integer, Integer> t2 = ts.deserialize(in);
			return new Tuple3<>(t2.f0, t2.f1, topic);
		}

		@Override
		public boolean isEndOfStream(Tuple3<Integer, Integer, String> nextElement) {
			return false;
		}

		@Override
		public TypeInformation<Tuple3<Integer, Integer, String>> getProducedType() {
			return TypeInformation.of(new TypeHint<Tuple3<Integer, Integer, String>>(){});
		}

		@Override
		public byte[] serializeKey(Tuple3<Integer, Integer, String> element) {
			return null;
		}

		@Override
		public byte[] serializeValue(Tuple3<Integer, Integer, String> element) {
			ByteArrayOutputStream by = new ByteArrayOutputStream();
			DataOutputView out = new DataOutputViewStreamWrapper(by);
			try {
				ts.serialize(new Tuple2<>(element.f0, element.f1), out);
			} catch (IOException e) {
				throw new RuntimeException("Error" , e);
			}
			return by.toByteArray();
		}

		@Override
		public String getTargetTopic(Tuple3<Integer, Integer, String> element) {
			return element.f2;
		}
	}
}