diff options
Diffstat (limited to 'test/test_partitioner.py')
| -rw-r--r-- | test/test_partitioner.py | 33 |
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') |