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.clients.consumer.internals;
  
  
  import org.apache.kafka.clients.consumer.ConsumerInterceptor;
  import org.apache.kafka.clients.consumer.ConsumerRecord;
  import org.apache.kafka.clients.consumer.ConsumerRecords;
  import org.apache.kafka.clients.consumer.OffsetAndMetadata;
  import org.apache.kafka.common.KafkaException;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.header.internals.RecordHeaders;
  import org.apache.kafka.common.record.TimestampType;
  
  import org.junit.jupiter.api.Test;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  public class ConsumerInterceptorsTest {
      private final int filterPartition1 = 5;
      private final int filterPartition2 = 6;
      private final String topic = "test";
      private final int partition = 1;
      private final TopicPartition tp = new TopicPartition(topic, partition);
      private final TopicPartition filterTopicPart1 = new TopicPartition("test5", filterPartition1);
      private final TopicPartition filterTopicPart2 = new TopicPartition("test6", filterPartition2);
      private final ConsumerRecord<Integer, Integer> consumerRecord = new ConsumerRecord<>(topic, partition, 0, 0L,
          TimestampType.CREATE_TIME, 0, 0, 1, 1, new RecordHeaders(), Optional.empty());
      private int onCommitCount = 0;
      private int onConsumeCount = 0;
  
      /**
       * Test consumer interceptor that filters records in onConsume() intercept
       */
      private class FilterConsumerInterceptor<K, V> implements ConsumerInterceptor<K, V> {
          private final int filterPartition;
          private boolean throwExceptionOnConsume = false;
          private boolean throwExceptionOnCommit = false;
  
          FilterConsumerInterceptor(int filterPartition) {
              this.filterPartition = filterPartition;
          }
  
          @Override
          public void configure(Map<String, ?> configs) {
          }
  
          @Override
          public ConsumerRecords<K, V> onConsume(ConsumerRecords<K, V> records) {
              onConsumeCount++;
              if (throwExceptionOnConsume)
                  throw new KafkaException("Injected exception in FilterConsumerInterceptor.onConsume.");
  
              // filters out topic/partitions with partition == FILTER_PARTITION
              Map<TopicPartition, List<ConsumerRecord<K, V>>> recordMap = new HashMap<>();
              Map<TopicPartition, OffsetAndMetadata> nextOffsetAndMetadata = new HashMap<>();
              for (TopicPartition tp : records.partitions()) {
                  if (tp.partition() != filterPartition) {
                      recordMap.put(tp, records.records(tp));
                      nextOffsetAndMetadata.put(tp, records.nextOffsets().get(tp));
                  }
              }
              return new ConsumerRecords<>(recordMap, nextOffsetAndMetadata);
          }
  
          @Override
          public void onCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {
              onCommitCount++;
              if (throwExceptionOnCommit)
                  throw new KafkaException("Injected exception in FilterConsumerInterceptor.onCommit.");
          }
  
          @Override
          public void close() {
          }
  
          // if 'on' is true, onConsume will always throw an exception
          public void injectOnConsumeError(boolean on) {
              throwExceptionOnConsume = on;
          }
  
          // if 'on' is true, onConsume will always throw an exception
          public void injectOnCommitError(boolean on) {
              throwExceptionOnCommit = on;
          }
      }
  
      @Test
      public void testOnConsumeChain() {
          List<ConsumerInterceptor<Integer, Integer>>  interceptorList = new ArrayList<>();
          // we are testing two different interceptors by configuring the same interceptor differently, which is not
          // how it would be done in KafkaConsumer, but ok for testing interceptor callbacks
          FilterConsumerInterceptor<Integer, Integer> interceptor1 = new FilterConsumerInterceptor<>(filterPartition1);
          FilterConsumerInterceptor<Integer, Integer> interceptor2 = new FilterConsumerInterceptor<>(filterPartition2);
          interceptorList.add(interceptor1);
          interceptorList.add(interceptor2);
          ConsumerInterceptors<Integer, Integer> interceptors = new ConsumerInterceptors<>(interceptorList, null);
  
          // verify that onConsumer modifies ConsumerRecords
          Map<TopicPartition, List<ConsumerRecord<Integer, Integer>>> records = new HashMap<>();
          Map<TopicPartition, OffsetAndMetadata> nextOffsets = new HashMap<>();
  
          List<ConsumerRecord<Integer, Integer>> list1 = new ArrayList<>();
          list1.add(consumerRecord);
          List<ConsumerRecord<Integer, Integer>> list2 = new ArrayList<>();
          list2.add(new ConsumerRecord<>(filterTopicPart1.topic(), filterTopicPart1.partition(), 0, 0L,
              TimestampType.CREATE_TIME, 0, 0, 1, 1, new RecordHeaders(), Optional.empty()));
          List<ConsumerRecord<Integer, Integer>> list3 = new ArrayList<>();
          list3.add(new ConsumerRecord<>(filterTopicPart2.topic(), filterTopicPart2.partition(), 0, 0L, TimestampType.CREATE_TIME,
              0, 0, 1, 1, new RecordHeaders(), Optional.empty()));
          records.put(tp, list1);
          nextOffsets.put(tp, new OffsetAndMetadata(1, Optional.empty(), ""));
  
          records.put(filterTopicPart1, list2);
          nextOffsets.put(filterTopicPart1, new OffsetAndMetadata(1, Optional.empty(), ""));
  
          records.put(filterTopicPart2, list3);
          nextOffsets.put(filterTopicPart2, new OffsetAndMetadata(1, Optional.empty(), ""));
  
          ConsumerRecords<Integer, Integer> consumerRecords = new ConsumerRecords<>(records, nextOffsets);
          ConsumerRecords<Integer, Integer> interceptedRecords = interceptors.onConsume(consumerRecords);
          assertEquals(1, interceptedRecords.count());
          assertTrue(interceptedRecords.partitions().contains(tp));
          assertFalse(interceptedRecords.partitions().contains(filterTopicPart1));
          assertFalse(interceptedRecords.partitions().contains(filterTopicPart2));
          assertEquals(2, onConsumeCount);
          validateNextOffsets(interceptedRecords.nextOffsets(), 1);
  
          // verify that even if one of the intermediate interceptors throws an exception, all interceptors' onConsume are called
          interceptor1.injectOnConsumeError(true);
          ConsumerRecords<Integer, Integer> partInterceptedRecs = interceptors.onConsume(consumerRecords);
          assertEquals(2, partInterceptedRecs.count());
          assertTrue(partInterceptedRecs.partitions().contains(filterTopicPart1));  // since interceptor1 threw exception
          assertFalse(partInterceptedRecs.partitions().contains(filterTopicPart2)); // interceptor2 should still be called
          assertEquals(4, onConsumeCount);
          validateNextOffsets(partInterceptedRecs.nextOffsets(), 2);
  
          // if all interceptors throw an exception, records should be unmodified
          interceptor2.injectOnConsumeError(true);
          ConsumerRecords<Integer, Integer> noneInterceptedRecs = interceptors.onConsume(consumerRecords);
          assertEquals(noneInterceptedRecs, consumerRecords);
          assertEquals(3, noneInterceptedRecs.count());
          assertEquals(6, onConsumeCount);
          validateNextOffsets(noneInterceptedRecs.nextOffsets(), 3);
  
          interceptors.close();
      }
  
      @Test
      public void testOnCommitChain() {
          List<ConsumerInterceptor<Integer, Integer>> interceptorList = new ArrayList<>();
          // we are testing two different interceptors by configuring the same interceptor differently, which is not
          // how it would be done in KafkaConsumer, but ok for testing interceptor callbacks
          FilterConsumerInterceptor<Integer, Integer> interceptor1 = new FilterConsumerInterceptor<>(filterPartition1);
          FilterConsumerInterceptor<Integer, Integer> interceptor2 = new FilterConsumerInterceptor<>(filterPartition2);
          interceptorList.add(interceptor1);
          interceptorList.add(interceptor2);
          ConsumerInterceptors<Integer, Integer> interceptors = new ConsumerInterceptors<>(interceptorList, null);
  
          // verify that onCommit is called for all interceptors in the chain
          Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
          offsets.put(tp, new OffsetAndMetadata(0));
          interceptors.onCommit(offsets);
          assertEquals(2, onCommitCount);
  
          // verify that even if one of the interceptors throws an exception, all interceptors' onCommit are called
          interceptor1.injectOnCommitError(true);
          interceptors.onCommit(offsets);
          assertEquals(4, onCommitCount);
  
          interceptors.close();
      }
  
      private void validateNextOffsets(Map<TopicPartition, OffsetAndMetadata> nextOffsetAndMetadataMap, int size) {
          assertEquals(size, nextOffsetAndMetadataMap.size());
          if (size == 1) {
              assertEquals(new OffsetAndMetadata(1, Optional.empty(), ""), nextOffsetAndMetadataMap.get(tp));
          } else if (size == 2) {
              assertEquals(new OffsetAndMetadata(1, Optional.empty(), ""), nextOffsetAndMetadataMap.get(tp));
              assertEquals(new OffsetAndMetadata(1, Optional.empty(), ""), nextOffsetAndMetadataMap.get(filterTopicPart1));
          } else if (size == 3) {
              assertEquals(new OffsetAndMetadata(1, Optional.empty(), ""), nextOffsetAndMetadataMap.get(tp));
              assertEquals(new OffsetAndMetadata(1, Optional.empty(), ""), nextOffsetAndMetadataMap.get(filterTopicPart1));
              assertEquals(new OffsetAndMetadata(1, Optional.empty(), ""), nextOffsetAndMetadataMap.get(filterTopicPart2));
          }
      }
  }
  