Update Partitioners for use with KafkaProducer (#827)

1 files changed, 33 insertions, 0 deletions

diff --git a/test/test_partitioner.py b/test/test_partitioner.py
index 52b6b81..e0398c6 100644
--- a/test/test_partitioner.py
+++ b/test/test_partitioner.py
@@ -3,6 +3,7 @@ import six
 
 from kafka.partitioner import Murmur2Partitioner
 from kafka.partitioner.default import DefaultPartitioner
+from kafka.partitioner import RoundRobinPartitioner
 
 
 def test_default_partitioner():
@@ -22,6 +23,38 @@ def test_default_partitioner():
     assert partitioner(None, all_partitions, []) in all_partitions
 
 
+def test_roundrobin_partitioner():
+    partitioner = RoundRobinPartitioner()
+    all_partitions = list(range(100))
+    available = all_partitions
+    # partitioner should cycle between partitions
+    i = 0
+    max_partition = all_partitions[len(all_partitions) - 1]
+    while i <= max_partition:
+        assert i == partitioner(None, all_partitions, available)
+        i += 1
+
+    i = 0
+    while i <= int(max_partition / 2):
+        assert i == partitioner(None, all_partitions, available)
+        i += 1
+
+    # test dynamic partition re-assignment
+    available = available[:-25]
+
+    while i <= max(available):
+        assert i == partitioner(None, all_partitions, available)
+        i += 1
+
+    all_partitions = list(range(200))
+    available = all_partitions
+
+    max_partition = all_partitions[len(all_partitions) - 1]
+    while i <= max_partition:
+        assert i == partitioner(None, all_partitions, available)
+        i += 1
+
+
 def test_hash_bytes():
     p = Murmur2Partitioner(range(1000))
     assert p.partition(bytearray(b'test')) == p.partition(b'test')