net.mguenther.kafka.junit.ObserveKeyValues Maven / Gradle / Ivy

Go to download
package net.mguenther.kafka.junit;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.ToString;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.serialization.StringDeserializer;

import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;

@Getter
@ToString
@RequiredArgsConstructor
public class ObserveKeyValues {

    public static final int DEFAULT_OBSERVATION_TIME_MILLIS = 30_000;

    public static class ObserveKeyValuesBuilder {

        private final String topic;
        private final int expected;
        private final Class clazzOfK;
        private final Class clazzOfV;
        private final Properties consumerProps = new Properties();
        private Predicate filterOnKeys = key -> true;
        private Predicate filterOnValues = value -> true;
        private Predicate filterOnHeaders = value -> true;
        private int observationTimeMillis = DEFAULT_OBSERVATION_TIME_MILLIS;
        private boolean includeMetadata = false;
        private Map seekTo = new HashMap<>();

        ObserveKeyValuesBuilder(final String topic, final int expected, final Class clazzOfK, final Class clazzOfV) {
            this.topic = topic;
            this.expected = expected;
            this.clazzOfK = clazzOfK;
            this.clazzOfV = clazzOfV;
        }

        public ObserveKeyValuesBuilder observeFor(final int duration, final TimeUnit unit) {
            this.observationTimeMillis = (int) unit.toMillis(duration);
            return this;
        }

        public ObserveKeyValuesBuilder filterOnKeys(final Predicate filterOnKeys) {
            this.filterOnKeys = filterOnKeys;
            return this;
        }

        public ObserveKeyValuesBuilder filterOnValues(final Predicate filterOnValues) {
            this.filterOnValues = filterOnValues;
            return this;
        }

        public ObserveKeyValuesBuilder filterOnHeaders(final Predicate filterOnHeaders) {
            this.filterOnHeaders = filterOnHeaders;
            return this;
        }

        public ObserveKeyValuesBuilder includeMetadata() {
            return withMetadata(true);
        }

        public ObserveKeyValuesBuilder withMetadata(final boolean modifier) {
            this.includeMetadata = modifier;
            return this;
        }

        public ObserveKeyValuesBuilder seekTo(final int partition, final long offset) {
            seekTo.put(partition, offset);
            return this;
        }

        public ObserveKeyValuesBuilder seekTo(final Map seekTo) {
            this.seekTo.putAll(seekTo);
            return this;
        }

        public  ObserveKeyValuesBuilder with(final String propertyName, final T value) {
            consumerProps.put(propertyName, value);
            return this;
        }

        public  ObserveKeyValuesBuilder withAll(final Properties consumerProps) {
            this.consumerProps.putAll(consumerProps);
            return this;
        }

        private  void ifNonExisting(final String propertyName, final T value) {
            if (consumerProps.get(propertyName) != null) return;
            consumerProps.put(propertyName, value);
        }

        public ObserveKeyValues useDefaults() {
            consumerProps.clear();
            observationTimeMillis = DEFAULT_OBSERVATION_TIME_MILLIS;
            return build();
        }

        public ObserveKeyValues build() {
            ifNonExisting(ConsumerConfig.GROUP_ID_CONFIG, UUID.randomUUID().toString());
            ifNonExisting(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
            ifNonExisting(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
            ifNonExisting(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
            ifNonExisting(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
            ifNonExisting(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 100);
            ifNonExisting(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_uncommitted");
            return new ObserveKeyValues<>(topic, expected, observationTimeMillis, includeMetadata, seekTo, consumerProps, filterOnKeys, filterOnValues, filterOnHeaders, clazzOfK, clazzOfV);
        }
    }

    private final String topic;
    private final int expected;
    private final int observationTimeMillis;
    private final boolean includeMetadata;
    private final Map seekTo;
    private final Properties consumerProps;
    private final Predicate filterOnKeys;
    private final Predicate filterOnValues;
    private final Predicate filterOnHeaders;
    private final Class clazzOfK;
    private final Class clazzOfV;

    public static ObserveKeyValuesBuilder on(final String topic, final int expected) {
        return on(topic, expected, String.class, String.class);
    }

    public static  ObserveKeyValuesBuilder on(final String topic,
                                                            final int expected,
                                                            final Class clazzOfV) {
        return on(topic, expected, String.class, clazzOfV);
    }

    public static  ObserveKeyValuesBuilder on(final String topic,
                                                          final int expected,
                                                          final Class clazzOfK,
                                                          final Class clazzOfV) {
        return new ObserveKeyValuesBuilder<>(topic, expected, clazzOfK, clazzOfV);
    }
}