def test_generate_sibling_count(self):
t1 = Topic('topic1', 2)
t2 = Topic('topic2', 2)
t3 = Topic('topic3', 1)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p12 = create_and_attach_partition(t1, 2)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p22 = create_and_attach_partition(t2, 2)
p30 = create_and_attach_partition(t3, 0)
p31 = create_and_attach_partition(t3, 1)
b1 = create_broker('b1', [p10, p11, p20, p21, p30, p31])
b2 = create_broker('b2', [p12, p21, p22])
b3 = create_broker('b3', [p10, p11, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
expected = {
b1: {
b2: {
t1: 1,
t2: 0,
t3: 2
},
b3: {
t1: 0,
t2: 1,
t3: 2
}
},
b2: {
b1: {
t1: -1,
t2: 0,
t3: -2
},
b3: {
t1: -1,
t2: 1,
t3: 0
}
},
b3: {
b1: {
t1: 0,
t2: -1,
t3: -2
},
b2: {
t1: 1,
t2: -1,
t3: 0
}
},
}
actual = rg.generate_sibling_distance()
assert dict(actual) == expected
def test_update_sibling_count(self):
t1 = Topic('topic1', 2)
t2 = Topic('topic2', 2)
t3 = Topic('topic3', 1)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p12 = create_and_attach_partition(t1, 2)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p22 = create_and_attach_partition(t2, 2)
p30 = create_and_attach_partition(t3, 0)
p31 = create_and_attach_partition(t3, 1)
b1 = create_broker('b1', [p10, p11, p20, p21, p30, p31])
b2 = create_broker('b2', [p12, p21, p22])
b3 = create_broker('b3', [p10, p11, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
sibling_distance = {
b2: {
b1: {
t1: -1,
t2: 0,
t3: -2
},
b3: {
t1: 1,
t2: -1,
t3: -2
}
},
}
# Move a p10 from b1 to b2
b1.move_partition(p10, b2)
# NOTE: b2: b1: t1: -1 -> 1 and b2: b3: t1: 1 -> 0
expected = {
b2: {
b1: {
t1: 1,
t2: 0,
t3: -2
},
b3: {
t1: 0,
t2: -1,
t3: -2
},
},
}
actual = rg.update_sibling_distance(sibling_distance, b2, t1)
assert dict(actual) == expected
def test_add_partition(self):
mock_partitions = set([
Mock(
spec=Partition,
replicas=[sentinel.r1, sentinel.r2],
),
Mock(
spec=Partition,
replicas=[sentinel.r4, sentinel.r3],
),
])
topic = Topic('t0', 2, mock_partitions)
new_partition = Mock(spec=Partition, replicas=[sentinel.r2])
topic.add_partition(new_partition)
assert topic.partitions == mock_partitions | set([new_partition])
def test_add_partition(self):
mock_partitions = set([
Mock(
spec=Partition,
replicas=[sentinel.r1, sentinel.r2],
),
Mock(
spec=Partition,
replicas=[sentinel.r4, sentinel.r3],
),
])
topic = Topic('t0', 2, mock_partitions)
new_partition = Mock(spec=Partition, replicas=[sentinel.r2])
topic.add_partition(new_partition)
assert topic.partitions == mock_partitions | set([new_partition])
def _add_partition(topic_id, partition_id, replication_factor=1):
if topic_id not in topics:
topics[topic_id] = Topic(topic_id, replication_factor)
topic = topics[topic_id]
partition = Partition(topic, partition_id)
topic.add_partition(partition)
return partition
def test_get_preferred_partition(self):
t1 = Topic('t1', 1)
t2 = Topic('t2', 1)
t3 = Topic('t3', 1)
# Naming convention: p10 -> topic t1 partition 0
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p30 = create_and_attach_partition(t3, 0)
source = create_broker('b1', [p10, p11, p20])
dest = create_broker('b2', [p21, p30])
sibling_distance = {t1: -2, t2: 0, t3: 1}
actual = source.get_preferred_partition(dest, sibling_distance)
assert actual in (p10, p11)
def test_get_preferred_partition_no_preferred(self):
t1 = Topic('t1', replication_factor=2)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
source = create_broker('b1', [p10, p11])
dest = create_broker('b2', [p10, p11])
sibling_distance = {t1: 0}
actual = source.get_preferred_partition(dest, sibling_distance)
assert actual is None
def topic(self):
return Topic('t0', 2, set([sentinel.p1, sentinel.p2]))