content: /*
   * 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.kafka.streams.tests;
  
  import org.apache.kafka.clients.consumer.ConsumerConfig;
  import org.apache.kafka.clients.consumer.ConsumerRecord;
  import org.apache.kafka.clients.consumer.ConsumerRecords;
  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.clients.producer.ProducerRecord;
  import org.apache.kafka.common.PartitionInfo;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.serialization.ByteArrayDeserializer;
  import org.apache.kafka.common.serialization.ByteArraySerializer;
  import org.apache.kafka.common.serialization.Deserializer;
  import org.apache.kafka.common.serialization.IntegerDeserializer;
  import org.apache.kafka.common.serialization.IntegerSerializer;
  import org.apache.kafka.common.serialization.Serde;
  import org.apache.kafka.common.serialization.Serdes;
  import org.apache.kafka.common.serialization.Serializer;
  import org.apache.kafka.streams.KafkaStreams;
  import org.apache.kafka.streams.KeyValue;
  import org.apache.kafka.streams.StreamsBuilder;
  import org.apache.kafka.streams.StreamsConfig;
  import org.apache.kafka.streams.Topology;
  import org.apache.kafka.streams.kstream.Consumed;
  import org.apache.kafka.streams.kstream.Grouped;
  import org.apache.kafka.streams.kstream.KTable;
  import org.apache.kafka.streams.kstream.Materialized;
  import org.apache.kafka.streams.kstream.ValueJoiner;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.nio.ByteBuffer;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.Arrays;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  import java.util.Random;
  import java.util.TreeMap;
  import java.util.UUID;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.stream.Stream;
  
  import static java.util.stream.Collectors.toList;
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.apache.kafka.common.utils.Utils.mkProperties;
  
  /**
   * This test builds on a basic relational data caricature:
   * a set of articles, each with a set of comments on them.
   * <p>
   * The particular relation in this data is that each Comment
   * has an articleId foreign-key reference to an Article.
   * <p>
   * We traverse this relation in both directions to verify
   * correct operation of Kafka Streams:
   * - Embed on each Comment a prefix of the text in the Article it references
   * aka SELECT * FROM Comment JOIN Article ON Comment.ArticleID = Article.ID
   * - Embed on each Article a count of all the Comments it has
   */
  public class RelationalSmokeTest extends SmokeTestUtil {
      private static final Logger LOG = LoggerFactory.getLogger(RelationalSmokeTest.class);
  
      static final String ARTICLE_SOURCE = "in-article";
      static final String COMMENT_SOURCE = "in-comment";
      static final String ARTICLE_RESULT_SINK = "out-augmented-article";
      static final String COMMENT_RESULT_SINK = "out-augmented-comment";
      private static final String[] TOPICS = {
          ARTICLE_SOURCE,
          COMMENT_SOURCE,
          ARTICLE_RESULT_SINK,
          COMMENT_RESULT_SINK
      };
  
      public static String[] topics() {
          return Arrays.copyOf(TOPICS, TOPICS.length);
      }
  
      public static class Article {
          private final int key;
          private final long timestamp;
          private final String text;
  
          private Article(final int key, final long timestamp, final String text) {
              this.key = key;
              this.timestamp = timestamp;
              this.text = text;
          }
  
          public int getKey() {
              return key;
          }
  
          public long getTimestamp() {
              return timestamp;
          }
  
          public String getText() {
              return text;
          }
  
          @Override
          public String toString() {
              return "Article{" +
                  "key=" + key +
                  ", timestamp=" + Instant.ofEpochMilli(timestamp) +
                  ", text='" + text + '\'' +
                  '}';
          }
  
          public static class ArticleSerializer implements Serializer<Article> {
              @Override
              public byte[] serialize(final String topic, final Article data) {
                  final byte[] serialText = stringSerde.serializer().serialize(topic, data.getText());
  
                  final int length = Integer.BYTES + Long.BYTES + Integer.BYTES + serialText.length;
  
                  final ByteBuffer buffer =
                      ByteBuffer.allocate(length)
                          .putInt(data.getKey())
                          .putLong(data.getTimestamp())
                          .putInt(serialText.length)
                          .put(serialText);
  
                  return Serdes.ByteBuffer().serializer().serialize(topic, buffer);
              }
          }
  
          public static class ArticleDeserializer implements Deserializer<Article> {
  
              public static Article deserialize(final String topic, final ByteBuffer buffer) {
                  final int key = buffer.getInt();
                  final long timestamp = buffer.getLong();
                  final int textLength = buffer.getInt();
                  final byte[] serialText = new byte[textLength];
                  buffer.get(serialText);
                  final String text = stringSerde.deserializer().deserialize(topic, serialText);
                  return new Article(key, timestamp, text);
              }
  
              @Override
              public Article deserialize(final String topic, final byte[] data) {
                  final ByteBuffer buffer = Serdes.ByteBuffer().deserializer().deserialize(topic, data);
                  return deserialize(topic, buffer);
              }
          }
  
          public static class ArticleSerde implements Serde<Article> {
              @Override
              public Serializer<Article> serializer() {
                  return new ArticleSerializer();
              }
  
              @Override
              public Deserializer<Article> deserializer() {
                  return new ArticleDeserializer();
              }
          }
      }
  
      public static class Comment {
          private final int key;
          private final long timestamp;
          private final String text;
          private final int articleId;
  
          private Comment(final int key, final long timestamp, final String text, final int articleId) {
              this.key = key;
              this.timestamp = timestamp;
              this.text = text;
              this.articleId = articleId;
          }
  
          public int getKey() {
              return key;
          }
  
          public long getTimestamp() {
              return timestamp;
          }
  
          public String getText() {
              return text;
          }
  
          public int getArticleId() {
              return articleId;
          }
  
          @Override
          public String toString() {
              return "Comment{" +
                  "key=" + key +
                  ", timestamp=" + Instant.ofEpochMilli(timestamp) +
                  ", text='" + text + '\'' +
                  ", articleId=" + articleId +
                  '}';
          }
  
          public static class CommentSerializer implements Serializer<Comment> {
  
              @Override
              public byte[] serialize(final String topic, final Comment data) {
                  final byte[] serialText = stringSerde.serializer().serialize(topic, data.text);
  
                  final int length = Integer.BYTES + Long.BYTES + (Integer.BYTES + serialText.length) + Integer.BYTES;
  
                  final ByteBuffer buffer =
                      ByteBuffer.allocate(length)
                          .putInt(data.key)
                          .putLong(data.timestamp)
                          .putInt(serialText.length)
                          .put(serialText)
                          .putInt(data.articleId);
  
                  return Serdes.ByteBuffer().serializer().serialize(topic, buffer);
              }
          }
  
          public static class CommentDeserializer implements Deserializer<Comment> {
  
              public static Comment deserialize(final String topic, final ByteBuffer buffer) {
                  final int key = buffer.getInt();
                  final long timestamp = buffer.getLong();
                  final int textLength = buffer.getInt();
                  final byte[] textBytes = new byte[textLength];
                  buffer.get(textBytes);
                  final String text = stringSerde.deserializer().deserialize(topic, textBytes);
                  final int articleId = buffer.getInt();
  
                  return new Comment(key, timestamp, text, articleId);
              }
  
              @Override
              public Comment deserialize(final String topic, final byte[] data) {
                  final ByteBuffer buffer = Serdes.ByteBuffer().deserializer().deserialize(topic, data);
                  return deserialize(topic, buffer);
              }
          }
  
          public static class CommentSerde implements Serde<Comment> {
  
              @Override
              public Serializer<Comment> serializer() {
                  return new CommentSerializer();
              }
  
              @Override
              public Deserializer<Comment> deserializer() {
                  return new CommentDeserializer();
              }
          }
      }
  
      public static final class DataSet {
          private final Article[] articles;
          private final Comment[] comments;
  
          private DataSet(final Article[] articles, final Comment[] comments) {
              this.articles = articles;
              this.comments = comments;
          }
  
          public Article[] getArticles() {
              return articles;
          }
  
          public Comment[] getComments() {
              return comments;
          }
  
          @Override
          public String toString() {
              final StringBuilder stringBuilder = new StringBuilder();
              stringBuilder.append(articles.length).append(" Articles").append('\n');
              for (final Article article : articles) {
                  stringBuilder.append("  ").append(article).append('\n');
              }
              stringBuilder.append(comments.length).append(" Comments").append("\n");
              for (final Comment comment : comments) {
                  stringBuilder.append("  ").append(comment).append('\n');
              }
              return stringBuilder.toString();
          }
  
          public static DataSet generate(final int numArticles, final int numComments) {
              // generate four days' worth of data, starting 4 days in the past (avoiding future records)
              final int timeSpan = 1000 * 60 * 60 * 24 * 4;
              final long dataStartTime = System.currentTimeMillis() - timeSpan;
              final long dataEndTime = System.currentTimeMillis();
  
              // Explicitly create a seed so we can log.
              // If we are debugging a failed run, we can deterministically produce the same dataset
              // by plugging in the seed from that run.
              final long seed = new Random().nextLong();
              final Random random = new Random(seed);
              LOG.info("Dataset PRNG seed: {}", seed);
              final Iterator<Integer> articlesToCommentOnSequence = zipfNormal(random, numArticles);
  
              final Article[] articles = new Article[numArticles];
              final Comment[] comments = new Comment[numComments];
  
              // first generate the articles (note: out of order)
              for (int i = 0; i < numArticles; i++) {
                  final long timestamp = random.nextInt(timeSpan) + dataStartTime;
                  final String text = randomText(random, 2000);
                  articles[i] = new Article(i, timestamp, text);
              }
  
              // then spend the rest of the time generating the comments
              for (int i = 0; i < numComments; i++) {
                  final int articleId = articlesToCommentOnSequence.next();
                  final long articleTimestamp = articles[articleId].getTimestamp();
                  // comments get written after articles
                  final long timestamp = random.nextInt((int) (dataEndTime - articleTimestamp)) + articleTimestamp;
                  final String text = randomText(random, 200);
                  final Comment comment = new Comment(i, timestamp, text, articleId);
                  comments[i] = comment;
              }
              return new DataSet(articles, comments);
          }
  
          /**
           * Rough-and-ready random text generator. Creates a text with a length normally
           * distributed about {@code avgLength} with a standard deviation of 1/3rd {@code avgLength}.
           * Each letter is drawn uniformly from a-z.
           */
          private static String randomText(final Random random, final int avgLength) {
              final int lowChar = 97; // letter 'a'
              final int highChar = 122; // letter 'z'
  
              final int length = Math.max(0, (int) (random.nextGaussian() * avgLength / 3.0) + avgLength);
              final char[] chars = new char[length];
              for (int i = 0; i < chars.length; i++) {
                  chars[i] = (char) (random.nextInt(highChar - lowChar) + lowChar);
              }
              return new String(chars);
          }
  
          /**
           * Generates a keySpace number of unique keys normally distributed,
           * with the mean at 0 and stdDev 1/3 of the way through the keySpace
           * any sample more than 3 standard deviations from the mean are assigned to the 0 key.
           * <p>
           * This is designed to roughly balance the key properties of two dominant real-world
           * data distribution: Zipfian and Normal, while also being efficient to generate.
           */
          private static Iterator<Integer> zipfNormal(final Random random, final int keySpace) {
              return new Iterator<>() {
                  @Override
                  public boolean hasNext() {
                      return true;
                  }
  
                  @Override
                  public Integer next() {
                      final double gaussian = Math.abs(random.nextGaussian());
                      final double scaled = gaussian / 3.0; //
                      final double sample = scaled > 1.0 ? 0.0 : scaled;
                      final double keyDouble = sample * keySpace;
                      return (int) keyDouble;
                  }
              };
          }
  
          public void produce(final String kafka, final Duration timeToSpend) throws InterruptedException {
              final Properties producerProps = new Properties();
              final String id = "RelationalSmokeTestProducer" + UUID.randomUUID();
              producerProps.put(ProducerConfig.CLIENT_ID_CONFIG, id);
              producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, kafka);
              producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, IntegerSerializer.class);
              producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class);
              producerProps.put(ProducerConfig.ACKS_CONFIG, "-1");
  
              final Article.ArticleSerializer articleSerializer = new Article.ArticleSerializer();
              final Comment.CommentSerializer commentSerializer = new Comment.CommentSerializer();
  
              final long pauseTime = timeToSpend.toMillis() / (articles.length + comments.length);
  
              try (final KafkaProducer<Integer, byte[]> producer = new KafkaProducer<>(producerProps)) {
                  int a = 0;
                  int c = 0;
                  while (a < articles.length || c < comments.length) {
                      final ProducerRecord<Integer, byte[]> producerRecord;
                      if (a < articles.length && c >= comments.length ||
                          a < articles.length && articles[a].getTimestamp() <= comments[c].timestamp) {
                          producerRecord = new ProducerRecord<>(
                              ARTICLE_SOURCE,
                              null,
                              articles[a].getTimestamp(),
                              articles[a].getKey(),
                              articleSerializer.serialize("", articles[a])
                          );
                          a++;
                      } else {
                          producerRecord = new ProducerRecord<>(
                              COMMENT_SOURCE,
                              null,
                              comments[c].getTimestamp(),
                              comments[c].getKey(),
                              commentSerializer.serialize("", comments[c])
                          );
                          c++;
                      }
                      producer.send(producerRecord);
                      producer.flush();
                      LOG.info("sent {} {}", producerRecord.topic(), producerRecord.key());
                      Thread.sleep(pauseTime);
                  }
              }
          }
      }
  
      public static final class AugmentedArticle extends Article {
          private final long commentCount;
  
          private AugmentedArticle(final int key, final long timestamp, final String text, final long commentCount) {
              super(key, timestamp, text);
              this.commentCount = commentCount;
          }
  
          public long getCommentCount() {
              return commentCount;
          }
  
          @Override
          public String toString() {
              return "AugmentedArticle{" +
                  "key=" + super.key +
                  ", timestamp=" + getTimestamp() +
                  ", text='" + getText() + '\'' +
                  ", commentCount=" + commentCount +
                  '}';
          }
  
          public static class AugmentedArticleSerializer implements Serializer<AugmentedArticle> {
              private final ArticleSerializer articleSerializer = new ArticleSerializer();
  
              @Override
              public byte[] serialize(final String topic, final AugmentedArticle data) {
                  final byte[] serializedArticle = articleSerializer.serialize(topic, data);
                  final int length = serializedArticle.length + Long.BYTES;
                  final ByteBuffer buffer =
                      ByteBuffer.allocate(length)
                          .put(serializedArticle)
                          .putLong(data.getCommentCount());
                  return Serdes.ByteBuffer().serializer().serialize(topic, buffer);
              }
          }
  
          public static class AugmentedArticleDeserializer implements Deserializer<AugmentedArticle> {
              @Override
              public AugmentedArticle deserialize(final String topic, final byte[] data) {
                  final ByteBuffer wrap = ByteBuffer.wrap(data);
  
                  final Article article = ArticleDeserializer.deserialize(topic, wrap);
                  final long commentCount = wrap.getLong();
                  return new AugmentedArticle(article.key, article.getTimestamp(), article.getText(), commentCount);
              }
          }
  
          public static class AugmentedArticleSerde implements Serde<AugmentedArticle> {
  
              @Override
              public Serializer<AugmentedArticle> serializer() {
                  return new AugmentedArticleSerializer();
              }
  
              @Override
              public Deserializer<AugmentedArticle> deserializer() {
                  return new AugmentedArticleDeserializer();
              }
          }
  
          public static ValueJoiner<Article, Long, AugmentedArticle> joiner() {
              return (article, commentCount) -> new AugmentedArticle(
                  article.getKey(),
                  article.getTimestamp(),
                  article.getText(), commentCount == null ? 0 : commentCount
              );
          }
      }
  
      public static final class AugmentedComment extends Comment {
          private final String articlePrefix;
  
          private AugmentedComment(final int key,
                                   final long timestamp,
                                   final String text,
                                   final int articleId,
                                   final String articlePrefix) {
              super(key, timestamp, text, articleId);
              this.articlePrefix = articlePrefix;
          }
  
          public String getArticlePrefix() {
              return articlePrefix;
          }
  
          @Override
          public String toString() {
              return "AugmentedComment{" +
                  "key=" + super.key +
                  ", timestamp=" + getTimestamp() +
                  ", text='" + getText() + '\'' +
                  ", articleId=" + getArticleId() +
                  ", articlePrefix='" + articlePrefix + '\'' +
                  '}';
          }
  
          public static class AugmentedCommentSerializer implements Serializer<AugmentedComment> {
              private final CommentSerializer commentSerializer = new CommentSerializer();
  
              @Override
              public byte[] serialize(final String topic, final AugmentedComment data) {
                  final byte[] serializedComment = commentSerializer.serialize(topic, data);
                  final byte[] serializedPrefix = stringSerde.serializer().serialize(topic, data.getArticlePrefix());
                  final int length = serializedComment.length + Integer.BYTES + serializedPrefix.length;
                  final ByteBuffer buffer =
                      ByteBuffer.allocate(length)
                          .put(serializedComment)
                          .putInt(serializedPrefix.length)
                          .put(serializedPrefix);
                  return Serdes.ByteBuffer().serializer().serialize(topic, buffer);
              }
          }
  
          public static class AugmentedCommentDeserializer implements Deserializer<AugmentedComment> {
              @Override
              public AugmentedComment deserialize(final String topic, final byte[] data) {
                  final ByteBuffer wrap = ByteBuffer.wrap(data);
  
                  final Comment comment = CommentDeserializer.deserialize(topic, wrap);
                  final int prefixLength = wrap.getInt();
                  final byte[] serializedPrefix = new byte[prefixLength];
                  wrap.get(serializedPrefix);
                  final String articlePrefix = stringSerde.deserializer().deserialize(topic, serializedPrefix);
                  return new AugmentedComment(
                      comment.key,
                      comment.getTimestamp(),
                      comment.getText(),
                      comment.getArticleId(),
                      articlePrefix
                  );
              }
          }
  
          public static class AugmentedCommentSerde implements Serde<AugmentedComment> {
  
              @Override
              public Serializer<AugmentedComment> serializer() {
                  return new AugmentedCommentSerializer();
              }
  
              @Override
              public Deserializer<AugmentedComment> deserializer() {
                  return new AugmentedCommentDeserializer();
              }
          }
  
          private static String prefix(final String text, final int length) {
              return text.length() < length ? text : text.substring(0, length);
          }
  
          public static ValueJoiner<Comment, Article, AugmentedComment> joiner() {
              return (comment, article) -> new AugmentedComment(
                  comment.key,
                  comment.getTimestamp(),
                  comment.getText(),
                  comment.getArticleId(),
                  prefix(article.getText(), 10)
              );
          }
      }
  
      public static final class App {
          public static Topology getTopology() {
              final StreamsBuilder streamsBuilder = new StreamsBuilder();
              final KTable<Integer, Article> articles =
                  streamsBuilder.table(ARTICLE_SOURCE, Consumed.with(intSerde, new Article.ArticleSerde()));
  
              final KTable<Integer, Comment> comments =
                  streamsBuilder.table(COMMENT_SOURCE, Consumed.with(intSerde, new Comment.CommentSerde()));
  
  
              final KTable<Integer, Long> commentCounts =
                  comments.groupBy(
                      (key, value) -> new KeyValue<>(value.getArticleId(), (short) 1),
                      Grouped.with(Serdes.Integer(), Serdes.Short())
                  )
                  .count();
  
              articles
                  .leftJoin(
                      commentCounts,
                      AugmentedArticle.joiner(),
                      Materialized.with(null, new AugmentedArticle.AugmentedArticleSerde())
                  )
                  .toStream()
                  .to(ARTICLE_RESULT_SINK);
  
              comments.join(articles,
                            Comment::getArticleId,
                            AugmentedComment.joiner(),
                            Materialized.with(null, new AugmentedComment.AugmentedCommentSerde()))
                  .toStream()
                  .to(COMMENT_RESULT_SINK);
  
              return streamsBuilder.build();
          }
  
          public static Properties getConfig(final String broker,
                                             final String application,
                                             final String id,
                                             final String processingGuarantee,
                                             final String groupProtocol,
                                             final String stateDir) {
              final Properties properties =
                  mkProperties(
                      mkMap(
                          mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, broker),
                          mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, application),
                          mkEntry(StreamsConfig.CLIENT_ID_CONFIG, id),
                          mkEntry(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, processingGuarantee),
                          mkEntry(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"),
                          mkEntry(StreamsConfig.STATE_DIR_CONFIG, stateDir),
                          mkEntry(StreamsConfig.GROUP_PROTOCOL_CONFIG, groupProtocol)
                      )
                  );
              properties.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000L);
              properties.put(StreamsConfig.WINDOW_STORE_CHANGE_LOG_ADDITIONAL_RETENTION_MS_CONFIG, Duration.ofDays(5).toMillis());
              return properties;
          }
  
          public static KafkaStreams startSync(final String broker,
                                               final String application,
                                               final String id,
                                               final String processingGuarantee,
                                               final String groupProtocol,
                                               final String stateDir) throws InterruptedException {
              final KafkaStreams kafkaStreams =
                  new KafkaStreams(getTopology(), getConfig(broker, application, id, processingGuarantee, groupProtocol, stateDir));
              final CountDownLatch startUpLatch = new CountDownLatch(1);
              kafkaStreams.setStateListener((newState, oldState) -> {
                  if (oldState == KafkaStreams.State.REBALANCING && newState == KafkaStreams.State.RUNNING) {
                      startUpLatch.countDown();
                  }
              });
              kafkaStreams.start();
              startUpLatch.await();
              LOG.info("Streams has started.");
              return kafkaStreams;
          }
  
          public static boolean verifySync(final String broker, final Instant deadline) throws InterruptedException {
              final Deserializer<Article> articleDeserializer = new Article.ArticleDeserializer();
  
              final Deserializer<AugmentedArticle> augmentedArticleDeserializer =
                  new AugmentedArticle.AugmentedArticleDeserializer();
  
              final Deserializer<Comment> commentDeserializer = new Comment.CommentDeserializer();
  
              final Deserializer<AugmentedComment> augmentedCommentDeserializer =
                  new AugmentedComment.AugmentedCommentDeserializer();
  
  
              final Properties consumerProperties = new Properties();
              final String id = "RelationalSmokeTestConsumer" + UUID.randomUUID();
              consumerProperties.put(ConsumerConfig.CLIENT_ID_CONFIG, id);
              consumerProperties.put(ConsumerConfig.GROUP_ID_CONFIG, id);
              consumerProperties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, broker);
              consumerProperties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, IntegerDeserializer.class);
              consumerProperties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
              consumerProperties.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
              consumerProperties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
  
              try (final KafkaConsumer<Integer, byte[]> consumer = new KafkaConsumer<>(consumerProperties)) {
                  final List<PartitionInfo> articlePartitions = consumer.partitionsFor(ARTICLE_SOURCE);
                  final List<PartitionInfo> augmentedArticlePartitions = consumer.partitionsFor(ARTICLE_RESULT_SINK);
                  final List<PartitionInfo> commentPartitions = consumer.partitionsFor(COMMENT_SOURCE);
                  final List<PartitionInfo> augmentedCommentPartitions = consumer.partitionsFor(COMMENT_RESULT_SINK);
                  final List<TopicPartition> assignment =
                      Stream.concat(
                          Stream.concat(
                              articlePartitions.stream().map(p -> new TopicPartition(p.topic(), p.partition())),
                              augmentedArticlePartitions.stream().map(p -> new TopicPartition(p.topic(), p.partition()))
                          ), 
                          Stream.concat(
                              commentPartitions.stream().map(p -> new TopicPartition(p.topic(), p.partition())),
                              augmentedCommentPartitions.stream().map(p -> new TopicPartition(p.topic(), p.partition()))
                          )
                      ).collect(toList());
                  consumer.assign(assignment);
                  consumer.seekToBeginning(assignment);
  
                  final Map<Integer, Article> consumedArticles = new TreeMap<>();
                  final Map<Integer, AugmentedArticle> consumedAugmentedArticles = new TreeMap<>();
                  final Map<Integer, Comment> consumedComments = new TreeMap<>();
                  final Map<Integer, AugmentedComment> consumedAugmentedComments = new TreeMap<>();
  
                  boolean printedConsumedArticle = false;
                  boolean printedConsumedAugmentedArticle = false;
                  boolean printedConsumedComment = false;
                  boolean printedConsumedAugmentedComment = false;
                  boolean passed = false;
  
                  while (!passed && Instant.now().isBefore(deadline)) {
                      boolean lastPollWasEmpty = false;
                      while (!lastPollWasEmpty) {
                          final ConsumerRecords<Integer, byte[]> poll = consumer.poll(Duration.ofSeconds(1));
                          lastPollWasEmpty = poll.isEmpty();
                          for (final ConsumerRecord<Integer, byte[]> record : poll) {
                              final Integer key = record.key();
                              switch (record.topic()) {
                                  case ARTICLE_SOURCE: {
                                      final Article article = articleDeserializer.deserialize("", record.value());
                                      if (consumedArticles.containsKey(key)) {
                                          LOG.warn("Duplicate article: {} and {}", consumedArticles.get(key), article);
                                      }
                                      consumedArticles.put(key, article);
                                      break;
                                  }
                                  case COMMENT_SOURCE: {
                                      final Comment comment = commentDeserializer.deserialize("", record.value());
                                      if (consumedComments.containsKey(key)) {
                                          LOG.warn("Duplicate comment: {} and {}", consumedComments.get(key), comment);
                                      }
                                      consumedComments.put(key, comment);
                                      break;
                                  }
                                  case ARTICLE_RESULT_SINK: {
                                      final AugmentedArticle article =
                                          augmentedArticleDeserializer.deserialize("", record.value());
                                      consumedAugmentedArticles.put(key, article);
                                      break;
                                  }
                                  case COMMENT_RESULT_SINK: {
                                      final AugmentedComment comment =
                                          augmentedCommentDeserializer.deserialize("", record.value());
                                      consumedAugmentedComments.put(key, comment);
                                      break;
                                  }
                                  default:
                                      throw new IllegalArgumentException(record.toString());
                              }
                          }
                          consumer.commitSync();
                      }
  
                      if (!printedConsumedArticle && !consumedArticles.isEmpty()) {
                          LOG.info("Consumed first Article");
                          printedConsumedArticle = true;
                      }
                      if (!printedConsumedComment && !consumedComments.isEmpty()) {
                          LOG.info("Consumed first Comment");
                          printedConsumedComment = true;
                      }
                      if (!printedConsumedAugmentedArticle && !consumedAugmentedArticles.isEmpty()) {
                          LOG.info("Consumed first AugmentedArticle");
                          printedConsumedAugmentedArticle = true;
                      }
                      if (!printedConsumedAugmentedComment && !consumedAugmentedComments.isEmpty()) {
                          LOG.info("Consumed first AugmentedComment");
                          printedConsumedAugmentedComment = true;
                      }
  
                      passed = verifySync(
                          false,
                          consumedArticles,
                          consumedComments,
                          consumedAugmentedArticles,
                          consumedAugmentedComments
                      );
                      if (!passed) {
                          LOG.info("Verification has not passed yet. ");
                          Thread.sleep(500);
                      }
                  }
                  return verifySync(
                      true,
                      consumedArticles,
                      consumedComments,
                      consumedAugmentedArticles,
                      consumedAugmentedComments
                  );
              }
          }
  
          public static void assertThat(final AtomicBoolean pass,
                                        final StringBuilder failures,
                                        final String message,
                                        final boolean passed) {
              if (!passed) {
                  if (failures != null) {
                      failures.append("\n").append(message);
                  }
                  pass.set(false);
              }
          }
  
          static boolean verifySync(final boolean logResults,
                                    final Map<Integer, Article> consumedArticles,
                                    final Map<Integer, Comment> consumedComments,
                                    final Map<Integer, AugmentedArticle> consumedAugmentedArticles,
                                    final Map<Integer, AugmentedComment> consumedAugmentedComments) {
              final AtomicBoolean pass = new AtomicBoolean(true);
              final StringBuilder report = logResults ? new StringBuilder() : null;
  
              assertThat(
                  pass,
                  report,
                  "Expected 1 article, got " + consumedArticles.size(),
                  !consumedArticles.isEmpty()
              );
              assertThat(
                  pass,
                  report,
                  "Expected 1 comment, got " + consumedComments.size(),
                  !consumedComments.isEmpty()
              );
  
              assertThat(
                  pass,
                  report,
                  "Mismatched article size between augmented articles (size "
                      + consumedAugmentedArticles.size() +
                      ") and consumed articles (size "
                      + consumedArticles.size() + ")",
                  consumedAugmentedArticles.size() == consumedArticles.size()
              );
              assertThat(
                  pass,
                  report,
                  "Mismatched comments size between augmented comments (size "
                      + consumedAugmentedComments.size() +
                      ") and consumed comments (size " +
                      consumedComments.size() + ")", 
                  consumedAugmentedComments.size() == consumedComments.size()
              );
  
              final Map<Integer, Long> commentCounts = new TreeMap<>();
  
              for (final RelationalSmokeTest.AugmentedComment augmentedComment : consumedAugmentedComments.values()) {
                  final int key = augmentedComment.getKey();
                  assertThat(
                      pass,
                      report,
                      "comment missing, but found in augmentedComment: " + key,
                      consumedComments.containsKey(key)
                  );
  
                  final Comment comment = consumedComments.get(key);
                  if (comment != null) {
                      assertThat(
                          pass,
                          report,
                          "comment missing, but found in augmentedComment: " + key,
                          consumedComments.containsKey(key)
                      );
                  }
                  commentCounts.put(
                      augmentedComment.getArticleId(),
                      commentCounts.getOrDefault(augmentedComment.getArticleId(), 0L) + 1
                  );
  
                  assertThat(
                      pass,
                      report,
                      "augmentedArticle [" + augmentedComment.getArticleId() + "] " +
                          "missing for augmentedComment [" + augmentedComment.getKey() + "]",
                      consumedAugmentedArticles.containsKey(augmentedComment.getArticleId())
                  );
                  final AugmentedArticle augmentedArticle =
                          consumedAugmentedArticles.get(augmentedComment.getArticleId());
                  if (augmentedArticle != null) {
                      assertThat(
                          pass,
                          report,
                          "articlePrefix didn't match augmentedArticle: " + augmentedArticle.getText(),
                          augmentedArticle.getText().startsWith(augmentedComment.getArticlePrefix())
                      );
                  }
  
                  assertThat(
                      pass,
                      report,
                      "article " + augmentedComment.getArticleId() + " missing from consumedArticles",
                      consumedArticles.containsKey(augmentedComment.getArticleId())
                  );
                  final Article article = consumedArticles.get(augmentedComment.getArticleId());
                  if (article != null) {
                      assertThat(
                          pass,
                          report,
                          "articlePrefix didn't match article: " + article.getText(),
                          article.getText().startsWith(augmentedComment.getArticlePrefix())
                      );
                  }
              }
  
  
              for (final RelationalSmokeTest.AugmentedArticle augmentedArticle : consumedAugmentedArticles.values()) {
                  assertThat(
                      pass,
                      report,
                      "article " + augmentedArticle.getKey() + " comment count mismatch",
                      augmentedArticle.getCommentCount() == commentCounts.getOrDefault(augmentedArticle.getKey(), 0L)
                  );
              }
  
              if (logResults) {
                  if (pass.get()) {
                      LOG.info(
                          "Evaluation passed ({}/{}) articles and ({}/{}) comments",
                          consumedAugmentedArticles.size(),
                          consumedArticles.size(),
                          consumedAugmentedComments.size(),
                          consumedComments.size()
                      );
                  } else {
                      LOG.error(
                          "Evaluation failed\nReport: {}\n" +
                              "Consumed Input Articles: {}\n" +
                              "Consumed Input Comments: {}\n" +
                              "Consumed Augmented Articles: {}\n" +
                              "Consumed Augmented Comments: {}",
                          report,
                          consumedArticles,
                          consumedComments,
                          consumedAugmentedArticles,
                          consumedAugmentedComments
                      );
                  }
              }
  
              return pass.get();
          }
      }
  
      /*
       * Used by the smoke tests.
       */
      public static void main(final String[] args) {
          System.out.println(Arrays.toString(args));
          final String mode = args[0];
          final String kafka = args[1];
  
          try {
              switch (mode) {
                  case "driver": {
                      // this starts the driver (data generation and result verification)
                      final int numArticles = 1_000;
                      final int numComments = 10_000;
                      final DataSet dataSet = DataSet.generate(numArticles, numComments);
                      // publish the data for at least one minute
                      dataSet.produce(kafka, Duration.ofMinutes(1));
                      LOG.info("Smoke test finished producing");
                      // let it soak in
                      Thread.sleep(1000);
                      LOG.info("Smoke test starting verification");
                      // wait for at most 10 minutes to get a passing result
                      final boolean pass = App.verifySync(kafka, Instant.now().plus(Duration.ofMinutes(10)));
                      if (pass) {
                          LOG.info("Smoke test complete: passed");
                      } else {
                          LOG.error("Smoke test complete: failed");
                      }
                      break;
                  }
                  case "application": {
                      final String nodeId = args[2];
                      final String processingGuarantee = args[3];
                      final String groupProtocol = args[4];
                      final String stateDir = args[5];
                      App.startSync(kafka, UUID.randomUUID().toString(), nodeId, processingGuarantee, groupProtocol, stateDir);
                      break;
                  }
                  default:
                      LOG.error("Unknown command: {}", mode);
                      throw new RuntimeException("Unknown command: " + mode);
              }
          } catch (final InterruptedException e) {
              LOG.error("Interrupted", e);
          }
      }
  }
  