Added Test*Consumer classes instead of functions

Author: pranavrth

tests/common/__init__.py

@@ -36,20 +36,24 @@ def use_kraft():
     return use_group_protocol_consumer() or _trivup_cluster_type_kraft()
 
 
-def _get_consumer_generic(consumer_clazz, conf=None, **kwargs):
-    if use_group_protocol_consumer():
-        if conf is not None and 'group.id' in conf:
-            conf['group.protocol'] = 'consumer'
-    return consumer_clazz(conf, **kwargs)
+def _update_conf_group_protocol(conf=None):
+    if conf is not None and 'group.id' in conf and use_group_protocol_consumer():
+        conf['group.protocol'] = 'consumer'
 
 
-def get_consumer(conf=None, **kwargs):
-    return _get_consumer_generic(Consumer, conf, **kwargs)
+class TestConsumer(Consumer):
+    def __init__(self, conf, **kwargs):
+        _update_conf_group_protocol(conf)
+        super(TestConsumer, self).__init__(conf, **kwargs)
 
 
-def get_avro_consumer(conf=None, **kwargs):
-    return _get_consumer_generic(AvroConsumer, conf, **kwargs)
+class TestDeserializingConsumer(DeserializingConsumer):
+    def __init__(self, conf, **kwargs):
+        _update_conf_group_protocol(conf)
+        super(TestDeserializingConsumer, self).__init__(conf, **kwargs)
 
 
-def get_deserializing_consumer(conf=None, **kwargs):
-    return _get_consumer_generic(DeserializingConsumer, conf, **kwargs)
+class TestAvroConsumer(AvroConsumer):
+    def __init__(self, conf, **kwargs):
+        _update_conf_group_protocol(conf)
+        super(TestAvroConsumer, self).__init__(conf, **kwargs)

tests/integration/cluster_fixture.py

@@ -23,7 +23,7 @@
 from confluent_kafka import Producer, SerializingProducer
 from confluent_kafka.admin import AdminClient, NewTopic
 from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
-from ..common import get_consumer, get_deserializing_consumer
+from ..common import TestDeserializingConsumer, TestConsumer
 
 
 class KafkaClusterFixture(object):
@@ -105,7 +105,7 @@ def cimpl_consumer(self, conf=None):
         if conf is not None:
             consumer_conf.update(conf)
 
-        return get_consumer(consumer_conf)
+        return TestConsumer(consumer_conf)
 
     def consumer(self, conf=None, key_deserializer=None, value_deserializer=None):
         """
@@ -138,7 +138,7 @@ def consumer(self, conf=None, key_deserializer=None, value_deserializer=None):
         if value_deserializer is not None:
             consumer_conf['value.deserializer'] = value_deserializer
 
-        return get_deserializing_consumer(consumer_conf)
+        return TestDeserializingConsumer(consumer_conf)
 
     def admin(self, conf=None):
         if conf:

tests/integration/integration_test.py

@@ -29,7 +29,7 @@
 import gc
 import struct
 import re
-from ..common import get_consumer, get_avro_consumer
+from ..common import TestConsumer, TestAvroConsumer
 
 try:
     # Memory tracker
@@ -374,7 +374,7 @@ def verify_consumer():
             'enable.partition.eof': True}
 
     # Create consumer
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
 
     def print_wmark(consumer, topic_parts):
         # Verify #294: get_watermark_offsets() should not fail on the first call
@@ -484,7 +484,7 @@ def print_wmark(consumer, topic_parts):
     c.close()
 
     # Start a new client and get the committed offsets
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
     offsets = c.committed(list(map(lambda p: confluent_kafka.TopicPartition(topic, p), range(0, 3))))
     for tp in offsets:
         print(tp)
@@ -501,7 +501,7 @@ def verify_consumer_performance():
             'error_cb': error_cb,
             'auto.offset.reset': 'earliest'}
 
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
 
     def my_on_assign(consumer, partitions):
         print('on_assign:', len(partitions), 'partitions:')
@@ -609,7 +609,7 @@ def verify_batch_consumer():
             'auto.offset.reset': 'earliest'}
 
     # Create consumer
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
 
     # Subscribe to a list of topics
     c.subscribe([topic])
@@ -666,7 +666,7 @@ def verify_batch_consumer():
     c.close()
 
     # Start a new client and get the committed offsets
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
     offsets = c.committed(list(map(lambda p: confluent_kafka.TopicPartition(topic, p), range(0, 3))))
     for tp in offsets:
         print(tp)
@@ -683,7 +683,7 @@ def verify_batch_consumer_performance():
             'error_cb': error_cb,
             'auto.offset.reset': 'earliest'}
 
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
 
     def my_on_assign(consumer, partitions):
         print('on_assign:', len(partitions), 'partitions:')
@@ -878,7 +878,7 @@ def run_avro_loop(producer_conf, consumer_conf):
         p.produce(**combo)
     p.flush()
 
-    c = get_avro_consumer(consumer_conf)
+    c = TestAvroConsumer(consumer_conf)
     c.subscribe([(t['topic']) for t in combinations])
 
     msgcount = 0
@@ -990,7 +990,7 @@ def stats_cb(stats_json_str):
             'statistics.interval.ms': 200,
             'auto.offset.reset': 'earliest'}
 
-    c = get_consumer(conf)
+    c = TestConsumer(conf)
     c.subscribe([topic])
 
     max_msgcnt = 1000000
@@ -1117,7 +1117,7 @@ def verify_avro_explicit_read_schema():
         p.produce(topic=avro_topic, **combo)
     p.flush()
 
-    c = get_avro_consumer(consumer_conf, reader_key_schema=reader_schema, reader_value_schema=reader_schema)
+    c = TestAvroConsumer(consumer_conf, reader_key_schema=reader_schema, reader_value_schema=reader_schema)
     c.subscribe([avro_topic])
 
     msgcount = 0

tests/integration/producer/test_transactions.py

@@ -18,7 +18,7 @@
 import inspect
 import sys
 from uuid import uuid1
-from ...common import get_consumer
+from ...common import TestConsumer
 
 from confluent_kafka import KafkaError
 
@@ -115,7 +115,7 @@ def test_send_offsets_committed_transaction(kafka_cluster):
         'error_cb': error_cb
     }
     consumer_conf.update(kafka_cluster.client_conf())
-    consumer = get_consumer(consumer_conf)
+    consumer = TestConsumer(consumer_conf)
 
     kafka_cluster.seed_topic(input_topic)
     consumer.subscribe([input_topic])
@@ -205,7 +205,7 @@ def consume_committed(conf, topic):
                      'error_cb': prefixed_error_cb(called_by()), }
 
     consumer_conf.update(conf)
-    consumer = get_consumer(consumer_conf)
+    consumer = TestConsumer(consumer_conf)
     consumer.subscribe([topic])
 
     msg_cnt = read_all_msgs(consumer)

tests/soak/soakclient.py

@@ -30,7 +30,7 @@
 from confluent_kafka.admin import AdminClient, NewTopic
 from collections import defaultdict
 from builtins import int
-from common import get_consumer
+from common import TestConsumer
 import argparse
 import threading
 import time
@@ -445,7 +445,7 @@ def filter_config(conf, filter_out, strip_prefix):
         cconf['error_cb'] = self.consumer_error_cb
         cconf['on_commit'] = self.consumer_commit_cb
         self.logger.info("consumer: using group.id {}".format(cconf['group.id']))
-        self.consumer = get_consumer(cconf)
+        self.consumer = TestConsumer(cconf)
 
         # Create and start producer thread
         self.producer_thread = threading.Thread(target=self.producer_thread_main)

tests/test_Consumer.py

@@ -4,21 +4,21 @@
                              KafkaException, TIMESTAMP_NOT_AVAILABLE,
                              OFFSET_INVALID, libversion)
 import pytest
-from .common import get_consumer
+from .common import TestConsumer
 
 
 def test_basic_api():
     """ Basic API tests, these wont really do anything since there is no
         broker configured. """
 
     with pytest.raises(TypeError) as ex:
-        kc = get_consumer()
+        kc = TestConsumer()
     assert ex.match('expected configuration dict')
 
     def dummy_commit_cb(err, partitions):
         pass
 
-    kc = get_consumer({'group.id': 'test', 'socket.timeout.ms': '100',
+    kc = TestConsumer({'group.id': 'test', 'socket.timeout.ms': '100',
                        'session.timeout.ms': 1000,  # Avoid close() blocking too long
                        'on_commit': dummy_commit_cb})
 
@@ -120,7 +120,7 @@ def dummy_assign_revoke(consumer, partitions):
 def test_store_offsets():
     """ Basic store_offsets() tests """
 
-    c = get_consumer({'group.id': 'test',
+    c = TestConsumer({'group.id': 'test',
                       'enable.auto.commit': True,
                       'enable.auto.offset.store': False,
                       'socket.timeout.ms': 50,
@@ -162,7 +162,7 @@ def commit_cb(cs, err, ps):
 
     cs = CommitState('test', 2)
 
-    c = get_consumer({'group.id': 'x',
+    c = TestConsumer({'group.id': 'x',
                       'enable.auto.commit': False, 'socket.timeout.ms': 50,
                       'session.timeout.ms': 100,
                       'on_commit': lambda err, ps: commit_cb(cs, err, ps)})
@@ -197,7 +197,7 @@ def poll(self, somearg):
 @pytest.mark.skipif(libversion()[1] < 0x000b0000,
                     reason="requires librdkafka >=0.11.0")
 def test_offsets_for_times():
-    c = get_consumer({'group.id': 'test',
+    c = TestConsumer({'group.id': 'test',
                       'enable.auto.commit': True,
                       'enable.auto.offset.store': False,
                       'socket.timeout.ms': 50,
@@ -217,7 +217,7 @@ def test_offsets_for_times():
 def test_multiple_close_does_not_throw_exception():
     """ Calling Consumer.close() multiple times should not throw Runtime Exception
     """
-    c = get_consumer({'group.id': 'test',
+    c = TestConsumer({'group.id': 'test',
                       'enable.auto.commit': True,
                       'enable.auto.offset.store': False,
                       'socket.timeout.ms': 50,
@@ -233,7 +233,7 @@ def test_multiple_close_does_not_throw_exception():
 def test_any_method_after_close_throws_exception():
     """ Calling any consumer method after close should throw a RuntimeError
     """
-    c = get_consumer({'group.id': 'test',
+    c = TestConsumer({'group.id': 'test',
                       'enable.auto.commit': True,
                       'enable.auto.offset.store': False,
                       'socket.timeout.ms': 50,
@@ -297,7 +297,7 @@ def test_any_method_after_close_throws_exception():
 def test_calling_store_offsets_after_close_throws_erro():
     """ calling store_offset after close should throw RuntimeError """
 
-    c = get_consumer({'group.id': 'test',
+    c = TestConsumer({'group.id': 'test',
                       'enable.auto.commit': True,
                       'enable.auto.offset.store': False,
                       'socket.timeout.ms': 50,
@@ -320,5 +320,5 @@ def test_consumer_without_groupid():
     """ Consumer should raise exception if group.id is not set """
 
     with pytest.raises(ValueError) as ex:
-        get_consumer({'bootstrap.servers': "mybroker:9092"})
+        TestConsumer({'bootstrap.servers': "mybroker:9092"})
     assert ex.match('group.id must be set')

tests/test_Producer.py

@@ -5,7 +5,7 @@
 from confluent_kafka import Producer, KafkaError, KafkaException, \
     TopicPartition, libversion
 from struct import pack
-from .common import get_consumer
+from .common import TestConsumer
 
 
 def error_cb(err):
@@ -212,7 +212,7 @@ def test_transaction_api():
     assert ex.value.args[0].fatal() is False
     assert ex.value.args[0].txn_requires_abort() is False
 
-    consumer = get_consumer({"group.id": "testgroup"})
+    consumer = TestConsumer({"group.id": "testgroup"})
     group_metadata = consumer.consumer_group_metadata()
     consumer.close()
 

tests/test_log.py

@@ -4,7 +4,7 @@
 import confluent_kafka
 import confluent_kafka.avro
 import logging
-from .common import get_consumer, get_avro_consumer
+from .common import TestConsumer, TestAvroConsumer
 
 
 class CountingFilter(logging.Filter):
@@ -25,7 +25,7 @@ def test_logging_consumer():
     logger.setLevel(logging.DEBUG)
     f = CountingFilter('consumer')
     logger.addFilter(f)
-    kc = get_consumer({'group.id': 'test',
+    kc = TestConsumer({'group.id': 'test',
                        'debug': 'all'},
                       logger=logger)
     while f.cnt == 0:
@@ -44,10 +44,10 @@ def test_logging_avro_consumer():
     f = CountingFilter('avroconsumer')
     logger.addFilter(f)
 
-    kc = get_avro_consumer({'schema.registry.url': 'http://example.com',
-                            'group.id': 'test',
-                            'debug': 'all'},
-                           logger=logger)
+    kc = TestAvroConsumer({'schema.registry.url': 'http://example.com',
+                           'group.id': 'test',
+                           'debug': 'all'},
+                          logger=logger)
     while f.cnt == 0:
         kc.poll(timeout=0.5)
 
@@ -149,7 +149,7 @@ def test_consumer_logger_logging_in_given_format():
     handler.setFormatter(logging.Formatter('%(name)s Logger | %(message)s'))
     logger.addHandler(handler)
 
-    c = get_consumer(
+    c = TestConsumer(
         {"bootstrap.servers": "test", "group.id": "test", "logger": logger, "debug": "msg"})
     c.poll(0)