<commit_before>package com.ameliant.tools.kafkaperf.drivers;

import com.ameliant.tools.kafkaperf.config.ConsumerConfigsBuilder;
import com.ameliant.tools.kafkaperf.config.ConsumerDefinition;
import com.ameliant.tools.kafkaperf.config.ProducerConfigsBuilder;
import com.ameliant.tools.kafkaperf.config.ProducerDefinition;
import com.ameliant.tools.kafkaperf.resources.EmbeddedKafkaBroker;
import com.ameliant.tools.kafkaperf.resources.EmbeddedZooKeeper;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.junit.Rule;
import org.junit.Test;

import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static java.lang.String.format;

/**
 * @author jkorab
 */
public class ProducerPartitioningTest {

    @Rule
    public EmbeddedZooKeeper zooKeeper = new EmbeddedZooKeeper();

    @Rule
    public EmbeddedKafkaBroker broker = new EmbeddedKafkaBroker.Builder()
            .zookeeperConnect("127.0.0.1:" + zooKeeper.getPort())
            .build();

        @Test
    public void testSend_partitions() throws InterruptedException {
        Map<String, Object> producerConfigs = getProducerConfigs();

        ProducerDefinition producerDefinition = new ProducerDefinition();
        producerDefinition.setConfig(producerConfigs);
        producerDefinition.setTopic("foo");
        producerDefinition.setMessageSize(1024);
        producerDefinition.setMessagesToSend(100);
        producerDefinition.setSendBlocking(true);
        producerDefinition.setUniqueKeyCount(2);

        int numWorkers = 3;
        CountDownLatch latch = new CountDownLatch(numWorkers);

        ProducerDriver producer = new ProducerDriver(producerDefinition, latch);
        ConsumerDriver consumer1 = createConsumerDriver(latch, "foo", "bar", 50);
        ConsumerDriver consumer2 = createConsumerDriver(latch, "foo", "bar", 50);

        ExecutorService executorService = Executors.newFixedThreadPool(numWorkers);
        try {
            // fill up the topic
            executorService.submit(producer);
            executorService.submit(consumer1);
            executorService.submit(consumer2);

            if (!latch.await(15, TimeUnit.SECONDS)) {
                String consumerResults = Arrays.asList(consumer1, consumer2).stream()
                        .map(consumer -> consumer.getMessagesExpected() + ":" + consumer.getMessagesReceived())
                        .reduce("", (results, expectedVsActual) -> {
                            String retval = results;
                            retval += (results == "") ? expectedVsActual : ", " + expectedVsActual;
                            return retval;
                        });
                throw new RuntimeException("Consumers did not receive expected message counts: [" + consumerResults + "]");
            }
        } finally {
            executorService.shutdownNow();
        }

    }


    private Map<String, Object> getProducerConfigs() {
        return new ProducerConfigsBuilder()
                    .bootstrapServers("127.0.0.1:" + broker.getPort())
                    .requestRequiredAcks(ProducerConfigsBuilder.RequestRequiredAcks.ackFromLeader)
                    .producerType(ProducerConfigsBuilder.ProducerType.sync)
                    .keySerializer(ByteArraySerializer.class)
                    .valueSerializer(ByteArraySerializer.class)
                    .batchSize(0)
                    .build();
    }
    private ConsumerDriver createConsumerDriver(CountDownLatch latch, String topic, String groupId, int messageCount) {
        Map<String, Object> configs = new ConsumerConfigsBuilder()
                .groupId(groupId)
                .bootstrapServers("127.0.0.1:" + broker.getPort())
                .enableAutoCommit(true)
                .autoCommitIntervalMs(1000)
                .sessionTimeoutMs(30000)
                .keyDeserializer(ByteArrayDeserializer.class)
                .valueDeserializer(ByteArrayDeserializer.class)
                .autoOffsetReset(ConsumerConfigsBuilder.OffsetReset.earliest)
                .build();

        ConsumerDefinition consumerDefinition = new ConsumerDefinition();
        consumerDefinition.setConfig(configs);
        consumerDefinition.setTopic(topic);
        consumerDefinition.setMessagesToReceive(messageCount);<commit_msg>Ignoring test, work in progress.<commit_after>package com.ameliant.tools.kafkaperf.drivers;

import com.ameliant.tools.kafkaperf.config.ConsumerConfigsBuilder;
import com.ameliant.tools.kafkaperf.config.ConsumerDefinition;
import com.ameliant.tools.kafkaperf.config.ProducerConfigsBuilder;
import com.ameliant.tools.kafkaperf.config.ProducerDefinition;
import com.ameliant.tools.kafkaperf.resources.EmbeddedKafkaBroker;
import com.ameliant.tools.kafkaperf.resources.EmbeddedZooKeeper;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;

import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static java.lang.String.format;

/**
 * @author jkorab
 */
public class ProducerPartitioningTest {

    @Rule
    public EmbeddedZooKeeper zooKeeper = new EmbeddedZooKeeper();

    @Rule
    public EmbeddedKafkaBroker broker = new EmbeddedKafkaBroker.Builder()
            .zookeeperConnect("127.0.0.1:" + zooKeeper.getPort())
            .build();

    @Test @Ignore
    public void testSend_partitions() throws InterruptedException {
        Map<String, Object> producerConfigs = getProducerConfigs();

        ProducerDefinition producerDefinition = new ProducerDefinition();
        producerDefinition.setConfig(producerConfigs);
        producerDefinition.setTopic("foo");
        producerDefinition.setMessageSize(1024);
        producerDefinition.setMessagesToSend(100);
        producerDefinition.setSendBlocking(true);
        producerDefinition.setUniqueKeyCount(2);

        int numWorkers = 3;
        CountDownLatch latch = new CountDownLatch(numWorkers);

        ProducerDriver producer = new ProducerDriver(producerDefinition, latch);
        ConsumerDriver consumer1 = createConsumerDriver(latch, "foo", "bar", 50);
        ConsumerDriver consumer2 = createConsumerDriver(latch, "foo", "bar", 50);

        ExecutorService executorService = Executors.newFixedThreadPool(numWorkers);
        try {
            // fill up the topic
            executorService.submit(producer);
            executorService.submit(consumer1);
            executorService.submit(consumer2);

            if (!latch.await(15, TimeUnit.SECONDS)) {
                String consumerResults = Arrays.asList(consumer1, consumer2).stream()
                        .map(consumer -> consumer.getMessagesExpected() + ":" + consumer.getMessagesReceived())
                        .reduce("", (results, expectedVsActual) -> {
                            String retval = results;
                            retval += (results == "") ? expectedVsActual : ", " + expectedVsActual;
                            return retval;
                        });
                throw new RuntimeException("Consumers did not receive expected message counts: [" + consumerResults + "]");
            }
        } finally {
            executorService.shutdownNow();
        }

    }

    private Map<String, Object> getProducerConfigs() {
        return new ProducerConfigsBuilder()
                    .bootstrapServers("127.0.0.1:" + broker.getPort())
                    .requestRequiredAcks(ProducerConfigsBuilder.RequestRequiredAcks.ackFromLeader)
                    .producerType(ProducerConfigsBuilder.ProducerType.sync)
                    .keySerializer(ByteArraySerializer.class)
                    .valueSerializer(ByteArraySerializer.class)
                    .batchSize(0)
                    .build();
    }

    private ConsumerDriver createConsumerDriver(CountDownLatch latch, String topic, String groupId, int messageCount) {
        Map<String, Object> configs = new ConsumerConfigsBuilder()
                .groupId(groupId)
                .bootstrapServers("127.0.0.1:" + broker.getPort())
                .enableAutoCommit(true)
                .autoCommitIntervalMs(1000)
                .sessionTimeoutMs(30000)
                .keyDeserializer(ByteArrayDeserializer.class)
                .valueDeserializer(ByteArrayDeserializer.class)
                .autoOffsetReset(ConsumerConfigsBuilder.OffsetReset.earliest)
                .build();

        ConsumerDefinition consumerDefinition = new ConsumerDefinition();
        consumerDefinition.setConfig(configs);
        consumerDefinition.setTopic(topic);
        consumerDefinition.setMessagesToReceive(messageCount);
{'fraction_non_alphanumeric': 0.09290577641115748, 'fraction_numerical': 0.010762134856138809, 'mean_word_length': 2.7647788342290203, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 6, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4267425', 'n_tokens_mistral': 2529, 'n_tokens_neox': 2340, 'n_words': 392}