def test_remove_zone(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'}, '7': {'name': 'node7', 'zone': 'a'}, '8': {'name': 'node8', 'zone': 'a'}, '9': {'name': 'node9', 'zone': 'b'}, '10': {'name': 'node10', 'zone': 'b'}, '11': {'name': 'node11', 'zone': 'c'}, '12': {'name': 'node12', 'zone': 'c'}, } cluster = Cluster(cluster_config) placements = {} for i in cluster.nodes: placements[i] = [] for i in range(10000): nodes = cluster.find_nodes(str(i)) for node in nodes: placements[node].append(i) cluster.remove_node('5', node_name='node5', node_zone='c') cluster.remove_node('6', node_name='node6', node_zone='c') cluster.remove_node('11', node_name='node11', node_zone='c') cluster.remove_node('12', node_name='node12', node_zone='c') self.assertRaises(ValueError, cluster.remove_node, '12', 'node12', 'c') new_placements = {} for i in cluster.nodes: new_placements[i] = [] for i in range(10000): 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 placements.items(): after = set(assignments) before = set(new_placements.get(node, [])) removed += len(before.difference(after)) added += len(after.difference(before)) self.assertEqual(added, removed) self.assertEqual(13332, (added + removed))
def test_find_nodes(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) self.assertEqual(['2', '3'], cluster.find_nodes('lol')) self.assertEqual(['6', '2'], cluster.find_nodes('wat')) self.assertEqual(['5', '2'], cluster.find_nodes('ok')) self.assertEqual(['1', '3'], cluster.find_nodes('bar')) self.assertEqual(['1', '3'], cluster.find_nodes('foo')) self.assertEqual(['2', '4'], cluster.find_nodes('slap'))