Java Code Examples for org.apache.kafka.clients.consumer.OffsetResetStrategy#EARLIEST

@Override
public Consumer<String, ByteBuffer> createConsumer(String subscription) {
  MockConsumer<String, ByteBuffer> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  if (!createdConsumers.containsKey(subscription)) {
    createdConsumers.put(subscription, new ArrayList<>());
  }
  createdConsumers.get(subscription).add(consumer);

  MockConsumerConfiguration configuration = consumerConfigurations.get(subscription);
  if (configuration != null) {
    consumer.updatePartitions(configuration.topic, configuration.partitionInfoList);
    consumer.updateBeginningOffsets(configuration.startOffsets);
    consumer.updateEndOffsets(configuration.endOffsets);
  }
  return consumer;
}
 

@Test
public void shouldApplyTopicRenameWhenCheckingHealth() {
  MockConsumer<byte[], byte[]> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  List<PartitionInfo> partitionInfoList =
      Arrays.asList(
          new PartitionInfo("replace1", 0, null, null, null),
          new PartitionInfo("replace1", 1, null, null, null));
  consumer.updatePartitions("replace1", partitionInfoList);

  SourcePartitionValidator sourcePartitionHealthChecker =
      new SourcePartitionValidator(
          consumer,
          SourcePartitionValidator.MatchingStrategy.TOPIC,
          t -> t.equals("topic1") ? "replace1" : t);
  assertThat(sourcePartitionHealthChecker.isHealthy(new TopicPartition("topic1", 0)), is(true));
  assertThat(sourcePartitionHealthChecker.isHealthy(new TopicPartition("topic1", 2)), is(true));
  assertThat(sourcePartitionHealthChecker.isHealthy(new TopicPartition("topic2", 0)), is(false));
}
 

/**
 * Returns the offset that would be used for the partition based on the consumer's configured offset reset strategy
 *
 * @return the offset that would be used for the partition based on the consumer's configured offset reset strategy
 */
protected long getResetOffset() {
    OffsetResetStrategy strategy = config.getOffsetResetStrategy();

    if (strategy == OffsetResetStrategy.EARLIEST) {
        LOGGER.debug("Looking up offset for partition [{}] using earliest reset", topicPartition);
        return getEarliestOffset();
    }
    else if (strategy == OffsetResetStrategy.LATEST) {
        LOGGER.debug("Looking up offset for partition [{}] using latest reset", topicPartition);
        return getLatestOffset();
    }
    else {
        throw new IllegalStateException("Unable to reset partition to previously committed offset as there is no"
                + " offset for partition [" + topicPartition + "] and offset reset strategy [" + strategy + "] is unknown");
    }
}
 

@Test
public void testConsume(TestContext ctx) throws Exception {
  MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  KafkaReadStream<String, String> consumer = createConsumer(vertx, mock);
  Async doneLatch = ctx.async();
  consumer.handler(record -> {
    ctx.assertEquals("the_topic", record.topic());
    ctx.assertEquals(0, record.partition());
    ctx.assertEquals("abc", record.key());
    ctx.assertEquals("def", record.value());
    consumer.close(v -> doneLatch.complete());
  });
  consumer.subscribe(Collections.singleton("the_topic"), v -> {
    mock.schedulePollTask(() -> {
      mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0)));
      mock.addRecord(new ConsumerRecord<>("the_topic", 0, 0L, "abc", "def"));
      mock.seek(new TopicPartition("the_topic", 0), 0L);
    });
  });
}
 

@Test
public void testConsumeWithHeader(TestContext ctx) {
  MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  KafkaReadStream<String, String> consumer = createConsumer(vertx, mock);
  Async doneLatch = ctx.async();
  consumer.handler(record -> {
    ctx.assertEquals("the_topic", record.topic());
    ctx.assertEquals(0, record.partition());
    ctx.assertEquals("abc", record.key());
    ctx.assertEquals("def", record.value());
    Header[] headers = record.headers().toArray();
    ctx.assertEquals(1, headers.length);
    Header header = headers[0];
    ctx.assertEquals("header_key", header.key());
    ctx.assertEquals("header_value", new String(header.value()));
    consumer.close(v -> doneLatch.complete());
  });
  consumer.subscribe(Collections.singleton("the_topic"), v -> {
    mock.schedulePollTask(() -> {
      mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0)));
      mock.addRecord(new ConsumerRecord<>("the_topic", 0, 0L, 0L, TimestampType.NO_TIMESTAMP_TYPE, 0L, 0, 0, "abc", "def",
        new RecordHeaders(Collections.singletonList(new RecordHeader("header_key", "header_value".getBytes())))));
      mock.seek(new TopicPartition("the_topic", 0), 0L);
    });
  });
}
 

public KafkaMockConsumer(final OffsetResetStrategy offsetResetStrategy) {
  super(OffsetResetStrategy.EARLIEST);

  assign(Arrays.asList(new TopicPartition("testtopic", 0)));

  HashMap<TopicPartition, Long> beginningOffsets = new HashMap<>();
  beginningOffsets.put(new TopicPartition("testtopic", 0), 0L);
  updateBeginningOffsets(beginningOffsets);

  for (int idx = 0; idx < 10; ++idx) {
    addRecord(
        new ConsumerRecord<>("testtopic",
            0, idx,
            ("mykey" + idx).getBytes(StandardCharsets.UTF_8),
            ("myvalue" + idx).getBytes(StandardCharsets.UTF_8)));
  }
}
 

private MockConsumer<byte[], byte[]> mockSourceConsumer() {
  MockConsumer<byte[], byte[]> mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  updateMockPartitions(mockConsumer, "topic1", 2);
  updateMockPartitions(mockConsumer, "topic2", 1);
  updateMockPartitions(mockConsumer, "topic3", 1);
  updateMockPartitions(mockConsumer, "topic4", 1);
  updateMockPartitions(mockConsumer, "topic5", 1);
  updateMockPartitions(mockConsumer, "reroute.outgoing", 1);
  return mockConsumer;
}
 

private MockConsumer<byte[], byte[]> mockDestinationConsumer() {
  // Topic 5 is NOT present in destination
  MockConsumer<byte[], byte[]> mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  updateMockPartitions(mockConsumer, "topic1", 2);
  updateMockPartitions(mockConsumer, "topic2", 1);
  updateMockPartitions(mockConsumer, "topic3", 1);
  updateMockPartitions(mockConsumer, "topic4", 1);
  updateMockPartitions(mockConsumer, "reroute.incoming", 1);
  return mockConsumer;
}
 

@Test
public void shouldContinueRunningWhenExceptionEncountered() throws InterruptedException {
  Map<String, String> properties = getBaseProperties();
  SourceConfig config = new SourceConfig(properties);
  TaskConfigBuilder taskConfigBuilder =
      new TaskConfigBuilder(new RoundRobinTaskAssignor(), config);

  // Require two thrown exceptions to ensure that the KafkaMonitor run loop executes more than
  // once
  CountDownLatch exceptionThrownLatch = new CountDownLatch(2);
  MockConsumer<byte[], byte[]> consumer =
      new MockConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {
        @Override
        public Map<String, List<PartitionInfo>> listTopics() {
          exceptionThrownLatch.countDown();
          throw new TimeoutException("KABOOM!");
        }
      };

  kafkaMonitor =
      new KafkaMonitor(
          mock(ConnectorContext.class),
          config,
          consumer,
          mockDestinationConsumer,
          taskConfigBuilder);
  Thread monitorThread = new Thread(kafkaMonitor);
  monitorThread.start();
  exceptionThrownLatch.await(2, TimeUnit.SECONDS);
  monitorThread.join(1);

  assertThat(monitorThread.getState(), not(State.TERMINATED));
  kafkaMonitor.stop();
  monitorThread.interrupt();
  monitorThread.join(5000);
}
 

@Before
public void setUp() {
  this.mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  List<PartitionInfo> partitionInfoList =
      Arrays.asList(
          new PartitionInfo("topic1", 0, null, null, null),
          new PartitionInfo("topic1", 1, null, null, null));
  mockConsumer.updatePartitions("topic1", partitionInfoList);
}
 

@Before
public void setUp() {
  mockConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  mockConsumer.updatePartitions(
      TOPIC,
      Arrays.asList(
          new PartitionInfo(TOPIC, 0, null, null, null),
          new PartitionInfo(TOPIC, 1, null, null, null)));
  mirusSourceTask = new MirusSourceTask(consumerProperties -> mockConsumer);

  // Always return offset = 0
  SourceTaskContext context =
      new SourceTaskContext() {
        @Override
        public Map<String, String> configs() {
          return null;
        }

        @Override
        public OffsetStorageReader offsetStorageReader() {
          return new OffsetStorageReader() {
            @Override
            public <T> Map<String, Object> offset(Map<String, T> partition) {
              return new HashMap<>(MirusSourceTask.offsetMap(0L));
            }

            @Override
            public <T> Map<Map<String, T>, Map<String, Object>> offsets(
                Collection<Map<String, T>> partitions) {
              return partitions.stream().collect(Collectors.toMap(p -> p, this::offset));
            }
          };
        }
      };
  mirusSourceTask.initialize(context);
  mirusSourceTask.start(mockTaskProperties());
}
 

@Test
public void testBatch(TestContext ctx) throws Exception {
  int num = 50;
  MockConsumer<String, String> mock = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  KafkaReadStream<String, String> consumer = createConsumer(vertx, mock);
  Async doneLatch = ctx.async();
  AtomicInteger count = new AtomicInteger();
  consumer.handler(record -> {
    int val = count.getAndIncrement();
    if (val < num) {
      ctx.assertEquals("the_topic", record.topic());
      ctx.assertEquals(0, record.partition());
      ctx.assertEquals("key-" + val, record.key());
      ctx.assertEquals("value-" + val, record.value());
      if (val == num - 1) {
        consumer.close(v -> doneLatch.complete());
      }
    }
  });
  consumer.subscribe(Collections.singleton("the_topic"), v -> {
    mock.schedulePollTask(() -> {
      mock.rebalance(Collections.singletonList(new TopicPartition("the_topic", 0)));
      mock.seek(new TopicPartition("the_topic", 0), 0);
      for (int i = 0; i < num; i++) {
        mock.addRecord(new ConsumerRecord<>("the_topic", 0, i, "key-" + i, "value-" + i));
      }
    });
  });
}
 

private MockConsumer<String, String> createMockConsumer(){
    MockConsumer<String, String> consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);

    Map<org.apache.kafka.common.TopicPartition, Long> beginningOffsets = new HashMap<>();
    beginningOffsets.put(new org.apache.kafka.common.TopicPartition(TOPIC, 0), 0L);
    consumer.updateBeginningOffsets(beginningOffsets);
    return consumer;
}
 

@Before
public void setUp() {
    GarmadonReader.GarmadonMessageHandler garmadonMessageHandler = mock(GarmadonReader.GarmadonMessageHandler.class);
    MockConsumer kafkaConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
    GarmadonReader.Builder builder = GarmadonReader.Builder.stream(kafkaConsumer);
    GarmadonReader garmadonReader = builder
            .intercept(GarmadonMessageFilter.ANY.INSTANCE, garmadonMessageHandler)
            .build(false);

    garmadonReaderBuilder = Mockito.mock(GarmadonReader.Builder.class);
    when(garmadonReaderBuilder.intercept(any(GarmadonMessageFilter.class), any(GarmadonReader.GarmadonMessageHandler.class))).thenReturn(garmadonReaderBuilder);
    when(garmadonReaderBuilder.build()).thenReturn(garmadonReader);

    bulkProcessor = Mockito.mock(BulkProcessor.class);
    prometheusHttpConsumerMetrics = Mockito.mock(PrometheusHttpConsumerMetrics.class);
    elasticSearchCacheManager = Mockito.mock(ElasticSearchCacheManager.class);

    elasticSearchReader = new ElasticSearchReader(garmadonReaderBuilder, bulkProcessor,
            "garmadon-index", prometheusHttpConsumerMetrics, elasticSearchCacheManager);

    header = EventHeaderProtos.Header.newBuilder()
            .setUsername("user")
            .setApplicationId("app_id")
            .setApplicationName("application_name")
            .setAttemptId("attempt_id")
            .setContainerId("container_id")
            .build();

    headerMap.put("pid", "");
    headerMap.put("main_class", "");
    headerMap.put("application_id", "app_id");
    headerMap.put("tags", new ArrayList<>());
    headerMap.put("hostname", "");
    headerMap.put("component", "");
    headerMap.put("application_name", "application_name");
    headerMap.put("framework", "");
    headerMap.put("attempt_id", "attempt_id");
    headerMap.put("container_id", "container_id");
    headerMap.put("username", "user");
    headerMap.put("executor_id", "");
    headerMap.put("timestamp", 0);
}
 

@Test
public void shouldApplyTopicRenameTransforms() {
  Map<String, String> properties = getBaseProperties();
  properties.put(SourceConfigDefinition.TOPICS_REGEX.getKey(), "reroute.*");
  properties.put("transforms", "reroute");
  properties.put("transforms.reroute.type", "org.apache.kafka.connect.transforms.RegexRouter");
  properties.put("transforms.reroute.regex", "^reroute\\.outgoing$");
  properties.put("transforms.reroute.replacement", "reroute.incoming");
  SourceConfig config = new SourceConfig(properties);

  MockConsumer<byte[], byte[]> mockSource = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  updateMockPartitions(mockSource, "reroute.outgoing", 1);

  MockConsumer<byte[], byte[]> mockDest = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
  updateMockPartitions(mockDest, "reroute.incoming", 1);

  TaskConfigBuilder taskConfigBuilder =
      new TaskConfigBuilder(new RoundRobinTaskAssignor(), config);
  KafkaMonitor monitor =
      new KafkaMonitor(
          mock(ConnectorContext.class), config, mockSource, mockDest, taskConfigBuilder);

  monitor.partitionsChanged();
  List<Map<String, String>> result = monitor.taskConfigs(3);

  List<TopicPartition> partitions = assignedTopicPartitionsFromTaskConfigs(result);
  assertThat(partitions, contains(new TopicPartition("reroute.outgoing", 0)));
}
 

@BeforeEach
void setUp() {
    consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
    updates = new ArrayList<>();
    countryPopulationConsumer = new CountryPopulationConsumer(consumer, ex -> this.pollException = ex, updates::add);
}