diff options
| -rw-r--r-- | kafka/consumer/multiprocess.py | 3 | ||||
| -rw-r--r-- | test/test_consumer_integration.py | 8 |
2 files changed, 5 insertions, 6 deletions
diff --git a/kafka/consumer/multiprocess.py b/kafka/consumer/multiprocess.py index cdfaeeb..5bc04cc 100644 --- a/kafka/consumer/multiprocess.py +++ b/kafka/consumer/multiprocess.py @@ -109,7 +109,7 @@ class MultiProcessConsumer(Consumer): auto_commit_every_n=AUTO_COMMIT_MSG_COUNT, auto_commit_every_t=AUTO_COMMIT_INTERVAL, num_procs=1, partitions_per_proc=0, - simple_consumer_options=None): + **simple_consumer_options): # Initiate the base consumer class super(MultiProcessConsumer, self).__init__( @@ -150,6 +150,7 @@ class MultiProcessConsumer(Consumer): for chunk in chunks: options = {'partitions': list(chunk)} if simple_consumer_options: + simple_consumer_options.pop('partitions', None) options.update(simple_consumer_options) args = (client.copy(), group, topic, self.queue, diff --git a/test/test_consumer_integration.py b/test/test_consumer_integration.py index 7dd4a79..d3df56a 100644 --- a/test/test_consumer_integration.py +++ b/test/test_consumer_integration.py @@ -61,10 +61,8 @@ class TestConsumerIntegration(KafkaIntegrationTestCase): group = kwargs.pop('group', self.id().encode('utf-8')) topic = kwargs.pop('topic', self.topic) - if consumer_class == SimpleConsumer: + if consumer_class in [SimpleConsumer, MultiProcessConsumer]: kwargs.setdefault('iter_timeout', 0) - elif consumer_class == MultiProcessConsumer: - kwargs.setdefault('simple_consumer_options', {'iter_timeout': 0}) return consumer_class(self.client, group, topic, **kwargs) @@ -245,8 +243,8 @@ class TestConsumerIntegration(KafkaIntegrationTestCase): self.send_messages(0, range(0, 10)) self.send_messages(1, range(10, 20)) - consumer = MultiProcessConsumer(self.client, "group1", self.topic, auto_commit=False, - simple_consumer_options={'iter_timeout': 0}) + consumer = MultiProcessConsumer(self.client, "group1", self.topic, + auto_commit=False, iter_timeout=0) self.assertEqual(consumer.pending(), 20) self.assertEqual(consumer.pending(partitions=[0]), 10) |