Merge pull request #380 from dpkp/multiprocess_consumer_partitions

Support optional partitions kwarg in MultiProcessConsumer

1 files changed, 11 insertions, 2 deletions

diff --git a/test/test_consumer.py b/test/test_consumer.py
index 7b8f370..a3d09a8 100644
--- a/test/test_consumer.py
+++ b/test/test_consumer.py
@@ -1,8 +1,8 @@
 
-from mock import MagicMock
+from mock import MagicMock, patch
 from . import unittest
 
-from kafka import SimpleConsumer, KafkaConsumer
+from kafka import SimpleConsumer, KafkaConsumer, MultiProcessConsumer
 from kafka.common import KafkaConfigurationError
 
 class TestKafkaConsumer(unittest.TestCase):
@@ -13,3 +13,12 @@ class TestKafkaConsumer(unittest.TestCase):
     def test_broker_list_required(self):
         with self.assertRaises(KafkaConfigurationError):
             KafkaConsumer()
+
+class TestMultiProcessConsumer(unittest.TestCase):
+    def test_partition_list(self):
+        client = MagicMock()
+        partitions = (0,)
+        with patch.object(MultiProcessConsumer, 'fetch_last_known_offsets') as fetch_last_known_offsets:
+            consumer = MultiProcessConsumer(client, 'testing-group', 'testing-topic', partitions=partitions)
+            self.assertEqual(fetch_last_known_offsets.call_args[0], (partitions,) )
+        self.assertEqual(client.get_partition_ids_for_topic.call_count, 0) # pylint: disable=no-member