flink-instrumentation/workspace/flink-1.4.1-instrumented/flink-1.4.1/flink-connectors/flink-connector-kafka-base/src/test/java/org/apache/flink/streaming/connectors/kafka/KafkaConsumerTestBase.java
/**
	 * 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.
	 */
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.
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 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 = TypeInfoParser.parse("Tuple2<Integer, Integer>");
		final TypeInformationSerializationSchema<Tuple2<Integer, Integer>> deser =
			new TypeInformationSerializationSchema<>(intIntTupleType, env.getConfig());
		cc.putAll(secureProps);
		FlinkKafkaConsumerBase<Tuple2<Integer, Integer>> consumer = kafkaServer.getConsumer(topicName, deser, cc);
		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;
		}
		DataStream<Tuple2<Integer, Integer>> source = env
			.addSource(consumer).setParallelism(sourceParallelism)
			.map(new ThrottledMapper<Tuple2<Integer, Integer>>(20)).setParallelism(sourceParallelism);
		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);
				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());
						}
					}
					throw new SuccessException();
				}
			}
		}).setParallelism(1);
		tryExecute(env, "Read data from Kafka");
		LOG.info("Successfully read sequence for verification");
	}
protected void readSequence ( final StreamExecutionEnvironment env , final StartupMode startupMode , final Map < KafkaTopicPartition , Long > specificStartupOffsets , 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 = TypeInfoParser . parse ( "Tuple2<Integer, Integer>" ) ; final TypeInformationSerializationSchema < Tuple2 < Integer , Integer > > deser = new TypeInformationSerializationSchema < > ( intIntTupleType , env . getConfig ( ) ) ; cc . putAll ( secureProps ) ; FlinkKafkaConsumerBase < Tuple2 < Integer , Integer > > consumer = kafkaServer . getConsumer ( topicName , deser , cc ) ; 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 ; } DataStream < Tuple2 < Integer , Integer > > source = env . addSource ( consumer ) . setParallelism ( sourceParallelism ) . map ( new ThrottledMapper < Tuple2 < Integer , Integer > > ( 20 ) ) . setParallelism ( sourceParallelism ) ; 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 ) ; 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 ( ) ) ; } } throw new SuccessException ( ) ; } } } ) . setParallelism ( 1 ) ; tryExecute ( env , "Read data from Kafka" ) ; LOG . info ( "Successfully read sequence for verification" ) ; }
Runs a job using the provided environment to read a sequence of records from a single Kafka topic.
Runs a job using the provided environment to read a sequence of records from a single Kafka topic .
// create the consumer // verify data // verify if we've seen everything // test has passed