kafkatest: add new command line arguments

edenhill · edenhill · commit cb5f882f44dc · 2020-04-03T16:52:07.000+02:00
diff --git a/confluent_kafka/kafkatest/verifiable_consumer.py b/confluent_kafka/kafkatest/verifiable_consumer.py
@@ -239,7 +239,9 @@ def to_dict(self):
     parser = argparse.ArgumentParser(description='Verifiable Python Consumer')
     parser.add_argument('--topic', action='append', type=str, required=True)
     parser.add_argument('--group-id', dest='conf_group.id', required=True)
+    parser.add_argument('--group-instance-id', dest='conf_group.instance.id')
     parser.add_argument('--broker-list', dest='conf_bootstrap.servers', required=True)
+    parser.add_argument('--bootstrap-server', dest='conf_bootstrap.servers')
     parser.add_argument('--session-timeout', type=int, dest='conf_session.timeout.ms', default=6000)
     parser.add_argument('--enable-autocommit', action='store_true', dest='conf_enable.auto.commit', default=False)
     parser.add_argument('--max-messages', type=int, dest='max_messages', default=-1)
diff --git a/confluent_kafka/kafkatest/verifiable_producer.py b/confluent_kafka/kafkatest/verifiable_producer.py
@@ -64,10 +64,12 @@ def dr_cb(self, err, msg):
     parser.add_argument('--topic', type=str, required=True)
     parser.add_argument('--throughput', type=int, default=0)
     parser.add_argument('--broker-list', dest='conf_bootstrap.servers', required=True)
+    parser.add_argument('--bootstrap-server', dest='conf_bootstrap.servers')
     parser.add_argument('--max-messages', type=int, dest='max_msgs', default=1000000)  # avoid infinite
     parser.add_argument('--value-prefix', dest='value_prefix', type=str, default=None)
     parser.add_argument('--acks', type=int, dest='topicconf_request.required.acks', default=-1)
     parser.add_argument('--message-create-time', type=int, dest='create_time', default=0)
+    parser.add_argument('--repeating-keys', type=int, dest='repeating_keys', default=0)
     parser.add_argument('--producer.config', dest='producer_config')
     parser.add_argument('-X', nargs=1, dest='extra_conf', action='append', help='Configuration property', default=[])
     args = vars(parser.parse_args())
@@ -92,6 +94,9 @@ def dr_cb(self, err, msg):
     else:
         value_fmt = '%d'
 
+    repeating_keys = args['repeating_keys']
+    key_counter = 0
+
     if throughput > 0:
         delay = 1.0/throughput
     else:
@@ -106,8 +111,14 @@ def dr_cb(self, err, msg):
 
             t_end = time.time() + delay
             while vp.run:
+                if repeating_keys != 0:
+                    key = '%d' % key_counter
+                    key_counter = (key_counter + 1) % repeating_keys
+                else:
+                    key = None
+
                 try:
-                    vp.producer.produce(topic, value=(value_fmt % i),
+                    vp.producer.produce(topic, value=(value_fmt % i), key=key,
                                         timestamp=args.get('create_time', 0))
                     vp.num_sent += 1
                 except KafkaException as e:
