Minor aesthetic cleanup of partitioner tests

1 files changed, 10 insertions, 11 deletions

diff --git a/test/test_partitioner.py b/test/test_partitioner.py
index 47470e1..3a5264b 100644
--- a/test/test_partitioner.py
+++ b/test/test_partitioner.py
@@ -1,13 +1,14 @@
 from __future__ import absolute_import
 
+import pytest
+
 from kafka.partitioner import DefaultPartitioner, Murmur2Partitioner, RoundRobinPartitioner
 from kafka.partitioner.hashed import murmur2
 
 
 def test_default_partitioner():
     partitioner = DefaultPartitioner()
-    all_partitions = list(range(100))
-    available = all_partitions
+    all_partitions = available = list(range(100))
     # partitioner should return the same partition for the same key
     p1 = partitioner(b'foo', all_partitions, available)
     p2 = partitioner(b'foo', all_partitions, available)
@@ -23,8 +24,7 @@ def test_default_partitioner():
 
 def test_roundrobin_partitioner():
     partitioner = RoundRobinPartitioner()
-    all_partitions = list(range(100))
-    available = all_partitions
+    all_partitions = available = list(range(100))
     # partitioner should cycle between partitions
     i = 0
     max_partition = all_partitions[len(all_partitions) - 1]
@@ -53,15 +53,14 @@ def test_roundrobin_partitioner():
         i += 1
 
 
-def test_murmur2_java_compatibility():
+@pytest.mark.parametrize("bytes_payload,partition_number", [
+    (b'', 681), (b'a', 524), (b'ab', 434), (b'abc', 107), (b'123456789', 566),
+    (b'\x00 ', 742)
+])
+def test_murmur2_java_compatibility(bytes_payload, partition_number):
     p = Murmur2Partitioner(range(1000))
     # compare with output from Kafka's org.apache.kafka.clients.producer.Partitioner
-    assert p.partition(b'') == 681
-    assert p.partition(b'a') == 524
-    assert p.partition(b'ab') == 434
-    assert p.partition(b'abc') == 107
-    assert p.partition(b'123456789') == 566
-    assert p.partition(b'\x00 ') == 742
+    assert p.partition(bytes_payload) == partition_number
 
 
 def test_murmur2_not_ascii():