/*
 * 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.tools;

import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.errors.AuthorizationException;
import org.apache.kafka.common.utils.Utils;

import net.sourceforge.argparse4j.inf.ArgumentParser;
import net.sourceforge.argparse4j.inf.ArgumentParserException;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;

import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.SplittableRandom;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.STRICT_STUBS)
public class ProducerPerformanceTest {

    @Mock
    KafkaProducer<byte[], byte[]> producerMock;

    @Spy
    ProducerPerformance producerPerformanceSpy;

    private File createTempFile(String contents) throws IOException {
        File file = File.createTempFile("ProducerPerformanceTest", ".tmp");
        file.deleteOnExit();
        Files.write(file.toPath(), contents.getBytes());
        return file;
    }

    @Test
    public void testReadPayloadFile() throws Exception {
        File payloadFile = createTempFile("Hello\nKafka");
        String payloadFilePath = payloadFile.getAbsolutePath();
        String payloadDelimiter = "\n";

        List<byte[]> payloadByteList = ProducerPerformance.readPayloadFile(payloadFilePath, payloadDelimiter);

        assertEquals(2, payloadByteList.size());
        assertEquals("Hello", new String(payloadByteList.get(0)));
        assertEquals("Kafka", new String(payloadByteList.get(1)));
        Utils.delete(payloadFile);
    }

    @Test
    public void testReadProps() throws Exception {
        List<String> producerProps = List.of("bootstrap.servers=localhost:9000");
        File producerConfig = createTempFile("acks=1");

        Properties prop = ProducerPerformance.readProps(producerProps, producerConfig.getAbsolutePath());

        assertNotNull(prop);
        assertEquals(5, prop.size());
        Utils.delete(producerConfig);
    }

    @Test
    public void testReadPayloadFileWithAlternateDelimiters() throws Exception {
        List<byte[]> payloadByteList;

        payloadByteList = generateListFromFileUsingDelimiter("Hello~~Kafka", "~~");
        assertEquals(2, payloadByteList.size());
        assertEquals("Hello", new String(payloadByteList.get(0)));
        assertEquals("Kafka", new String(payloadByteList.get(1)));

        payloadByteList = generateListFromFileUsingDelimiter("Hello,Kafka,", ",");
        assertEquals(2, payloadByteList.size());
        assertEquals("Hello", new String(payloadByteList.get(0)));
        assertEquals("Kafka", new String(payloadByteList.get(1)));

        payloadByteList = generateListFromFileUsingDelimiter("Hello\t\tKafka", "\t");
        assertEquals(3, payloadByteList.size());
        assertEquals("Hello", new String(payloadByteList.get(0)));
        assertEquals("Kafka", new String(payloadByteList.get(2)));

        payloadByteList = generateListFromFileUsingDelimiter("Hello\n\nKafka\n", "\n");
        assertEquals(3, payloadByteList.size());
        assertEquals("Hello", new String(payloadByteList.get(0)));
        assertEquals("Kafka", new String(payloadByteList.get(2)));

        payloadByteList = generateListFromFileUsingDelimiter("Hello::Kafka::World", "\\s*::\\s*");
        assertEquals(3, payloadByteList.size());
        assertEquals("Hello", new String(payloadByteList.get(0)));
        assertEquals("Kafka", new String(payloadByteList.get(1)));

    }

    @Test
    public void testCompareStringSplitWithScannerDelimiter() throws Exception {

        String contents = "Hello~~Kafka";
        String payloadDelimiter = "~~";
        compareList(generateListFromFileUsingDelimiter(contents, payloadDelimiter), contents.split(payloadDelimiter));

        contents = "Hello,Kafka,";
        payloadDelimiter = ",";
        compareList(generateListFromFileUsingDelimiter(contents, payloadDelimiter), contents.split(payloadDelimiter));

        contents = "Hello\t\tKafka";
        payloadDelimiter = "\t";
        compareList(generateListFromFileUsingDelimiter(contents, payloadDelimiter), contents.split(payloadDelimiter));

        contents = "Hello\n\nKafka\n";
        payloadDelimiter = "\n";
        compareList(generateListFromFileUsingDelimiter(contents, payloadDelimiter), contents.split(payloadDelimiter));

        contents = "Hello::Kafka::World";
        payloadDelimiter = "\\s*::\\s*";
        compareList(generateListFromFileUsingDelimiter(contents, payloadDelimiter), contents.split(payloadDelimiter));

    }

    private void compareList(List<byte[]> payloadByteList, String[] payloadByteListFromSplit) {
        assertEquals(payloadByteListFromSplit.length, payloadByteList.size());
        for (int i = 0; i < payloadByteListFromSplit.length; i++) {
            assertEquals(payloadByteListFromSplit[i], new String(payloadByteList.get(i)));
        }
    }

    private List<byte[]> generateListFromFileUsingDelimiter(String fileContent, String payloadDelimiter) throws Exception {
        File payloadFile = null;
        List<byte[]> payloadByteList;
        try {
            payloadFile = createTempFile(fileContent);
            payloadByteList = ProducerPerformance.readPayloadFile(payloadFile.getAbsolutePath(), payloadDelimiter);
        } finally {
            Utils.delete(payloadFile);
        }
        return payloadByteList;
    }

    @Test
    public void testNumberOfCallsForSendAndClose() throws IOException {
        doReturn(null).when(producerMock).send(any(), any());
        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));

        String[] args = new String[] {
            "--topic", "Hello-Kafka", 
            "--num-records", "5", 
            "--throughput", "100", 
            "--record-size", "100", 
            "--bootstrap-server", "localhost:9000"};
        producerPerformanceSpy.start(args);
        verify(producerMock, times(5)).send(any(), any());
        verify(producerMock, times(1)).close();
    }

    @Test
    public void testEnableTransaction() throws IOException {
        doReturn(null).when(producerMock).send(any(), any());
        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));

        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--transactional-id", "foobar",
            "--bootstrap-server", "localhost:9000"};
        producerPerformanceSpy.start(args);
        verify(producerMock, times(1)).beginTransaction();
        verify(producerMock, times(1)).commitTransaction();
        verify(producerMock, times(1)).close();
    }

    @Test
    public void testNumberOfSuccessfulSendAndClose() throws IOException {
        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
        doAnswer(invocation -> {
            producerPerformanceSpy.cb.onCompletion(null, null);
            return null;
        }).when(producerMock).send(any(), any());

        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "10",
            "--throughput", "1",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        producerPerformanceSpy.start(args);

        verify(producerMock, times(10)).send(any(), any());
        assertEquals(10, producerPerformanceSpy.stats.totalCount());
        verify(producerMock, times(1)).close();
    }

    @Test
    public void testNumberOfFailedSendAndClose() throws IOException {
        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
        doAnswer(invocation -> {
            producerPerformanceSpy.cb.onCompletion(null, new AuthorizationException("not authorized."));
            return null;
        }).when(producerMock).send(any(), any());

        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "10",
            "--throughput", "1",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        producerPerformanceSpy.start(args);

        verify(producerMock, times(10)).send(any(), any());
        assertEquals(0, producerPerformanceSpy.stats.currentWindowCount());
        assertEquals(0, producerPerformanceSpy.stats.totalCount());
        verify(producerMock, times(1)).close();
    }

    @Test
    public void testMutuallyExclusiveGroup() {
        String[] args1 = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--payload-monotonic",
            "--bootstrap-server", "localhost:9000"};
        ArgumentParser parser1 = ProducerPerformance.argParser();
        ArgumentParserException thrown = assertThrows(ArgumentParserException.class, () ->  parser1.parseArgs(args1));
        assertEquals("argument --payload-monotonic: not allowed with argument --record-size", thrown.getMessage());

        String[] args2 = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--payload-file",  "abc.txt",
            "--payload-monotonic",
            "--bootstrap-server", "localhost:9000"};
        ArgumentParser parser2 = ProducerPerformance.argParser();
        thrown = assertThrows(ArgumentParserException.class, () -> parser2.parseArgs(args2));
        assertEquals("argument --payload-monotonic: not allowed with argument --payload-file", thrown.getMessage());
    }

    @Test
    public void testUnexpectedArg() {
        String[] args = new String[] {
            "--test", "test", 
            "--topic", "Hello-Kafka", 
            "--num-records", "5", 
            "--throughput", "100", 
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        ArgumentParser parser = ProducerPerformance.argParser();
        ArgumentParserException thrown = assertThrows(ArgumentParserException.class, () -> parser.parseArgs(args));
        assertEquals("unrecognized arguments: '--test'", thrown.getMessage());
    }

    @Test
    public void testFractionalThroughput() {
        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "1.25",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        ArgumentParser parser = ProducerPerformance.argParser();
        assertDoesNotThrow(() -> parser.parseArgs(args));
    }

    @Test
    public void testGenerateRandomPayloadByPayloadFile() {
        Integer recordSize = null;
        String inputString = "Hello Kafka";
        byte[] byteArray = inputString.getBytes(StandardCharsets.UTF_8);
        List<byte[]> payloadByteList = new ArrayList<>();
        payloadByteList.add(byteArray);
        byte[] payload = null;
        SplittableRandom random = new SplittableRandom(0);

        payload = ProducerPerformance.generateRandomPayload(recordSize, payloadByteList, payload, random, false, 0L);
        assertEquals(inputString, new String(payload));
    }

    @Test
    public void testGenerateRandomPayloadByRecordSize() {
        Integer recordSize = 100;
        byte[] payload = new byte[recordSize];
        List<byte[]> payloadByteList = new ArrayList<>();
        SplittableRandom random = new SplittableRandom(0);

        payload = ProducerPerformance.generateRandomPayload(recordSize, payloadByteList, payload, random, false, 0L);
        for (byte b : payload) {
            assertNotEquals(0, b);
        }
    }

    @Test
    public void testGenerateMonotonicPayload() {
        byte[] payload = null;
        List<byte[]> payloadByteList = new ArrayList<>();
        SplittableRandom random = new SplittableRandom(0);

        for (int i = 0; i < 10; i++) {
            payload = ProducerPerformance.generateRandomPayload(null, payloadByteList, payload, random, true, i);
            assertEquals(1, payload.length);
            assertEquals(i + '0', payload[0]);
        }
    }

    @Test
    public void testGenerateRandomPayloadException() {
        Integer recordSize = null;
        byte[] payload = null;
        List<byte[]> payloadByteList = new ArrayList<>();
        SplittableRandom random = new SplittableRandom(0);

        IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, () -> ProducerPerformance.generateRandomPayload(recordSize, payloadByteList, payload, random, false, 0L));
        assertEquals("No payload file, record size or payload-monotonic option provided.", thrown.getMessage());
    }

    @Test
    public void testClientIdOverride()  throws Exception {
        List<String> producerProps = List.of("client.id=producer-1");

        Properties prop = ProducerPerformance.readProps(producerProps, null);

        assertNotNull(prop);
        assertEquals("producer-1", prop.getProperty("client.id"));
    }

    @Test
    public void testDefaultClientId() throws Exception {
        List<String> producerProps = List.of("acks=1");

        Properties prop = ProducerPerformance.readProps(producerProps, null);

        assertNotNull(prop);
        assertEquals("perf-producer-client", prop.getProperty("client.id"));
    }

    @Test
    public void testStatsInitializationWithLargeNumRecords() {
        long numRecords = Long.MAX_VALUE;
        assertDoesNotThrow(() -> new ProducerPerformance.Stats(numRecords, 5000L, false));
    }

    @Test
    public void testStatsCorrectness() throws Exception {
        ExecutorService singleThreaded = Executors.newSingleThreadExecutor();
        final long numRecords = 1000000;
        ProducerPerformance.Stats stats = new ProducerPerformance.Stats(numRecords, 5000L, false);
        for (long i = 0; i < numRecords; i++) {
            final Callback callback = new ProducerPerformance.PerfCallback(0, 100, stats, null);
            CompletableFuture.runAsync(() -> callback.onCompletion(null, null), singleThreaded);
        }

        singleThreaded.shutdown();
        final boolean success = singleThreaded.awaitTermination(60, TimeUnit.SECONDS);

        assertTrue(success, "should have terminated");
        assertEquals(numRecords, stats.totalCount());
        assertEquals(numRecords, stats.iteration());
        assertEquals(500000, stats.index());
        assertEquals(1000000 * 100, stats.bytes());
    }

    @Test
    public void testConfigPostProcessor() throws IOException, ArgumentParserException {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--print-metrics",
            "--bootstrap-server", "localhost:9000",
            "--transactional-id", "foobar",
            "--transaction-duration-ms", "5000",
        };
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("localhost:9000", configs.bootstrapServer);
        assertEquals("Hello-Kafka", configs.topicName);
        assertEquals(5, configs.numRecords);
        assertEquals(100, configs.throughput);
        assertEquals(100, configs.recordSize);
        assertFalse(configs.payloadMonotonic);
        assertTrue(configs.shouldPrintMetrics);
        assertTrue(configs.payloadByteList.isEmpty());
        Properties props = configs.producerProps;
        assertEquals(5, props.size());
        assertTrue(configs.transactionsEnabled);
        assertEquals(5000, configs.transactionDurationMs);
    }

    @Test
    public void testInvalidConfigPostProcessor() {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] invalidProducerProps = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100"};
        assertEquals("At least one of --bootstrap-server, --command-property, --producer-props, --producer.config or --command-config must be specified.",
            assertThrows(ArgumentParserException.class,
                () -> new ProducerPerformance.ConfigPostProcessor(parser, invalidProducerProps)).getMessage());

        String[] invalidTransactionDurationMs = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--transaction-duration-ms", "0"};
        assertEquals("--transaction-duration-ms should be greater than zero.",
            assertThrows(ArgumentParserException.class,
                () -> new ProducerPerformance.ConfigPostProcessor(parser, invalidTransactionDurationMs)).getMessage());

        String[] invalidNumRecords = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "-5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        assertEquals("--num-records should be greater than zero.",
            assertThrows(ArgumentParserException.class,
                () -> new ProducerPerformance.ConfigPostProcessor(parser, invalidNumRecords)).getMessage());

        String[] invalidRecordSize = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "-100",
            "--bootstrap-server", "localhost:9000"};
        assertEquals("--record-size should be greater than zero.",
            assertThrows(ArgumentParserException.class,
                () -> new ProducerPerformance.ConfigPostProcessor(parser, invalidRecordSize)).getMessage());

        String[] invalidReportingInterval = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--reporting-interval", "0",
            "--bootstrap-server", "localhost:9000"};
        assertEquals("--reporting-interval should be greater than zero.",
            assertThrows(ArgumentParserException.class,
                () -> new ProducerPerformance.ConfigPostProcessor(parser, invalidReportingInterval)).getMessage());
    }

    @Test
    public void testBootstrapServer() throws IOException, ArgumentParserException {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("localhost:9000", configs.producerProps.get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--command-property", "bootstrap.servers=localhost:9001"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("localhost:9001", configs.producerProps.get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--command-property", "bootstrap.servers=localhost:9001"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("localhost:9000", configs.producerProps.get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--producer-props", "bootstrap.servers=localhost:9001"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("localhost:9001", configs.producerProps.get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--producer-props", "bootstrap.servers=localhost:9001"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("localhost:9000", configs.producerProps.get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG));
    }

    @Test
    public void testNoTransactionRelatedConfigs() throws IOException, ArgumentParserException {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertFalse(configs.transactionsEnabled);
        assertNull(configs.transactionDurationMs);
        assertFalse(configs.producerProps.contains(ProducerConfig.TRANSACTIONAL_ID_CONFIG));
    }

    @Test
    public void testEnableTransactionByProducerProperty() throws IOException, ArgumentParserException {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--command-property", "transactional.id=foobar"};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("foobar", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));
    }

    @Test
    public void testEnableTransactionByTransactionId() throws IOException, ArgumentParserException {
        File producerConfigFile = createTempFile("transactional.id=foobar");
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--command-config", producerConfigFile.getAbsolutePath()};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("foobar", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--command-config", producerConfigFile.getAbsolutePath(),
            "--command-property", "transactional.id=hello_kafka"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("hello_kafka", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--transactional-id", "kafka_hello",
            "--bootstrap-server", "localhost:9000",
            "--command-config", producerConfigFile.getAbsolutePath(),
            "--command-property", "transactional.id=hello_kafka"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("kafka_hello", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

        Utils.delete(producerConfigFile);
    }

    @Test
    public void testEnableTransactionByTransactionIdDeprecated() throws IOException, ArgumentParserException {
        File producerConfigFile = createTempFile("transactional.id=foobar");
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--producer.config", producerConfigFile.getAbsolutePath()};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("foobar", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            "--producer.config", producerConfigFile.getAbsolutePath(),
            "--producer-props", "transactional.id=hello_kafka"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("hello_kafka", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

        args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--transactional-id", "kafka_hello",
            "--bootstrap-server", "localhost:9000",
            "--producer.config", producerConfigFile.getAbsolutePath(),
            "--producer-props", "transactional.id=hello_kafka"};
        configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(ProducerPerformance.DEFAULT_TRANSACTION_DURATION_MS, configs.transactionDurationMs);
        assertEquals("kafka_hello", configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));

        Utils.delete(producerConfigFile);
    }

    @Test
    public void testEnsureDeprecatedAndModernArgumentsNotBothSpecified() throws IOException {
        File producerConfigFile = createTempFile("bootstrap.servers=localhost:9000");
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--producer.config", producerConfigFile.getAbsolutePath(),
            "--command-config", producerConfigFile.getAbsolutePath()};
        ArgumentParser parser = ProducerPerformance.argParser();
        assertThrows(ArgumentParserException.class, () -> new ProducerPerformance.ConfigPostProcessor(parser, args));

        String[] args2 = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--producer-props", "bootstrap.servers=localhost:9090",
            "--command-property", "bootstrap.servers=localhost:9090"};
        assertThrows(ArgumentParserException.class, () -> new ProducerPerformance.ConfigPostProcessor(parser, args2));

        Utils.delete(producerConfigFile);
    }

    @Test
    public void testEnableTransactionByTransactionDurationMs() throws IOException, ArgumentParserException {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--transaction-duration-ms", "5000",
            "--bootstrap-server", "localhost:9000"};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertTrue(configs.transactionsEnabled);
        assertEquals(5000, configs.transactionDurationMs);
        assertTrue(configs.producerProps.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG).toString()
                .startsWith(ProducerPerformance.DEFAULT_TRANSACTION_ID_PREFIX));
    }

    @Test
    public void testWarmupRecordsFractionalValue() {
        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "10",
            "--warmup-records", "1.5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        ArgumentParser parser = ProducerPerformance.argParser();
        ArgumentParserException thrown = assertThrows(ArgumentParserException.class, () -> parser.parseArgs(args));
        thrown.printStackTrace();
    }

    @Test
    public void testWarmupRecordsString() {
        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "10",
            "--warmup-records", "foo",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        ArgumentParser parser = ProducerPerformance.argParser();
        ArgumentParserException thrown = assertThrows(ArgumentParserException.class, () -> parser.parseArgs(args));
        thrown.printStackTrace();
    }

    @Test
    public void testWarmupNumberOfSuccessfulSendAndClose() throws IOException {
        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
        doAnswer(invocation -> {
            producerPerformanceSpy.cb.onCompletion(null, null);
            return null;
        }).when(producerMock).send(any(), any());

        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "10",
            "--warmup-records", "2",
            "--throughput", "1",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        producerPerformanceSpy.start(args);

        verify(producerMock, times(10)).send(any(), any());
        assertEquals(10, producerPerformanceSpy.stats.totalCount());
        assertEquals(10 - 2, producerPerformanceSpy.steadyStateStats.totalCount());
        verify(producerMock, times(1)).close();
    }

    @Test
    public void testWarmupNegativeRecordsNormalTest() throws IOException {
        doReturn(producerMock).when(producerPerformanceSpy).createKafkaProducer(any(Properties.class));
        doAnswer(invocation -> {
            producerPerformanceSpy.cb.onCompletion(null, null);
            return null;
        }).when(producerMock).send(any(), any());

        String[] args = new String[] {
            "--topic", "Hello-Kafka",
            "--num-records", "10",
            "--warmup-records", "-1",
            "--throughput", "1",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000"};
        producerPerformanceSpy.start(args);

        verify(producerMock, times(10)).send(any(), any());
        assertEquals(10, producerPerformanceSpy.stats.totalCount());
        verify(producerMock, times(1)).close();
    }

    @ParameterizedTest
    @ValueSource(strings = {"--command-property", "--producer-props"})
    public void testMultipleProducerProperty(String configKey) throws IOException, ArgumentParserException {
        ArgumentParser parser = ProducerPerformance.argParser();
        String[] args = new String[]{
            "--topic", "Hello-Kafka",
            "--num-records", "5",
            "--throughput", "100",
            "--record-size", "100",
            "--bootstrap-server", "localhost:9000",
            configKey, "linger.ms=10", "batch.size=32768"};
        ProducerPerformance.ConfigPostProcessor configs = new ProducerPerformance.ConfigPostProcessor(parser, args);
        assertEquals("10", configs.producerProps.get(ProducerConfig.LINGER_MS_CONFIG));
        assertEquals("32768", configs.producerProps.get(ProducerConfig.BATCH_SIZE_CONFIG));
    }
}
