def test_grow(self):
cluster_config = {
'1': {'name': 'node1', 'zone': 'a'},
'2': {'name': 'node2', 'zone': 'a'},
'3': {'name': 'node3', 'zone': 'b'},
'4': {'name': 'node4', 'zone': 'b'},
'5': {'name': 'node5', 'zone': 'c'},
'6': {'name': 'node6', 'zone': 'c'},
}
cluster = Cluster(cluster_config)
placements = {}
for i in cluster.nodes:
placements[i] = []
for i in range(1000):
nodes = cluster.find_nodes(str(i))
for node in nodes:
placements[node].append(i)
cluster.add_node('7', node_name='node7', node_zone='a')
cluster.add_node('8', node_name='node8', node_zone='b')
cluster.add_node('9', node_name='node9', node_zone='c')
new_placements = {}
for i in cluster.nodes:
new_placements[i] = []
for i in range(1000):
nodes = cluster.find_nodes(str(i))
for node in nodes:
new_placements[node].append(i)
keys = [k for sublist in placements.values() for k in sublist]
new_keys = [k for sublist in new_placements.values() for k in sublist]
self.assertEqual(sorted(keys), sorted(new_keys))
added = 0
removed = 0
for node, assignments in new_placements.items():
after = set(assignments)
before = set(placements.get(node, []))
removed += len(before.difference(after))
added += len(after.difference(before))
self.assertEqual(added, removed)
self.assertEqual(1384, (added + removed))
def test_add_node(self):
cluster = Cluster()
self.assertEqual(0, len(cluster.nodes))
self.assertEqual([], cluster.zones)
self.assertEqual(0, len(cluster.zone_members))
self.assertEqual(0, len(cluster.rings))
cluster.add_node('2', node_zone='b')
self.assertEqual(1, len(cluster.nodes))
self.assertEqual(['b'], cluster.zones)
self.assertEqual(1, len(cluster.zone_members))
self.assertEqual(['2'], sorted(cluster.zone_members['b']))
self.assertEqual(1, len(cluster.rings))
cluster.add_node('1', node_zone='a')
self.assertEqual(2, len(cluster.nodes))
self.assertEqual(['a', 'b'], cluster.zones)
self.assertEqual(2, len(cluster.zone_members))
self.assertEqual(['1'], sorted(cluster.zone_members['a']))
self.assertEqual(['2'], sorted(cluster.zone_members['b']))
self.assertEqual(2, len(cluster.rings))
cluster.add_node('21', node_zone='b')
self.assertEqual(3, len(cluster.nodes))
self.assertEqual(['a', 'b'], cluster.zones)
self.assertEqual(2, len(cluster.zone_members))
self.assertEqual(['1'], sorted(cluster.zone_members['a']))
self.assertEqual(['2', '21'], sorted(cluster.zone_members['b']))
self.assertEqual(2, len(cluster.rings))
self.assertRaises(ValueError, cluster.add_node, '21')
self.assertRaises(ValueError, cluster.add_node, '21', None, None)
self.assertRaises(ValueError, cluster.add_node, '21', None, 'b')
cluster.add_node('22', node_zone='c')
self.assertEqual(4, len(cluster.nodes))
self.assertEqual(['a', 'b', 'c'], cluster.zones)
self.assertEqual(3, len(cluster.zone_members))
self.assertEqual(['1'], sorted(cluster.zone_members['a']))
self.assertEqual(['2', '21'], sorted(cluster.zone_members['b']))
self.assertEqual(['22'], sorted(cluster.zone_members['c']))
self.assertEqual(3, len(cluster.rings))
