messageFields = messageSchema.getFields();
+ if (messageFields.size() > 1) {
+ for (Schema.Field messageField : messageFields) {
+ collector.addFailure(
+ "Without a format, the schema must contain just a single message field of type bytes or nullable bytes.",
+ String.format("Remove field '%s'.", messageField.getName()))
+ .withOutputSchemaField(messageField.getName()).withConfigProperty(NAME_FORMAT);
+ }
+ return;
+ }
+
+ Schema.Field messageField = messageFields.get(0);
+ Schema messageFieldSchema = messageField.getSchema().isNullable() ? messageField.getSchema().getNonNullable() :
+ messageField.getSchema();
+ Schema.Type messageFieldType = messageFieldSchema.getType();
+ if (messageFieldType != Schema.Type.BYTES || messageFieldSchema.getLogicalType() != null) {
+ collector.addFailure(
+ String.format("Without a format, the message field must be of type bytes or nullable bytes, " +
+ "but field '%s' is of type '%s'.",
+ messageField.getName(), messageField.getSchema().getDisplayName()), null)
+ .withOutputSchemaField(messageField.getName()).withConfigProperty(NAME_FORMAT);
+ }
+ }
+ }
+}
diff --git a/confluent-kafka-plugins/src/main/java/io/cdap/plugin/confluent/streaming/source/ConfluentStreamingSourceUtil.java b/confluent-kafka-plugins/src/main/java/io/cdap/plugin/confluent/streaming/source/ConfluentStreamingSourceUtil.java
new file mode 100644
index 0000000..bec9944
--- /dev/null
+++ b/confluent-kafka-plugins/src/main/java/io/cdap/plugin/confluent/streaming/source/ConfluentStreamingSourceUtil.java
@@ -0,0 +1,417 @@
+/*
+ * Copyright © 2019 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.plugin.confluent.streaming.source;
+
+import com.google.common.base.Joiner;
+import com.google.common.base.Strings;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Sets;
+import io.cdap.cdap.api.data.format.FormatSpecification;
+import io.cdap.cdap.api.data.format.RecordFormat;
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.format.UnexpectedFormatException;
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.etl.api.FailureCollector;
+import io.cdap.cdap.etl.api.streaming.StreamingContext;
+import io.cdap.cdap.format.RecordFormats;
+import io.cdap.plugin.confluent.common.KafkaHelpers;
+import io.cdap.plugin.format.avro.AvroToStructuredTransformer;
+import io.confluent.kafka.serializers.KafkaAvroDeserializer;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.requests.ListOffsetRequest;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.streaming.Time;
+import org.apache.spark.streaming.api.java.JavaDStream;
+import org.apache.spark.streaming.kafka010.ConsumerStrategies;
+import org.apache.spark.streaming.kafka010.KafkaUtils;
+import org.apache.spark.streaming.kafka010.LocationStrategies;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Set;
+import javax.annotation.Nonnull;
+
+/**
+ * Util method for {@link ConfluentStreamingSource}.
+ *
+ * This class contains methods for {@link ConfluentStreamingSource} that require spark classes because during validation
+ * spark classes are not available. Refer CDAP-15912 for more information.
+ */
+final class ConfluentStreamingSourceUtil {
+ private static final Logger LOG = LoggerFactory.getLogger(ConfluentStreamingSourceUtil.class);
+
+ private ConfluentStreamingSourceUtil() {
+ // no-op
+ }
+
+ /**
+ * Returns {@link JavaDStream} for {@link ConfluentStreamingSource}.
+ * @param context streaming context
+ * @param conf kafka conf
+ * @param outputSchema source output schema
+ * @param collector failure collector
+ */
+ static JavaDStream getStructuredRecordJavaDStream(
+ StreamingContext context, ConfluentStreamingSourceConfig conf, Schema outputSchema, FailureCollector collector) {
+ String pipelineName = context.getPipelineName();
+ Map kafkaParams = getConsumerParams(conf, pipelineName);
+ Properties properties = new Properties();
+ properties.putAll(kafkaParams);
+ try (Consumer consumer = new KafkaConsumer<>(properties, new ByteArrayDeserializer(),
+ new ByteArrayDeserializer())) {
+ Map offsets = getOffsets(conf, collector, consumer);
+ LOG.info("Using initial offsets {}", offsets);
+
+ if (conf.getSchemaRegistryUrl() != null) {
+ AvroRecordTransform transform = new AvroRecordTransform(conf, outputSchema);
+ return createKafkaDirectStream(context, conf, kafkaParams, offsets, transform);
+ }
+ return createKafkaDirectStream(context, conf, kafkaParams, offsets, new RecordTransform(conf, outputSchema));
+ }
+ }
+
+ private static JavaDStream createKafkaDirectStream(
+ StreamingContext context,
+ ConfluentStreamingSourceConfig conf,
+ Map kafkaParams,
+ Map offsets,
+ Function2>, Time, JavaRDD> transform
+ ) {
+ return KafkaUtils.createDirectStream(
+ context.getSparkStreamingContext(), LocationStrategies.PreferConsistent(),
+ ConsumerStrategies.Subscribe(Collections.singleton(conf.getTopic()), kafkaParams, offsets)
+ ).transform(transform);
+ }
+
+ @Nonnull
+ private static Map getConsumerParams(ConfluentStreamingSourceConfig conf, String pipelineName) {
+ Map kafkaParams = new HashMap<>();
+ kafkaParams.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, conf.getBrokers());
+ // Spark saves the offsets in checkpoints, no need for Kafka to save them
+ kafkaParams.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "https");
+ kafkaParams.put(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG, "500");
+ kafkaParams.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_SSL");
+ kafkaParams.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
+ kafkaParams.put(SaslConfigs.SASL_JAAS_CONFIG, "org.apache.kafka.common.security.plain.PlainLoginModule required " +
+ "username=" + conf.getClusterApiKey() + " password=" + conf.getClusterApiSecret() + ";");
+
+ if (!Strings.isNullOrEmpty(conf.getSchemaRegistryUrl())) {
+ kafkaParams.put("schema.registry.url", conf.getSchemaRegistryUrl());
+ kafkaParams.put("basic.auth.credentials.source", "USER_INFO");
+ kafkaParams.put("schema.registry.basic.auth.user.info",
+ conf.getSchemaRegistryApiKey() + ":" + conf.getSchemaRegistryApiSecret());
+ kafkaParams.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, KafkaAvroDeserializer.class.getCanonicalName());
+ kafkaParams.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, KafkaAvroDeserializer.class.getCanonicalName());
+ } else {
+ kafkaParams.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getCanonicalName());
+ kafkaParams.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getCanonicalName());
+ }
+ kafkaParams.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
+ // Create a unique string for the group.id using the pipeline name and the topic.
+ // group.id is a Kafka consumer property that uniquely identifies the group of
+ // consumer processes to which this consumer belongs.
+ String groupId = Joiner.on("-")
+ .join(pipelineName.length(), conf.getTopic().length(), pipelineName, conf.getTopic());
+ kafkaParams.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+ kafkaParams.putAll(conf.getKafkaProperties());
+ // change the request timeout to fetch the metadata to be 15 seconds or 1 second greater than session time out ms,
+ // since this config has to be greater than the session time out, which is by default 10 seconds
+ // the KafkaConsumer at runtime should still use the default timeout 305 seconds or whatever the user provides in
+ // kafkaConf
+ int requestTimeout =
+ Integer.parseInt(conf.getKafkaProperties().getOrDefault(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG, "20000"));
+ if (conf.getKafkaProperties().containsKey(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG)) {
+ int sessionTimeout =
+ Integer.parseInt(conf.getKafkaProperties().get(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG) + 1000);
+ requestTimeout = Math.max(requestTimeout, sessionTimeout);
+ }
+ kafkaParams.put(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG, requestTimeout);
+ return kafkaParams;
+ }
+
+ @Nonnull
+ private static Map getOffsets(ConfluentStreamingSourceConfig conf, FailureCollector collector,
+ Consumer consumer) {
+ Map offsets = conf.getInitialPartitionOffsets(
+ getPartitions(consumer, conf, collector), collector);
+ collector.getOrThrowException();
+
+ // KafkaUtils doesn't understand -1 and -2 as smallest offset and latest offset.
+ // so we have to replace them with the actual smallest and latest
+ List earliestOffsetRequest = new ArrayList<>();
+ List latestOffsetRequest = new ArrayList<>();
+ for (Map.Entry entry : offsets.entrySet()) {
+ TopicPartition topicAndPartition = entry.getKey();
+ Long offset = entry.getValue();
+ if (offset == ListOffsetRequest.EARLIEST_TIMESTAMP) {
+ earliestOffsetRequest.add(topicAndPartition);
+ } else if (offset == ListOffsetRequest.LATEST_TIMESTAMP) {
+ latestOffsetRequest.add(topicAndPartition);
+ }
+ }
+
+ Set allOffsetRequest =
+ Sets.newHashSet(Iterables.concat(earliestOffsetRequest, latestOffsetRequest));
+ Map offsetsFound = new HashMap<>();
+ offsetsFound.putAll(KafkaHelpers.getEarliestOffsets(consumer, earliestOffsetRequest));
+ offsetsFound.putAll(KafkaHelpers.getLatestOffsets(consumer, latestOffsetRequest));
+ for (TopicPartition topicAndPartition : allOffsetRequest) {
+ offsets.put(topicAndPartition, offsetsFound.get(topicAndPartition));
+ }
+
+ Set missingOffsets = Sets.difference(allOffsetRequest, offsetsFound.keySet());
+ if (!missingOffsets.isEmpty()) {
+ throw new IllegalStateException(String.format(
+ "Could not find offsets for %s. Please check all brokers were included in the broker list.", missingOffsets));
+ }
+ return offsets;
+ }
+
+ private static Set getPartitions(Consumer consumer, ConfluentStreamingSourceConfig conf,
+ FailureCollector collector) {
+ Set partitions = conf.getPartitions(collector);
+ collector.getOrThrowException();
+
+ if (!partitions.isEmpty()) {
+ return partitions;
+ }
+
+ partitions = new HashSet<>();
+ for (PartitionInfo partitionInfo : consumer.partitionsFor(conf.getTopic())) {
+ partitions.add(partitionInfo.partition());
+ }
+ return partitions;
+ }
+
+ /**
+ * Applies the format function to each rdd.
+ */
+ private static class AvroRecordTransform
+ implements Function2>, Time, JavaRDD> {
+
+ private final ConfluentStreamingSourceConfig conf;
+ private final Schema outputSchema;
+
+ AvroRecordTransform(ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ this.conf = conf;
+ this.outputSchema = outputSchema;
+ }
+
+ @Override
+ public JavaRDD call(JavaRDD> input, Time batchTime) {
+ return input.map(new AvroFunction(batchTime.milliseconds(), conf, outputSchema));
+ }
+ }
+
+ /**
+ * Applies the format function to each rdd.
+ */
+ private static class RecordTransform
+ implements Function2>, Time, JavaRDD> {
+
+ private final ConfluentStreamingSourceConfig conf;
+ private final Schema outputSchema;
+
+ RecordTransform(ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ this.conf = conf;
+ this.outputSchema = outputSchema;
+ }
+
+ @Override
+ public JavaRDD call(JavaRDD> input, Time batchTime) {
+ Function, StructuredRecord> recordFunction = conf.getFormat() == null ?
+ new BytesFunction(batchTime.milliseconds(), conf, outputSchema) :
+ new FormatFunction(batchTime.milliseconds(), conf, outputSchema);
+ return input.map(recordFunction);
+ }
+ }
+
+ /**
+ * Common logic for transforming kafka key, message, partition, and offset into a structured record.
+ * Everything here should be serializable, as Spark Streaming will serialize all functions.
+ */
+ private abstract static class BaseFunction implements Function, StructuredRecord> {
+ protected final ConfluentStreamingSourceConfig conf;
+ private final long ts;
+ private final Schema outputSchema;
+
+ BaseFunction(long ts, ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ this.ts = ts;
+ this.conf = conf;
+ this.outputSchema = outputSchema;
+ }
+
+ @Override
+ public StructuredRecord call(ConsumerRecord in) throws Exception {
+ String timeField = conf.getTimeField();
+ String keyField = conf.getKeyField();
+ String partitionField = conf.getPartitionField();
+ String offsetField = conf.getOffsetField();
+ StructuredRecord.Builder builder = StructuredRecord.builder(outputSchema);
+ if (timeField != null) {
+ builder.set(timeField, ts);
+ }
+ if (keyField != null) {
+ builder.set(keyField, convertKey(in.key()));
+ }
+ if (partitionField != null) {
+ builder.set(partitionField, in.partition());
+ }
+ if (offsetField != null) {
+ builder.set(offsetField, in.offset());
+ }
+ addMessage(builder, in.value());
+ return builder.build();
+ }
+
+ protected abstract Object convertKey(K key);
+
+ protected abstract void addMessage(StructuredRecord.Builder builder, V message) throws Exception;
+ }
+
+ private abstract static class BinaryBaseFunction extends BaseFunction {
+ BinaryBaseFunction(long ts, ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ super(ts, conf, outputSchema);
+ }
+
+ @Override
+ protected Object convertKey(byte[] key) {
+ if (key == null) {
+ return null;
+ }
+ Schema keySchemaNullable = conf.getSchema().getField(conf.getKeyField()).getSchema();
+ Schema keySchema = keySchemaNullable.isNullable() ? keySchemaNullable.getNonNullable() : keySchemaNullable;
+ if (keySchema.getType() == Schema.Type.STRING) {
+ return new String(key, StandardCharsets.UTF_8);
+ }
+ if (keySchema.getType() == Schema.Type.BYTES) {
+ return key;
+ }
+ throw new IllegalStateException(String.format("Unexpected key type '%s'", keySchema.getDisplayName()));
+ }
+ }
+
+ /**
+ * Transforms kafka key and message into a structured record when message format is not given.
+ * Everything here should be serializable, as Spark Streaming will serialize all functions.
+ */
+ private static class BytesFunction extends BinaryBaseFunction {
+ private transient String messageField;
+
+ BytesFunction(long ts, ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ super(ts, conf, outputSchema);
+ }
+
+ @Override
+ protected void addMessage(StructuredRecord.Builder builder, byte[] message) {
+ builder.set(getMessageField(), message);
+ }
+
+ private String getMessageField() {
+ if (messageField == null) {
+ for (Schema.Field field : conf.getSchema().getFields()) {
+ String name = field.getName();
+ if (!name.equals(conf.getTimeField()) && !name.equals(conf.getKeyField())
+ && !name.equals(conf.getOffsetField()) && !name.equals(conf.getPartitionField())) {
+ messageField = name;
+ break;
+ }
+ }
+ if (messageField == null) {
+ throw new IllegalStateException("No message field found in schema");
+ }
+ }
+ return messageField;
+ }
+ }
+
+ /**
+ * Transforms kafka key and message into a structured record when message format and schema are given.
+ * Everything here should be serializable, as Spark Streaming will serialize all functions.
+ */
+ private static class FormatFunction extends BinaryBaseFunction {
+ private transient RecordFormat recordFormat;
+
+ FormatFunction(long ts, ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ super(ts, conf, outputSchema);
+ }
+
+ @Override
+ protected void addMessage(StructuredRecord.Builder builder, byte[] message) throws Exception {
+ // first time this was called, initialize record format
+ if (recordFormat == null) {
+ Schema messageSchema = conf.getMessageSchema();
+ FormatSpecification spec = new FormatSpecification(conf.getFormat(), messageSchema, new HashMap<>());
+ recordFormat = RecordFormats.createInitializedFormat(spec);
+ }
+
+ StructuredRecord messageRecord = recordFormat.read(ByteBuffer.wrap(message));
+ for (Schema.Field field : messageRecord.getSchema().getFields()) {
+ String fieldName = field.getName();
+ builder.set(fieldName, messageRecord.get(fieldName));
+ }
+ }
+ }
+
+ private static class AvroFunction extends BaseFunction {
+ private transient AvroToStructuredTransformer transformer;
+
+ AvroFunction(long ts, ConfluentStreamingSourceConfig conf, Schema outputSchema) {
+ super(ts, conf, outputSchema);
+ }
+
+ @Override
+ protected Object convertKey(Object key) {
+ return key;
+ }
+
+ @Override
+ protected void addMessage(StructuredRecord.Builder builder, Object message) throws Exception {
+ if (transformer == null) {
+ transformer = new AvroToStructuredTransformer();
+ }
+ if (!(message instanceof GenericRecord)) {
+ throw new UnexpectedFormatException(
+ String.format("Unexpected message class '%s'", message.getClass().getName()));
+ }
+ GenericRecord genericRecord = (GenericRecord) message;
+ StructuredRecord messageRecord = transformer.transform(genericRecord);
+ builder.set(conf.getValueField(), messageRecord);
+ }
+ }
+}
diff --git a/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/StructuredRecordRepresentation.java b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/StructuredRecordRepresentation.java
new file mode 100644
index 0000000..53c6a33
--- /dev/null
+++ b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/StructuredRecordRepresentation.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright © 2019 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.plugin.confluent;
+
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.format.StructuredRecordStringConverter;
+import org.assertj.core.presentation.StandardRepresentation;
+
+import java.io.IOException;
+
+public class StructuredRecordRepresentation extends StandardRepresentation {
+ @Override
+ public String toStringOf(Object object) {
+ try {
+ if (object instanceof StructuredRecord) {
+ return StructuredRecordStringConverter.toJsonString((StructuredRecord) object);
+ }
+ return super.toStringOf(object);
+ } catch (IOException e) {
+ throw new IllegalArgumentException(e);
+ }
+ }
+}
diff --git a/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/KafkaTestUtils.java b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/KafkaTestUtils.java
new file mode 100644
index 0000000..fc5c109
--- /dev/null
+++ b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/KafkaTestUtils.java
@@ -0,0 +1,174 @@
+/*
+ * Copyright © 2019 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.plugin.confluent.integration;
+
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.plugin.format.avro.StructuredToAvroTransformer;
+import io.confluent.kafka.serializers.KafkaAvroDeserializer;
+import io.confluent.kafka.serializers.KafkaAvroSerializer;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.config.SaslConfigs;
+import org.apache.kafka.common.config.SslConfigs;
+import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.serialization.StringDeserializer;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+
+public class KafkaTestUtils {
+ // private static final String KAFKA_SERVER = System.getProperty("test.kafka_server");
+ // private static final String CLUSTER_API_KEY = System.getProperty("test.cluster_api_key");
+ // private static final String CLUSTER_API_SECRET = System.getProperty("test.cluster_api_secret");
+ public static final String KAFKA_SERVER = "pkc-4yyd6.us-east1.gcp.confluent.cloud:9092";
+ public static final String CLUSTER_API_KEY = "GI4SVX36F5QTTRKW";
+ public static final String CLUSTER_API_SECRET = "A6S7gXegaDVnTYQS7aaK2TR4R3nzrVlg19ValdS1dplLfXupuTMSeMGN91kI4x6B";
+ public static final String SR_URL = "https://psrc-lz3xz.us-central1.gcp.confluent.cloud";
+ public static final String SR_API_KEY = "DSIGEYOAXB47HERZ";
+ public static final String SR_API_SECRET = "gmMj3rHhI/JqEVrc1R4htJxU/O4A9KkNf2YcN0NpY/EZY0eP0GFNZnsXGpkSHKkr";
+
+ private KafkaTestUtils() {
+ throw new AssertionError("Can not be initialized");
+ }
+
+ public static Consumer createConsumer() {
+ Map props = new HashMap<>();
+ props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
+ props.put(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG, "500");
+ props.put(ConsumerConfig.GROUP_ID_CONFIG, UUID.randomUUID().toString());
+ props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+ props.putAll(getSecurityProps());
+ return new KafkaConsumer<>(props, new StringDeserializer(), new StringDeserializer());
+ }
+
+ public static Consumer createConsumerForSchemaRegistry() {
+ Map props = new HashMap<>();
+ props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
+ props.put(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG, "500");
+ props.put(ConsumerConfig.GROUP_ID_CONFIG, UUID.randomUUID().toString());
+ props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+ props.putAll(getSecurityProps());
+ props.put("schema.registry.url", SR_URL);
+ props.put("basic.auth.credentials.source", "USER_INFO");
+ props.put("schema.registry.basic.auth.user.info", SR_API_KEY + ":" + SR_API_SECRET);
+ props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, KafkaAvroDeserializer.class.getCanonicalName());
+ props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, KafkaAvroDeserializer.class.getCanonicalName());
+ return new KafkaConsumer<>(props);
+ }
+
+ public static KafkaProducer createProducer() {
+ Map props = new HashMap<>();
+ props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
+ props.put(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG, "500");
+ props.put(ProducerConfig.ACKS_CONFIG, "all");
+ props.putAll(getSecurityProps());
+ return new KafkaProducer<>(props, new ByteArraySerializer(), new ByteArraySerializer());
+ }
+
+ public static KafkaProducer createProducerForSchemaRegistry() {
+ Map props = new HashMap<>();
+ props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
+ props.put(CommonClientConfigs.RETRY_BACKOFF_MS_CONFIG, "500");
+ props.put(ProducerConfig.ACKS_CONFIG, "all");
+ props.putAll(getSecurityProps());
+ props.put("schema.registry.url", SR_URL);
+ props.put("basic.auth.credentials.source", "USER_INFO");
+ props.put("schema.registry.basic.auth.user.info", SR_API_KEY + ":" + SR_API_SECRET);
+ props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, KafkaAvroSerializer.class.getCanonicalName());
+ props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, KafkaAvroSerializer.class.getCanonicalName());
+ return new KafkaProducer<>(props);
+ }
+
+ private static AdminClient createAdminClient() {
+ Map props = new HashMap<>();
+ props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, KAFKA_SERVER);
+ props.putAll(getSecurityProps());
+ return AdminClient.create(props);
+ }
+
+ private static Map getSecurityProps() {
+ Map props = new HashMap<>();
+ props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SASL_SSL");
+ props.put(SslConfigs.SSL_ENDPOINT_IDENTIFICATION_ALGORITHM_CONFIG, "https");
+ props.put(SaslConfigs.SASL_MECHANISM, "PLAIN");
+ props.put(SaslConfigs.SASL_JAAS_CONFIG, "org.apache.kafka.common.security.plain.PlainLoginModule required " +
+ "username=" + CLUSTER_API_KEY + " password=" + CLUSTER_API_SECRET + ";");
+ return props;
+ }
+
+ // Create topic in Confluent Cloud
+ public static void createTopic(String topic, int partitions, int replication) {
+ NewTopic newTopic = new NewTopic(topic, partitions, (short) replication);
+ try (AdminClient adminClient = createAdminClient()) {
+ adminClient.createTopics(Collections.singletonList(newTopic)).all().get();
+ } catch (InterruptedException | ExecutionException e) {
+ // Ignore if TopicExistsException, which may be valid if topic exists
+ if (!(e.getCause() instanceof TopicExistsException)) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ public static void deleteTopic(String topic) {
+ try (AdminClient adminClient = createAdminClient()) {
+ adminClient.deleteTopics(Collections.singletonList(topic)).all().get();
+ } catch (InterruptedException | ExecutionException e) {
+ // Ignore if UnknownTopicOrPartitionException, which may be valid if topic does not exist
+ if (!(e.getCause() instanceof UnknownTopicOrPartitionException)) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ public static GenericRecord toGenericRecord(StructuredRecord structuredRecord, Schema schema) {
+ StructuredToAvroTransformer transformer = new StructuredToAvroTransformer(schema);
+ try {
+ return transformer.transform(structuredRecord);
+ } catch (IOException e) {
+ throw new IllegalStateException("Failed to convert records", e);
+ }
+ }
+
+ public static List toGenericRecords(List structuredRecords, Schema schema) {
+ StructuredToAvroTransformer transformer = new StructuredToAvroTransformer(schema);
+ try {
+ List genericRecords = new ArrayList<>();
+ for (StructuredRecord structuredRecord : structuredRecords) {
+ genericRecords.add(transformer.transform(structuredRecord));
+ }
+ return genericRecords;
+ } catch (IOException e) {
+ throw new IllegalStateException("Failed to convert records", e);
+ }
+ }
+}
diff --git a/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/ConfluentStreamingTestBase.java b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/ConfluentStreamingTestBase.java
new file mode 100644
index 0000000..a76bdb3
--- /dev/null
+++ b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/ConfluentStreamingTestBase.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright © 2019 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.plugin.confluent.integration.streaming;
+
+import io.cdap.cdap.api.artifact.ArtifactSummary;
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.dataset.table.Table;
+import io.cdap.cdap.common.conf.Constants;
+import io.cdap.cdap.datastreams.DataStreamsApp;
+import io.cdap.cdap.datastreams.DataStreamsSparkLauncher;
+import io.cdap.cdap.etl.mock.batch.MockSink;
+import io.cdap.cdap.etl.mock.test.HydratorTestBase;
+import io.cdap.cdap.etl.proto.v2.DataStreamsConfig;
+import io.cdap.cdap.etl.proto.v2.ETLPlugin;
+import io.cdap.cdap.etl.proto.v2.ETLStage;
+import io.cdap.cdap.etl.spark.Compat;
+import io.cdap.cdap.proto.artifact.AppRequest;
+import io.cdap.cdap.proto.id.ApplicationId;
+import io.cdap.cdap.proto.id.ArtifactId;
+import io.cdap.cdap.proto.id.NamespaceId;
+import io.cdap.cdap.test.ApplicationManager;
+import io.cdap.cdap.test.DataSetManager;
+import io.cdap.cdap.test.SparkManager;
+import io.cdap.cdap.test.TestConfiguration;
+import io.cdap.plugin.confluent.StructuredRecordRepresentation;
+import io.cdap.plugin.confluent.streaming.sink.ConfluentStreamingSink;
+import io.cdap.plugin.confluent.streaming.source.ConfluentStreamingSource;
+import io.confluent.kafka.serializers.KafkaAvroDeserializer;
+import io.confluent.kafka.serializers.KafkaAvroSerializer;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.spark.streaming.kafka010.KafkaUtils;
+import org.assertj.core.api.Assertions;
+import org.awaitility.Awaitility;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+public abstract class ConfluentStreamingTestBase extends HydratorTestBase {
+ @ClassRule
+ public static final TestConfiguration CONFIG =
+ new TestConfiguration(Constants.Explore.EXPLORE_ENABLED, false,
+ Constants.AppFabric.SPARK_COMPAT, Compat.SPARK_COMPAT);
+ private static final Logger LOG = LoggerFactory.getLogger(ConfluentStreamingTestBase.class);
+ private static final ArtifactId APP_ARTIFACT_ID = NamespaceId.DEFAULT.artifact("data-streams", "1.0.0");
+ private static final ArtifactSummary APP_ARTIFACT = new ArtifactSummary("data-streams", "1.0.0");
+
+ @BeforeClass
+ public static void setupBasic() throws Exception {
+ LOG.info("Setting up application");
+
+ setupStreamingArtifacts(APP_ARTIFACT_ID, DataStreamsApp.class);
+
+ LOG.info("Setting up plugins");
+ addPluginArtifact(
+ NamespaceId.DEFAULT.artifact("confluent-kafka-plugins", "1.0.0"),
+ APP_ARTIFACT_ID,
+ ConfluentStreamingSource.class, ConfluentStreamingSink.class,
+ KafkaUtils.class, TopicPartition.class,
+ ByteArrayDeserializer.class, ByteArraySerializer.class,
+ KafkaAvroDeserializer.class, KafkaAvroSerializer.class
+ );
+ }
+
+ protected SparkManager deployETL(ETLPlugin sourcePlugin, ETLPlugin sinkPlugin, String appName) throws Exception {
+ ETLStage source = new ETLStage("source", sourcePlugin);
+ ETLStage sink = new ETLStage("sink", sinkPlugin);
+ DataStreamsConfig etlConfig = DataStreamsConfig.builder()
+ .addStage(source)
+ .addStage(sink)
+ .addConnection(source.getName(), sink.getName())
+ .setBatchInterval("1s")
+ .setStopGracefully(true)
+ .build();
+
+ AppRequest appRequest = new AppRequest<>(APP_ARTIFACT, etlConfig);
+ ApplicationId appId = NamespaceId.DEFAULT.app(appName);
+ ApplicationManager applicationManager = deployApplication(appId, appRequest);
+ return getProgramManager(applicationManager);
+ }
+
+ protected List waitForRecords(String outputTable, int messageCount) throws Exception {
+ DataSetManager outputManager = getDataset(outputTable);
+ Awaitility.await().atMost(60, TimeUnit.SECONDS).untilAsserted(() -> {
+ List output = MockSink.readOutput(outputManager);
+ Assertions.assertThat(output)
+ .withRepresentation(new StructuredRecordRepresentation())
+ .hasSizeGreaterThanOrEqualTo(messageCount);
+ });
+
+ List output = MockSink.readOutput(outputManager);
+ Assertions.assertThat(output)
+ .withRepresentation(new StructuredRecordRepresentation())
+ .hasSize(messageCount);
+ return output;
+ }
+
+ protected void waitForRecords(String outputTable, List expectedRecords) throws Exception {
+ DataSetManager outputManager = getDataset(outputTable);
+ Awaitility.await().atMost(60, TimeUnit.SECONDS).untilAsserted(() -> {
+ List output = MockSink.readOutput(outputManager);
+ Assertions.assertThat(output)
+ .withRepresentation(new StructuredRecordRepresentation())
+ .hasSizeGreaterThanOrEqualTo(expectedRecords.size());
+ });
+
+ List output = MockSink.readOutput(outputManager);
+
+ Assertions.assertThat(output)
+ .withRepresentation(new StructuredRecordRepresentation())
+ .containsExactlyInAnyOrderElementsOf(expectedRecords);
+ }
+
+ private SparkManager getProgramManager(ApplicationManager appManager) {
+ return appManager.getSparkManager(DataStreamsSparkLauncher.NAME);
+ }
+}
diff --git a/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/sink/ConfluentStreamingSinkTest.java b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/sink/ConfluentStreamingSinkTest.java
new file mode 100644
index 0000000..16b33f5
--- /dev/null
+++ b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/sink/ConfluentStreamingSinkTest.java
@@ -0,0 +1,291 @@
+/*
+ * Copyright © 2019 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.plugin.confluent.integration.streaming.sink;
+
+import com.google.common.base.Stopwatch;
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.etl.api.batch.SparkSink;
+import io.cdap.cdap.etl.mock.spark.streaming.MockSource;
+import io.cdap.cdap.etl.proto.v2.ETLPlugin;
+import io.cdap.cdap.proto.ProgramRunStatus;
+import io.cdap.cdap.test.SparkManager;
+import io.cdap.plugin.common.Constants;
+import io.cdap.plugin.confluent.integration.KafkaTestUtils;
+import io.cdap.plugin.confluent.integration.streaming.ConfluentStreamingTestBase;
+import io.cdap.plugin.confluent.integration.streaming.source.ConfluentStreamingSourceTest;
+import io.cdap.plugin.confluent.streaming.sink.ConfluentStreamingSink;
+import io.cdap.plugin.confluent.streaming.sink.ConfluentStreamingSinkConfig;
+import io.cdap.plugin.confluent.streaming.source.ConfluentStreamingSourceConfig;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.assertj.core.api.Assertions;
+import org.assertj.core.api.SoftAssertions;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Tests for Confluent Streaming Sink plugin.
+ */
+public class ConfluentStreamingSinkTest extends ConfluentStreamingTestBase {
+
+ private static Consumer kafkaConsumer;
+ private static Consumer kafkaAvroConsumer;
+
+ @Rule
+ public TestName testName = new TestName();
+
+ private String topic;
+ private SparkManager programManager;
+
+ @BeforeClass
+ public static void setupTestClass() {
+ kafkaConsumer = KafkaTestUtils.createConsumer();
+ kafkaAvroConsumer = KafkaTestUtils.createConsumerForSchemaRegistry();
+ }
+
+ @AfterClass
+ public static void cleanupTestClass() {
+ kafkaConsumer.close();
+ kafkaAvroConsumer.close();
+ }
+
+ @Before
+ public void setUp() {
+ topic = ConfluentStreamingSourceTest.class.getSimpleName() + "_" + testName.getMethodName();
+ KafkaTestUtils.deleteTopic(topic);
+ KafkaTestUtils.createTopic(topic, 2, 3);
+ kafkaConsumer.subscribe(Collections.singletonList(topic));
+ kafkaAvroConsumer.subscribe(Collections.singletonList(topic));
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ KafkaTestUtils.deleteTopic(topic);
+ if (programManager != null) {
+ programManager.stop();
+ programManager.waitForStopped(10, TimeUnit.SECONDS);
+ programManager.waitForRun(ProgramRunStatus.KILLED, 10, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ public void testWritesWithFormat() throws Exception {
+ String keyField = "key";
+ String partitionField = "partition";
+
+ Schema schema = Schema.recordOf(
+ "user",
+ Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
+ Schema.Field.of("first", Schema.of(Schema.Type.STRING)),
+ Schema.Field.of("last", Schema.of(Schema.Type.STRING)),
+ Schema.Field.of(keyField, Schema.nullableOf(Schema.of(Schema.Type.STRING))),
+ Schema.Field.of(partitionField, Schema.of(Schema.Type.INT))
+ );
+ List records = Arrays.asList(
+ StructuredRecord.builder(schema)
+ .set("id", 1L)
+ .set("first", "samuel")
+ .set("last", "jackson")
+ .set(keyField, "a")
+ .set(partitionField, 0)
+ .build(),
+ StructuredRecord.builder(schema)
+ .set("id", 2L)
+ .set("first", "dwayne")
+ .set("last", "johnson")
+ .set(keyField, "b")
+ .set(partitionField, 1)
+ .build(),
+ StructuredRecord.builder(schema)
+ .set("id", 3L)
+ .set("first", "christopher")
+ .set("last", "walken")
+ .set(keyField, "c")
+ .set(partitionField, 1)
+ .build()
+ );
+ ETLPlugin sourcePlugin = MockSource.getPlugin(schema, records);
+ Map properties = getConfigProperties();
+ properties.put(ConfluentStreamingSourceConfig.NAME_FORMAT, "csv");
+ properties.put(ConfluentStreamingSourceConfig.NAME_KEYFIELD, keyField);
+ properties.put(ConfluentStreamingSourceConfig.NAME_PARTITION_FIELD, partitionField);
+ programManager = deploySourcePlugin(sourcePlugin, properties);
+ programManager.startAndWaitForRun(ProgramRunStatus.RUNNING, 10, TimeUnit.SECONDS);
+
+ List> actualRecords = waitForRecordsInKafka(kafkaConsumer, 3);
+ SoftAssertions.assertSoftly(softly -> {
+ ConsumerRecord record1 = findRecordWithKey(actualRecords, "a");
+ softly.assertThat(record1.value()).isEqualTo("1,samuel,jackson");
+ softly.assertThat(record1.partition()).isEqualTo(0);
+ softly.assertThat(record1.offset()).isEqualTo(0);
+
+ ConsumerRecord record2 = findRecordWithKey(actualRecords, "b");
+ softly.assertThat(record2.value()).isEqualTo("2,dwayne,johnson");
+ softly.assertThat(record2.partition()).isEqualTo(1);
+ softly.assertThat(record2.offset()).isEqualTo(0);
+
+ ConsumerRecord record3 = findRecordWithKey(actualRecords, "c");
+ softly.assertThat(record3.value()).isEqualTo("3,christopher,walken");
+ softly.assertThat(record3.partition()).isEqualTo(1);
+ softly.assertThat(record3.offset()).isEqualTo(1);
+ });
+ }
+
+ @Test
+ public void testWritesWithSchemaRegistry() throws Exception {
+ String messageField = "message";
+ String keyField = "key";
+ String partitionField = "partition";
+
+ Schema valueSchema = Schema.recordOf(
+ "user",
+ Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
+ Schema.Field.of("first", Schema.of(Schema.Type.STRING)),
+ Schema.Field.of("last", Schema.of(Schema.Type.STRING))
+ );
+ Schema schema = Schema.recordOf(
+ "confluent",
+ Schema.Field.of(messageField, valueSchema),
+ Schema.Field.of(keyField, Schema.nullableOf(Schema.of(Schema.Type.STRING))),
+ Schema.Field.of(partitionField, Schema.of(Schema.Type.INT))
+ );
+ StructuredRecord value1 = StructuredRecord.builder(valueSchema)
+ .set("id", 1L)
+ .set("first", "samuel")
+ .set("last", "jackson")
+ .build();
+ StructuredRecord value2 = StructuredRecord.builder(valueSchema)
+ .set("id", 2L)
+ .set("first", "dwayne")
+ .set("last", "johnson")
+ .build();
+ StructuredRecord value3 = StructuredRecord.builder(valueSchema)
+ .set("id", 3L)
+ .set("first", "christopher")
+ .set("last", "walken")
+ .build();
+ List records = Arrays.asList(
+ StructuredRecord.builder(schema)
+ .set(messageField, value1)
+ .set(keyField, "a")
+ .set(partitionField, 0)
+ .build(),
+ StructuredRecord.builder(schema)
+ .set(messageField, value2)
+ .set(keyField, "b")
+ .set(partitionField, 1)
+ .build(),
+ StructuredRecord.builder(schema)
+ .set(messageField, value3)
+ .set(keyField, "c")
+ .set(partitionField, 1)
+ .build()
+ );
+ ETLPlugin sourcePlugin = MockSource.getPlugin(schema, records);
+ Map properties = getConfigProperties();
+ properties.put(ConfluentStreamingSourceConfig.NAME_KEYFIELD, keyField);
+ properties.put(ConfluentStreamingSourceConfig.NAME_PARTITION_FIELD, partitionField);
+ properties.put(ConfluentStreamingSourceConfig.NAME_SR_URL, KafkaTestUtils.SR_URL);
+ properties.put(ConfluentStreamingSourceConfig.NAME_SR_API_KEY, KafkaTestUtils.SR_API_KEY);
+ properties.put(ConfluentStreamingSourceConfig.NAME_SR_API_SECRET, KafkaTestUtils.SR_API_SECRET);
+ programManager = deploySourcePlugin(sourcePlugin, properties);
+ programManager.startAndWaitForRun(ProgramRunStatus.RUNNING, 10, TimeUnit.SECONDS);
+
+ GenericRecord expectedValue1 = KafkaTestUtils.toGenericRecord(value1, value1.getSchema());
+ GenericRecord expectedValue2 = KafkaTestUtils.toGenericRecord(value2, value2.getSchema());
+ GenericRecord expectedValue3 = KafkaTestUtils.toGenericRecord(value3, value3.getSchema());
+
+ List> actualRecords = waitForRecordsInKafka(kafkaAvroConsumer, 3);
+ SoftAssertions.assertSoftly(softly -> {
+ ConsumerRecord record1 = findRecordWithKey(actualRecords, "a");
+ softly.assertThat(record1.value()).isEqualTo(expectedValue1);
+ softly.assertThat(record1.partition()).isEqualTo(0);
+ softly.assertThat(record1.offset()).isEqualTo(0);
+
+ ConsumerRecord record2 = findRecordWithKey(actualRecords, "b");
+ softly.assertThat(record2.value()).isEqualTo(expectedValue2);
+ softly.assertThat(record2.partition()).isEqualTo(1);
+ softly.assertThat(record2.offset()).isEqualTo(0);
+
+ ConsumerRecord record3 = findRecordWithKey(actualRecords, "c");
+ softly.assertThat(record3.value()).isEqualTo(expectedValue3);
+ softly.assertThat(record3.partition()).isEqualTo(1);
+ softly.assertThat(record3.offset()).isEqualTo(1);
+ });
+ }
+
+ private ConsumerRecord findRecordWithKey(List> records, Object key) {
+ Optional> recordOptional = records.stream()
+ .filter(record -> Objects.equals(record.key(), key))
+ .findAny();
+ Assertions.assertThat(recordOptional).isPresent();
+ return recordOptional.get();
+ }
+
+ private Map getConfigProperties() {
+ Map properties = new HashMap<>();
+ properties.put(Constants.Reference.REFERENCE_NAME, "confluent");
+ properties.put(ConfluentStreamingSinkConfig.NAME_BROKERS, KafkaTestUtils.KAFKA_SERVER);
+ properties.put(ConfluentStreamingSinkConfig.NAME_TOPIC, topic);
+ properties.put(ConfluentStreamingSinkConfig.NAME_ASYNC, "false");
+ properties.put(ConfluentStreamingSinkConfig.NAME_COMPRESSION_TYPE, "none");
+ properties.put(ConfluentStreamingSinkConfig.NAME_CLUSTER_API_KEY, KafkaTestUtils.CLUSTER_API_KEY);
+ properties.put(ConfluentStreamingSinkConfig.NAME_CLUSTER_API_SECRET, KafkaTestUtils.CLUSTER_API_SECRET);
+ return properties;
+ }
+
+ private SparkManager deploySourcePlugin(ETLPlugin sourcePlugin, Map properties) throws Exception {
+ return deployETL(
+ sourcePlugin,
+ new ETLPlugin(ConfluentStreamingSink.PLUGIN_NAME, SparkSink.PLUGIN_TYPE, properties, null),
+ "KafkaSinkApp"
+ );
+ }
+
+ private List> waitForRecordsInKafka(Consumer consumer, int expectedMessages)
+ throws Exception {
+ List> result = new ArrayList<>();
+ Stopwatch stopwatch = new Stopwatch();
+ while (result.size() < expectedMessages && stopwatch.elapsed(TimeUnit.SECONDS) < 10) {
+ ConsumerRecords records = consumer.poll(Duration.ofMillis(100));
+ for (ConsumerRecord record : records) {
+ result.add(record);
+ }
+ }
+ Assertions.assertThat(result).hasSize(expectedMessages);
+ return result;
+ }
+}
diff --git a/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/source/ConfluentStreamingSourceTest.java b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/source/ConfluentStreamingSourceTest.java
new file mode 100644
index 0000000..fcc3d8f
--- /dev/null
+++ b/confluent-kafka-plugins/src/test/java/io/cdap/plugin/confluent/integration/streaming/source/ConfluentStreamingSourceTest.java
@@ -0,0 +1,453 @@
+/*
+ * Copyright © 2019 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.plugin.confluent.integration.streaming.source;
+
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.api.dataset.table.Table;
+import io.cdap.cdap.etl.api.streaming.StreamingSource;
+import io.cdap.cdap.etl.mock.batch.MockSink;
+import io.cdap.cdap.etl.proto.v2.ETLPlugin;
+import io.cdap.cdap.proto.ProgramRunStatus;
+import io.cdap.cdap.test.DataSetManager;
+import io.cdap.cdap.test.SparkManager;
+import io.cdap.plugin.common.Constants;
+import io.cdap.plugin.confluent.integration.KafkaTestUtils;
+import io.cdap.plugin.confluent.integration.streaming.ConfluentStreamingTestBase;
+import io.cdap.plugin.confluent.streaming.source.ConfluentStreamingSource;
+import io.cdap.plugin.confluent.streaming.source.ConfluentStreamingSourceConfig;
+import io.cdap.plugin.format.avro.StructuredToAvroTransformer;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.requests.ListOffsetRequest;
+import org.assertj.core.api.Assertions;
+import org.assertj.core.api.SoftAssertions;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+
+/**
+ * Tests for Confluent Streaming Source plugin.
+ */
+public class ConfluentStreamingSourceTest extends ConfluentStreamingTestBase {
+
+ private static KafkaProducer kafkaProducer;
+ private static KafkaProducer kafkaAvroProducer;
+
+ @Rule
+ public TestName testName = new TestName();
+
+ private String topic;
+ private String outputTable;
+ private SparkManager programManager;
+
+ @BeforeClass
+ public static void setupTestClass() {
+ kafkaProducer = KafkaTestUtils.createProducer();
+ kafkaAvroProducer = KafkaTestUtils.createProducerForSchemaRegistry();
+ }
+
+ @AfterClass
+ public static void cleanupTestClass() {
+ kafkaProducer.close();
+ kafkaAvroProducer.close();
+ }
+
+ @Before
+ public void setUp() {
+ outputTable = testName.getMethodName() + "_out";
+ topic = ConfluentStreamingSourceTest.class.getSimpleName() + "_" + testName.getMethodName();
+ KafkaTestUtils.deleteTopic(topic);
+ KafkaTestUtils.createTopic(topic, 2, 3);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ KafkaTestUtils.deleteTopic(topic);
+ if (programManager != null) {
+ programManager.stop();
+ programManager.waitForStopped(10, TimeUnit.SECONDS);
+ programManager.waitForRun(ProgramRunStatus.KILLED, 10, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ public void testConfluentStreamingSource() throws Exception {
+ Schema schema = Schema.recordOf(
+ "user",
+ Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
+ Schema.Field.of("first", Schema.of(Schema.Type.STRING)),
+ Schema.Field.of("last", Schema.of(Schema.Type.STRING))
+ );
+ Map properties = getConfigProperties(schema);
+ properties.put(ConfluentStreamingSourceConfig.NAME_FORMAT, "csv");
+ programManager = deploySourcePlugin(properties);
+ programManager.startAndWaitForRun(ProgramRunStatus.RUNNING, 10, TimeUnit.SECONDS);
+
+ sendKafkaMessage(topic, 0, "a", "1,samuel,jackson");
+
+ List expectedRecords = Collections.singletonList(
+ StructuredRecord.builder(schema)
+ .set("id", 1L)
+ .set("first", "samuel")
+ .set("last", "jackson")
+ .build()
+ );
+ waitForRecords(outputTable, expectedRecords);
+
+ programManager.stop();
+ programManager.waitForStopped(10, TimeUnit.SECONDS);
+
+ // clear the output table
+ DataSetManager outputManager = getDataset(outputTable);
+ MockSink.clear(outputManager);
+
+ // now write some more messages to kafka and start the program again to make sure it picks up where it left off
+ sendKafkaMessage(topic, 1, "b", "2,dwayne,johnson");
+
+ programManager.startAndWaitForRun(ProgramRunStatus.RUNNING, 10, TimeUnit.SECONDS);
+
+ List expectedRecords2 = Collections.singletonList(
+ StructuredRecord.builder(schema)
+ .set("id", 2L)
+ .set("first", "dwayne")
+ .set("last", "johnson")
+ .build()
+ );
+ waitForRecords(outputTable, expectedRecords2);
+ }
+
+ @Test
+ public void testConfluentStreamingSourceAdditionalFieldsWithoutFormat() throws Exception {
+ String keyField = "key";
+ String timeField = "time";
+ String partitionField = "partition";
+ String offsetField = "offset";
+ Schema schema = Schema.recordOf(
+ "user",
+ Schema.Field.of("message", Schema.of(Schema.Type.BYTES)),
+ Schema.Field.of(keyField, Schema.of(Schema.Type.BYTES)),
+ Schema.Field.of(timeField, Schema.of(Schema.Type.LONG)),
+ Schema.Field.of(partitionField, Schema.of(Schema.Type.INT)),
+ Schema.Field.of(offsetField, Schema.of(Schema.Type.LONG))
+ );
+ Map properties = getConfigProperties(schema);
+ properties.put(ConfluentStreamingSourceConfig.NAME_KEYFIELD, keyField);
+ properties.put(ConfluentStreamingSourceConfig.NAME_TIMEFIELD, timeField);
+ properties.put(ConfluentStreamingSourceConfig.NAME_PARTITION_FIELD, partitionField);
+ properties.put(ConfluentStreamingSourceConfig.NAME_OFFSET_FIELD, offsetField);
+ programManager = deploySourcePlugin(properties);
+ programManager.startAndWaitForRun(ProgramRunStatus.RUNNING, 10, TimeUnit.SECONDS);
+
+ sendKafkaMessage(topic, 1, "a", "payload_1");
+ sendKafkaMessage(topic, 0, "b", "payload_2");
+ sendKafkaMessage(topic, 0, "c", "payload_3");
+
+ List actualRecords = waitForRecords(outputTable, 3);
+
+ SoftAssertions.assertSoftly(softly -> {
+ StructuredRecord output1 =
+ findRecordWithField(actualRecords, keyField, ByteBuffer.wrap("a".getBytes(StandardCharsets.UTF_8)));
+ softly.assertThat(output1.