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))