def test_decide_partitions(self):
"""Test partition assignment for a number of partitions/consumers."""
consumer_group = 'testgroup'
# 100 test iterations
for i in xrange(100):
# Set up partitions, cluster, etc
num_participants = i + 1
num_partitions = 100 - i
participants = ['test-debian:{p}'.format(p=p)
for p in xrange(num_participants)]
topic = mock.Mock()
topic.name = 'testtopic'
topic.partitions = {}
for k in xrange(num_partitions):
part = mock.Mock(name='part-{part}'.format(part=k))
part.id = k
part.topic = topic
part.leader = mock.Mock()
part.leader.id = k % num_participants
topic.partitions[k] = part
cluster = mock.MagicMock()
zk = mock.MagicMock()
cns = BalancedConsumer(topic, cluster, consumer_group,
zookeeper=zk, auto_start=False)
# Simulate each participant to ensure they're correct
assigned_parts = []
for p_id in xrange(num_participants):
cns._consumer_id = participants[p_id] # override consumer id
# Decide partitions then validate
partitions = cns._decide_partitions(participants)
assigned_parts.extend(partitions)
remainder_ppc = num_partitions % num_participants
idx = participants.index(cns._consumer_id)
parts_per_consumer = num_partitions / num_participants
parts_per_consumer = math.floor(parts_per_consumer)
num_parts = parts_per_consumer + (0 if (idx + 1 > remainder_ppc) else 1)
self.assertEqual(len(partitions), num_parts)
# Validate all partitions were assigned once and only once
all_partitions = topic.partitions.values()
all_partitions.sort()
assigned_parts.sort()
self.assertListEqual(assigned_parts, all_partitions)
def buildMockConsumer(self,
consumer_group=b'testgroup',
num_partitions=10,
num_participants=1,
timeout=2000):
topic = mock.Mock()
topic.name = 'testtopic'
topic.partitions = {}
for k in range(num_partitions):
part = mock.Mock(name='part-{part}'.format(part=k))
part.id = k
part.topic = topic
part.leader = mock.Mock()
part.leader.id = k % num_participants
topic.partitions[k] = part
cluster = mock.MagicMock()
zk = mock.MagicMock()
return BalancedConsumer(topic,
cluster,
consumer_group,
zookeeper=zk,
auto_start=False,
use_rdkafka=False,
consumer_timeout_ms=timeout), topic
def test_decide_partitions(self):
for i in xrange(100):
num_participants = i + 1
num_partitions = 100 - i
consumer_group = 'testgroup'
topic = mock.Mock()
topic.name = 'testtopic'
topic.partitions = {}
for k in xrange(num_partitions):
part = mock.Mock()
part.id = k
part.topic = topic
part.leader = mock.Mock()
part.leader.id = k % num_participants
topic.partitions[k] = part
cluster = mock.MagicMock()
zk = mock.MagicMock()
cns = BalancedConsumer(topic,
cluster,
consumer_group,
zookeeper=zk,
auto_start=False)
participants = [
'test-debian:{}'.format(uuid4())
for i in xrange(num_participants - 1)
]
participants.append(cns._consumer_id)
participants.sort()
partitions = cns._decide_partitions(participants)
remainder_ppc = num_partitions % num_participants
idx = participants.index(cns._consumer_id)
parts_per_consumer = num_partitions / num_participants
num_parts = parts_per_consumer + (0 if
(idx + 1 > remainder_ppc) else 1)
self.assertEqual(len(partitions), num_parts)
def get_balanced_consumer(self, consumer_group, **kwargs):
"""Return a BalancedConsumer of this topic
"""
return BalancedConsumer(self, self._cluster, consumer_group, **kwargs)