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)