Java Code Examples for com.salesforce.kafka.test.KafkaTestUtils#produceRecords()

@Test
public void testKafka011Source() throws Exception {
    final String topicName = "topic_2";
    KafkaTestUtils kafkaTestUtils = SHARED_KAFKA_TEST_RESOURCE.getKafkaTestUtils();
    kafkaTestUtils.createTopic(topicName, 1, (short) 1);
    kafkaTestUtils.produceRecords(4, topicName, 0);

    StreamOperator data = new Kafka011SourceStreamOp()
        .setBootstrapServers(SHARED_KAFKA_TEST_RESOURCE.getKafkaConnectString())
        .setGroupId("g")
        .setStartupMode("earliest")
        .setTopic(topicName);

    Assert.assertEquals(data.getColNames().length, 5);
}
 

@Test
public void testKafkaSource() throws Exception {
    final String topicName = "topic_2";
    KafkaTestUtils kafkaTestUtils = SHARED_KAFKA_TEST_RESOURCE.getKafkaTestUtils();
    kafkaTestUtils.createTopic(topicName, 1, (short) 1);
    kafkaTestUtils.produceRecords(4, topicName, 0);

    StreamOperator data = new KafkaSourceStreamOp()
        .setBootstrapServers(SHARED_KAFKA_TEST_RESOURCE.getKafkaConnectString())
        .setGroupId("g")
        .setStartupMode("earliest")
        .setTopic(topicName);

    Assert.assertEquals(data.getColNames().length, 5);
}
 

/**
 * Integration test validates that streams can be used against KafkaTestServer.
 */
@Test
void testStreamConsumer() throws Exception {
    // Define topic to test with.
    final String inputTopic = "stream-input-topic" + System.currentTimeMillis();
    final String outputTopic = "stream-output-topic" + System.currentTimeMillis();

    // Define how many records
    final int numberOfRecords = 25;
    final int partitionId = 0;

    // Tracks how many records the Stream consumer has processed.
    final AtomicInteger recordCounter = new AtomicInteger(0);

    // Create our test server instance.
    try (final KafkaTestServer kafkaTestServer = new KafkaTestServer()) {
        // Start it and create our topic.
        kafkaTestServer.start();

        // Create test utils instance.
        final KafkaTestUtils kafkaTestUtils = new KafkaTestUtils(kafkaTestServer);

        // Create topics
        kafkaTestUtils.createTopic(inputTopic, 1, (short) 1);
        kafkaTestUtils.createTopic(outputTopic, 1, (short) 1);

        // Produce random data into input topic
        kafkaTestUtils.produceRecords(numberOfRecords, inputTopic, partitionId);

        // Define stream consumer properties.
        final Properties config = new Properties();
        config.put(StreamsConfig.APPLICATION_ID_CONFIG, "testStreamProcessor");
        config.setProperty(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaTestServer.getKafkaConnectString());
        config.put("group.id", "test-stream-group");
        config.put("auto.offset.reset", "earliest");

        // Build the stream
        final StreamsBuilder streamsBuilder = new StreamsBuilder();
        streamsBuilder
            // Read from input topic.
            .stream(inputTopic)

            // For each record processed, increment our counter
            .map((key, word) -> {
                recordCounter.incrementAndGet();
                return new KeyValue<>(word, word);
            })

            // Write to output topic.
            .to(outputTopic);

        // Create stream
        final KafkaStreams kafkaStreams = new KafkaStreams(streamsBuilder.build(), new StreamsConfig(config));
        try {
            // Start the stream consumer
            kafkaStreams.start();

            // Since stream processing is async, we need to wait for the Stream processor to start, consume messages
            // from the input topic, and process them. We'll wait for Wait for it to do its thing up to 10 seconds.
            for (int timeoutCounter = 0; timeoutCounter <= 10; timeoutCounter++) {
                // If we've processed all of our records
                if (recordCounter.get() >= numberOfRecords) {
                    // Break out of sleep loop.
                    break;
                }
                // Otherwise, we need to wait longer, sleep 1 second.
                Thread.sleep(1000L);
            }
        } finally {
            // Close the stream consumer.
            kafkaStreams.close();
        }

        // Validation.
        Assertions.assertEquals(numberOfRecords, recordCounter.get(), "Should have 25 records processed");

        // Consume records from output topic.
        final List<ConsumerRecord<String, String>> outputRecords =
            kafkaTestUtils.consumeAllRecordsFromTopic(outputTopic, StringDeserializer.class, StringDeserializer.class);

        // Validate we got the correct number of records.
        Assertions.assertEquals(numberOfRecords, outputRecords.size());
    }
}
 

/**
 * Integration test validates that streams can be used against KafkaTestServer.
 */
@Test
void testStreamConsumer() throws Exception {
    // Define topic to test with.
    final String inputTopic = "stream-input-topic" + System.currentTimeMillis();
    final String outputTopic = "stream-output-topic" + System.currentTimeMillis();

    // Define how many records
    final int numberOfRecords = 25;
    final int partitionId = 0;

    // Tracks how many records the Stream consumer has processed.
    final AtomicInteger recordCounter = new AtomicInteger(0);

    // Create our test server instance.
    try (final KafkaTestServer kafkaTestServer = new KafkaTestServer()) {
        // Start it and create our topic.
        kafkaTestServer.start();

        // Create test utils instance.
        final KafkaTestUtils kafkaTestUtils = new KafkaTestUtils(kafkaTestServer);

        // Create topics
        kafkaTestUtils.createTopic(inputTopic, 1, (short) 1);
        kafkaTestUtils.createTopic(outputTopic, 1, (short) 1);

        // Produce random data into input topic
        kafkaTestUtils.produceRecords(numberOfRecords, inputTopic, partitionId);

        // Define stream consumer properties.
        final Properties config = new Properties();
        config.put(StreamsConfig.APPLICATION_ID_CONFIG, "testStreamProcessor");
        config.setProperty(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaTestServer.getKafkaConnectString());
        config.put("group.id", "test-stream-group");
        config.put("auto.offset.reset", "earliest");

        // Serialization definition.
        final Serde<String> stringSerde = Serdes.String();

        // Build the stream
        final KStreamBuilder kStreamBuilder = new KStreamBuilder();
        kStreamBuilder
            // Read from input topic.
            .stream(stringSerde, stringSerde, inputTopic)

            // For each record processed, increment our counter
            .map((key, word) -> {
                recordCounter.incrementAndGet();
                return new KeyValue<>(word, word);
            })

            // Write to output topic.
            .to(stringSerde, stringSerde, outputTopic);

        // Create stream
        final KafkaStreams kafkaStreams = new KafkaStreams(kStreamBuilder, new StreamsConfig(config));
        try {
            // Start the stream consumer
            kafkaStreams.start();

            // Since stream processing is async, we need to wait for the Stream processor to start, consume messages
            // from the input topic, and process them. We'll wait for Wait for it to do its thing up to 10 seconds.
            for (int timeoutCounter = 0; timeoutCounter <= 10; timeoutCounter++) {
                // If we've processed all of our records
                if (recordCounter.get() >= numberOfRecords) {
                    // Break out of sleep loop.
                    break;
                }
                // Otherwise, we need to wait longer, sleep 1 second.
                Thread.sleep(1000L);
            }
        } finally {
            // Close the stream consumer.
            kafkaStreams.close();
        }

        // Validation.
        Assertions.assertEquals(numberOfRecords, recordCounter.get(), "Should have 25 records processed");

        // Consume records from output topic.
        final List<ConsumerRecord<String, String>> outputRecords =
            kafkaTestUtils.consumeAllRecordsFromTopic(outputTopic, StringDeserializer.class, StringDeserializer.class);

        // Validate we got the correct number of records.
        Assertions.assertEquals(numberOfRecords, outputRecords.size());
    }
}
 

/**
 * Main entry point
 * @param args command line args.
 */
public static void main(final String[] args) throws Exception {
    // Parse command line arguments
    final CommandLine cmd = parseArguments(args);

    // Right now we accept one parameter, the number of nodes in the cluster.
    final int clusterSize = Integer.parseInt(cmd.getOptionValue("size"));
    logger.info("Starting up kafka cluster with {} brokers", clusterSize);

    // Default to plaintext listener.
    BrokerListener listener = new PlainListener();

    final URL trustStore = DevCluster.class.getClassLoader().getResource("kafka.truststore.jks");
    final URL keyStore = DevCluster.class.getClassLoader().getResource("kafka.keystore.jks");

    final Properties properties = new Properties();
    if (cmd.hasOption("sasl") && cmd.hasOption("ssl")) {
        listener = new SaslSslListener()
            // SSL Options
            .withClientAuthRequired()
            .withTrustStoreLocation(trustStore.getFile())
            .withTrustStorePassword("password")
            .withKeyStoreLocation(keyStore.getFile())
            .withKeyStorePassword("password")
            .withKeyPassword("password")
            // SASL Options.
            .withUsername("kafkaclient")
            .withPassword("client-secret");
    } else if (cmd.hasOption("sasl")) {
        listener = new SaslPlainListener()
            .withUsername("kafkaclient")
            .withPassword("client-secret");
    } else if (cmd.hasOption("ssl")) {
        listener = new SslListener()
            .withClientAuthRequired()
            .withTrustStoreLocation(trustStore.getFile())
            .withTrustStorePassword("password")
            .withKeyStoreLocation(keyStore.getFile())
            .withKeyStorePassword("password")
            .withKeyPassword("password");
    }

    // Create a test cluster
    final KafkaTestCluster kafkaTestCluster = new KafkaTestCluster(
        clusterSize,
        properties,
        Collections.singletonList(listener)
    );

    // Start the cluster.
    kafkaTestCluster.start();

    // Create topics
    String[] topicNames = null;
    if (cmd.hasOption("topic")) {
        topicNames = cmd.getOptionValues("topic");

        for (final String topicName : topicNames) {
            final KafkaTestUtils utils = new KafkaTestUtils(kafkaTestCluster);
            utils.createTopic(topicName, clusterSize, (short) clusterSize);

            // Publish some data into that topic
            for (int partition = 0; partition < clusterSize; partition++) {
                utils.produceRecords(1000, topicName, partition);
            }
        }
    }

    // Log topic names created.
    if (topicNames != null) {
        logger.info("Created topics: {}", String.join(", ", topicNames));
    }

    // Log how to connect to cluster brokers.
    kafkaTestCluster
        .getKafkaBrokers()
        .stream()
        .forEach((broker) -> {
            logger.info("Started broker with Id {} at {}", broker.getBrokerId(), broker.getConnectString());
        });

    if (topicNames != null && topicNames.length > 0) {
        final KafkaTestUtils testUtils = new KafkaTestUtils(kafkaTestCluster);
        if (cmd.hasOption("consumer")) {
            runEndlessConsumer(Arrays.asList(topicNames), testUtils);
        }

        if (cmd.hasOption("producer")) {
            runEndlessProducer(Arrays.asList(topicNames), testUtils);
        }
    }

    // Log cluster connect string.
    logger.info("Cluster started at: {}", kafkaTestCluster.getKafkaConnectString());

    // Wait forever.
    Thread.currentThread().join();
}