net.mguenther.kafka.junit.SendKeyValuesTransactional 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.producer.ProducerConfig;
import org.apache.kafka.common.serialization.StringSerializer;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.UUID;

@ToString
@RequiredArgsConstructor
public class SendKeyValuesTransactional {

    public static class SendKeyValuesTransactionalBuilder {

        private final Map>> recordsPerTopic = new HashMap<>();
        private final Properties producerProps = new Properties();
        private boolean failTransaction = false;

        SendKeyValuesTransactionalBuilder(final String topic, final Collection> records) {
            recordsPerTopic.put(topic, records);
        }

        SendKeyValuesTransactionalBuilder(final Map>> recordsPerTopic) {
            this.recordsPerTopic.putAll(recordsPerTopic);
        }

        public SendKeyValuesTransactionalBuilder inTransaction(final String topic, final Collection> records) {
            final Collection> existingRecordsForTopic = recordsPerTopic.getOrDefault(topic, new ArrayList<>());
            existingRecordsForTopic.addAll(records);
            recordsPerTopic.put(topic, existingRecordsForTopic);
            return this;
        }

        public SendKeyValuesTransactionalBuilder failTransaction() {
            return withFailTransaction(true);
        }

        public SendKeyValuesTransactionalBuilder withFailTransaction(final boolean modifier) {
            failTransaction = modifier;
            return this;
        }

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

        public SendKeyValuesTransactionalBuilder withAll(final Properties transactionalProps) {
            this.producerProps.putAll(transactionalProps);
            return this;
        }

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

        public SendKeyValuesTransactional useDefaults() {
            producerProps.clear();
            failTransaction = false;
            return build();
        }

        public SendKeyValuesTransactional build() {
            ifNonExisting(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
            ifNonExisting(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
            ifNonExisting(ProducerConfig.RETRIES_CONFIG, Integer.MAX_VALUE);
            ifNonExisting(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, 1);
            ifNonExisting(ProducerConfig.ENABLE_IDEMPOTENCE_CONFIG, true);
            ifNonExisting(ProducerConfig.TRANSACTIONAL_ID_CONFIG, UUID.randomUUID().toString());
            ifNonExisting(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, 60_000);
            ifNonExisting(ProducerConfig.ACKS_CONFIG, "all");
            return new SendKeyValuesTransactional<>(recordsPerTopic, failTransaction, producerProps);
        }
    }

    @Getter
    private final Map>> recordsPerTopic;

    private final boolean failTransaction;

    @Getter
    private final Properties producerProps;

    public boolean shouldFailTransaction() {
        return failTransaction;
    }

    public static  SendKeyValuesTransactionalBuilder inTransaction(final String topic, final Collection> records) {
        return new SendKeyValuesTransactionalBuilder<>(topic, records);
    }

    public static  SendKeyValuesTransactionalBuilder inTransaction(final Map>> recordsPerTopic) {
        return new SendKeyValuesTransactionalBuilder<>(recordsPerTopic);
    }
}