def setUp(self): self.cluster = Cluster() self.cluster.add_broker(Broker(1, "brokerhost1.example.com")) self.cluster.add_broker(Broker(2, "brokerhost2.example.com")) self.cluster.add_broker(Broker(3, "brokerhost3.example.com")) self.cluster.add_topic(Topic("testTopic1", 2)) self.cluster.add_topic(Topic("testTopic2", 2))
def test_check_partition_swappable_racks_ok(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(1, 'brokerhost3.example.com') b4 = Broker(1, 'brokerhost4.example.com') b2.rack = "a" b3.rack = "b" b4.rack = "b" replicas_a = [b1, b2] replicas_b = [b3, b4] assert check_partition_swappable(replicas_a, replicas_b, 0) is True assert check_partition_swappable(replicas_b, replicas_a, 0) is True
def test_try_pick_new_broker(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(3, 'brokerhost3.example.com') self.cluster.add_broker(b3) b3.rack = "b" action = ActionBalanceRackAware(self.args, self.cluster) # Firmly order the deque action._random_brokers = deque([b3, b1, b2]) newbroker = action._try_pick_new_broker(self.cluster.topics['testTopic1'].partitions[0], 1) assert newbroker == b3 assert action._random_brokers == deque([b1, b2, b3])
def set_up_cluster_4broker(): cluster = Cluster() cluster.add_broker(Broker(1, "brokerhost1.example.com")) cluster.add_broker(Broker(2, "brokerhost2.example.com")) cluster.add_broker(Broker(3, "brokerhost3.example.com")) cluster.add_broker(Broker(4, "brokerhost4.example.com")) cluster.brokers[1].rack = "a" cluster.brokers[2].rack = "a" cluster.brokers[3].rack = "b" cluster.brokers[4].rack = "b" cluster.add_topic(Topic("testTopic1", 4)) cluster.add_topic(Topic("testTopic2", 4)) cluster.add_topic(Topic("testTopic3", 4)) partition = cluster.topics['testTopic1'].partitions[0] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) partition = cluster.topics['testTopic1'].partitions[1] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[3], 1) partition = cluster.topics['testTopic1'].partitions[2] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[3], 1) partition = cluster.topics['testTopic1'].partitions[3] partition.add_replica(cluster.brokers[4], 0) partition.add_replica(cluster.brokers[1], 1) partition = cluster.topics['testTopic2'].partitions[0] partition.add_replica(cluster.brokers[4], 0) partition.add_replica(cluster.brokers[3], 1) partition = cluster.topics['testTopic2'].partitions[1] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[4], 1) partition = cluster.topics['testTopic2'].partitions[2] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[1], 1) partition = cluster.topics['testTopic2'].partitions[3] partition.add_replica(cluster.brokers[3], 0) partition.add_replica(cluster.brokers[1], 1) partition = cluster.topics['testTopic3'].partitions[0] partition.add_replica(cluster.brokers[3], 0) partition.add_replica(cluster.brokers[2], 1) partition = cluster.topics['testTopic3'].partitions[1] partition.add_replica(cluster.brokers[4], 0) partition.add_replica(cluster.brokers[2], 1) partition = cluster.topics['testTopic3'].partitions[2] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) partition = cluster.topics['testTopic3'].partitions[3] partition.add_replica(cluster.brokers[3], 0) partition.add_replica(cluster.brokers[4], 1) return cluster
def test_try_pick_new_broker(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(3, 'brokerhost3.example.com') self.cluster.add_broker(b3) b3.rack = "b" action = ActionBalanceRackAware(self.args, self.cluster) # Firmly order the deque action._random_brokers = deque([b3, b1, b2]) newbroker = action._try_pick_new_broker( self.cluster.topics['testTopic1'].partitions[0], 1) assert newbroker == b3 assert action._random_brokers == deque([b1, b2, b3])
def setUp(self): self.topic = Topic('testTopic', 10) self.broker = Broker(1, 'brokerhost1.example.com') for i in range(10): self.topic.partitions[i].replicas = [self.broker] self.reassignment = Reassignment(self.topic.partitions, pause_time=0) self.null_plugin = PluginModule()
def test_process_cluster_empty_broker(self): self.cluster.add_broker(Broker(3, 'brokerhost3.example.com')) b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] self.cluster.add_topic(Topic("testTopic3", 2)) partition = self.cluster.topics['testTopic3'].partitions[0] partition.size = 1000 partition.add_replica(b1, 0) partition.add_replica(b2, 1) partition = self.cluster.topics['testTopic3'].partitions[1] partition.add_replica(b2, 0) partition.add_replica(b1, 1) partition.size = 2000 action = ActionBalanceSize(self.args, self.cluster) action.process_cluster() assert sum([p.size for p in self.cluster.brokers[1].partitions[0]], 0) == 3000 assert sum([p.size for p in self.cluster.brokers[1].partitions[1]], 0) == 3000 assert sum([p.size for p in self.cluster.brokers[2].partitions[0]], 0) == 3000 assert sum([p.size for p in self.cluster.brokers[2].partitions[1]], 0) == 3000 assert sum([p.size for p in self.cluster.brokers[3].partitions[0]], 0) == 3000 assert sum([p.size for p in self.cluster.brokers[3].partitions[1]], 0) == 3000
def add_brokers_from_zk(cluster, zk): for b in zk.get_children("/brokers/ids"): broker_data, bstat = zk.get("/brokers/ids/{0}".format(b)) cluster.add_broker(Broker.create_from_json(int(b), broker_data)) if cluster.num_brokers() == 0: raise ZookeeperException( "The cluster specified does not have any brokers")
def test_partition_copy_with_replicas(self): broker = Broker(1, 'testhost1') self.topic.partitions[0].replicas = [broker] partition2 = self.topic.partitions[0].copy() assert self.topic.partitions[0] == partition2 assert partition2.replicas == [] assert self.topic.partitions[0] is not partition2
def test_process_cluster_empty_one(self): self.cluster.add_broker(Broker(3, 'brokerhost3.example.com')) b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = self.cluster.brokers[3] self.cluster.add_topic(Topic("testTopic3", 2)) partition = self.cluster.topics['testTopic3'].partitions[0] partition.add_replica(b3, 0) partition.add_replica(b2, 1) partition = self.cluster.topics['testTopic3'].partitions[1] partition.add_replica(b2, 0) partition.add_replica(b3, 1) self.cluster.topics['testTopic1'].partitions[0].swap_replicas(b1, b3) self.cluster.topics['testTopic1'].partitions[1].swap_replicas(b1, b3) self.cluster.topics['testTopic2'].partitions[0].swap_replicas(b1, b3) self.cluster.topics['testTopic2'].partitions[1].swap_replicas(b1, b3) action = ActionBalanceCount(self.args, self.cluster) action.process_cluster() assert len(b1.partitions[0]) == 2 assert len(b1.partitions[1]) == 2 assert len(b2.partitions[0]) == 2 assert len(b2.partitions[1]) == 2 assert len(b3.partitions[0]) == 2 assert len(b3.partitions[1]) == 2
def test_process_partitions_at_pos_swap_partition(self, mock_pick): action = ActionBalanceRackAware(self.args, self.cluster) b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(3, 'brokerhost3.example.com') b4 = Broker(4, 'brokerhost4.example.com') self.cluster.add_broker(b3) self.cluster.add_broker(b4) b3.rack = "a" b4.rack = "c" self.cluster.topics['testTopic2'].partitions[0].swap_replicas(b2, b3) self.cluster.topics['testTopic1'].partitions[1].swap_replicas(b1, b4) mock_pick.return_value = self.cluster.topics['testTopic1'].partitions[1] action._process_partitions_at_pos(0) assert self.cluster.topics['testTopic1'].partitions[1].replicas == [b3, b4] assert self.cluster.topics['testTopic2'].partitions[0].replicas == [b2, b1]
def test_broker_create_from_json_extended(self): jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}' broker2 = Broker.create_from_json(1, jsonstr) assert broker2.jmx_port == -1 assert broker2.timestamp == "1466985807242" assert broker2.endpoints == ["PLAINTEXT://10.0.0.10:9092"] assert broker2.version == 3 assert broker2.port == 9092
def add_topic_with_replicas(cluster, topic, topic_data): newtopic = Topic(topic, len(topic_data['partitions'])) for partition in topic_data['partitions']: for i, replica in enumerate(topic_data['partitions'][partition]): if replica not in cluster.brokers: # Hit a replica that's not in the ID list (which means it's dead) # We'll add it, but trying to get sizes will fail as we don't have a hostname cluster.add_broker(Broker(replica, None)) newtopic.partitions[int(partition)].add_replica( cluster.brokers[replica], i) cluster.add_topic(newtopic)
def set_up_cluster(): cluster = Cluster() cluster.retention = 100000 cluster.add_broker(Broker(1, "brokerhost1.example.com")) cluster.add_broker(Broker(2, "brokerhost2.example.com")) cluster.brokers[1].rack = "a" cluster.brokers[2].rack = "b" cluster.add_topic(Topic("testTopic1", 2)) cluster.add_topic(Topic("testTopic2", 2)) partition = cluster.topics['testTopic1'].partitions[0] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) partition = cluster.topics['testTopic1'].partitions[1] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[1], 1) partition = cluster.topics['testTopic2'].partitions[0] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[1], 1) partition = cluster.topics['testTopic2'].partitions[1] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) return cluster
def test_process_cluster_no_change(self): self.cluster.add_broker(Broker(3, "brokerhost3.example.com")) self.args.brokers = [3] self.args.to_broker = 1 action = ActionClone(self.args, self.cluster) action.process_cluster() b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] assert self.cluster.topics['testTopic1'].partitions[0].replicas == [b1, b2] assert self.cluster.topics['testTopic1'].partitions[1].replicas == [b2, b1] assert self.cluster.topics['testTopic2'].partitions[0].replicas == [b2, b1] assert self.cluster.topics['testTopic2'].partitions[1].replicas == [b1, b2]
def create_cluster_onehost(self): cluster = Cluster() cluster.add_broker(Broker(1, "brokerhost1.example.com")) cluster.add_topic(Topic("testTopic1", 2)) cluster.add_topic(Topic("testTopic2", 2)) partition = cluster.topics['testTopic1'].partitions[0] partition.add_replica(cluster.brokers[1], 0) partition = cluster.topics['testTopic1'].partitions[1] partition.add_replica(cluster.brokers[1], 0) partition = cluster.topics['testTopic2'].partitions[0] partition.add_replica(cluster.brokers[1], 0) partition = cluster.topics['testTopic2'].partitions[1] partition.add_replica(cluster.brokers[1], 0) return cluster
def set_up_cluster(): cluster = Cluster() cluster.add_broker(Broker(1, "brokerhost1.example.com")) cluster.add_broker(Broker(2, "brokerhost2.example.com")) cluster.add_topic(Topic("testTopic1", 2)) cluster.add_topic(Topic("testTopic2", 2)) partition = cluster.topics['testTopic1'].partitions[0] partition.add_replica(cluster.brokers[2], 0) partition.add_replica(cluster.brokers[1], 1) partition.size = 1001 partition = cluster.topics['testTopic1'].partitions[1] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) partition.size = 1002 partition = cluster.topics['testTopic2'].partitions[0] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) partition.size = 2001 partition = cluster.topics['testTopic2'].partitions[1] partition.add_replica(cluster.brokers[1], 0) partition.add_replica(cluster.brokers[2], 1) partition.size = 2002 return cluster
def test_process_cluster(self): self.cluster.add_broker(Broker(3, "brokerhost3.example.com")) self.args.brokers = [1] self.args.to_brokers = [] action = ActionRemove(self.args, self.cluster) action.process_cluster() b2 = self.cluster.brokers[2] b3 = self.cluster.brokers[3] assert self.cluster.topics['testTopic1'].partitions[0].replicas == [ b3, b2 ] assert self.cluster.topics['testTopic1'].partitions[1].replicas == [ b2, b3 ] assert self.cluster.topics['testTopic2'].partitions[0].replicas == [ b2, b3 ] assert self.cluster.topics['testTopic2'].partitions[1].replicas == [ b3, b2 ]
def test_process_cluster_increase(self): self.cluster.add_broker(Broker(3, "brokerhost3.example.com")) self.args.topics = ['testTopic1'] self.args.replication_factor = 3 action = ActionSetRF(self.args, self.cluster) action.process_cluster() b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = self.cluster.brokers[3] assert self.cluster.topics['testTopic1'].partitions[0].replicas == [ b1, b2, b3 ] assert self.cluster.topics['testTopic1'].partitions[1].replicas == [ b2, b1, b3 ] assert self.cluster.topics['testTopic2'].partitions[0].replicas == [ b2, b1 ] assert self.cluster.topics['testTopic2'].partitions[1].replicas == [ b1, b2 ]
def create_from_zookeeper(cls, zkconnect): log.info("Connecting to zookeeper {0}".format(zkconnect)) try: zk = KazooClient(zkconnect) zk.start() except Exception as e: raise ZookeeperException("Cannot connect to Zookeeper: {0}".format(e)) # Get broker list cluster = cls() for b in zk.get_children("/brokers/ids"): broker_data, bstat = zk.get("/brokers/ids/{0}".format(b)) cluster.add_broker(Broker.create_from_json(int(b), broker_data)) if cluster.num_brokers() == 0: raise ZookeeperException("The cluster specified does not have any brokers") # Get current partition state log.info("Getting partition list from Zookeeper") for topic in zk.get_children("/brokers/topics"): zdata, zstat = zk.get("/brokers/topics/{0}".format(topic)) zj = json.loads(zdata) newtopic = Topic(topic, len(zj['partitions'])) for partition in zj['partitions']: for i, replica in enumerate(zj['partitions'][partition]): if replica not in cluster.brokers: # Hit a replica that's not in the ID list (which means it's dead) # We'll add it, but trying to get sizes will fail as we don't have a hostname cluster.add_broker(Broker(replica, None)) newtopic.partitions[int(partition)].add_replica(cluster.brokers[replica], i) cluster.add_topic(newtopic) if cluster.num_topics() == 0: raise ZookeeperException("The cluster specified does not have any topics") log.info("Closing connection to zookeeper") zk.stop() zk.close() return cluster
def test_process_partitions_at_pos_swap_partition(self, mock_pick): action = ActionBalanceRackAware(self.args, self.cluster) b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker(3, 'brokerhost3.example.com') b4 = Broker(4, 'brokerhost4.example.com') self.cluster.add_broker(b3) self.cluster.add_broker(b4) b3.rack = "a" b4.rack = "c" self.cluster.topics['testTopic2'].partitions[0].swap_replicas(b2, b3) self.cluster.topics['testTopic1'].partitions[1].swap_replicas(b1, b4) mock_pick.return_value = self.cluster.topics['testTopic1'].partitions[ 1] action._process_partitions_at_pos(0) assert self.cluster.topics['testTopic1'].partitions[1].replicas == [ b3, b4 ] assert self.cluster.topics['testTopic2'].partitions[0].replicas == [ b2, b1 ]
def setUp(self): self.topic = Topic('testTopic', 10) self.broker = Broker(1, 'brokerhost1.example.com') for i in range(10): self.topic.partitions[i].replicas = [self.broker] self.replica_election = ReplicaElection(self.topic.partitions, pause_time=0)
def setUp(self): self.topic = Topic('testTopic', 10) self.broker = Broker(1, 'brokerhost1.example.com') for i in range(10): self.topic.partitions[i].replicas = [self.broker]
def test_broker_create_from_json_bad_jmx_port(self): jsonstr = '{"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}' broker2 = Broker.create_from_json(1, jsonstr) assert broker2.hostname == '10.0.0.10'
class BrokerTests(unittest.TestCase): def setUp(self): self.broker = Broker(1, 'brokerhost1.example.com') def add_partitions(self, pos, num): topic = Topic('testTopic', num) self.broker.partitions[pos] = [] for i in range(num): self.broker.partitions[pos].append(topic.partitions[i]) def test_broker_create(self): assert self.broker.id == 1 assert self.broker.hostname == 'brokerhost1.example.com' assert self.broker.partitions == {} def test_broker_create_from_json_basic(self): jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}' broker2 = Broker.create_from_json(1, jsonstr) assert broker2.id == 1 assert broker2.hostname == '10.0.0.10' assert broker2.partitions == {} def test_broker_create_from_json_extended(self): jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}' broker2 = Broker.create_from_json(1, jsonstr) assert broker2.jmx_port == -1 assert broker2.timestamp == "1466985807242" assert broker2.endpoints == ["PLAINTEXT://10.0.0.10:9092"] assert broker2.version == 3 assert broker2.port == 9092 def test_broker_create_from_json_bad_host(self): jsonstr = '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"hostname":"10.0.0.10","version":3,"port":9092}' self.assertRaises(ConfigurationException, Broker.create_from_json, 1, jsonstr) def test_broker_create_from_json_bad_jmx_port(self): jsonstr = '{"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092"],"host":"10.0.0.10","version":3,"port":9092}' broker2 = Broker.create_from_json(1, jsonstr) assert broker2.hostname == '10.0.0.10' def test_broker_copy_without_partitions(self): broker2 = self.broker.copy() assert broker2.id == 1 assert broker2.hostname == 'brokerhost1.example.com' assert self.broker is not broker2 def test_broker_copy_with_partitions(self): self.add_partitions(0, 1) broker2 = self.broker.copy() assert broker2.id == 1 assert broker2.hostname == 'brokerhost1.example.com' assert broker2.partitions == {} assert self.broker is not broker2 def test_broker_num_leaders(self): self.add_partitions(0, 2) assert self.broker.num_leaders() == 2 def test_broker_num_partitions_at_position_zero(self): self.add_partitions(1, 2) assert self.broker.num_partitions_at_position(0) == 0 def test_broker_num_partitions_single_position(self): self.add_partitions(0, 2) assert self.broker.num_partitions() == 2 def test_broker_num_partitions_two_position(self): self.add_partitions(0, 2) self.add_partitions(1, 1) assert self.broker.num_partitions() == 3 def test_broker_percent_leaders(self): self.add_partitions(0, 2) self.add_partitions(1, 2) assert self.broker.percent_leaders() == 50.0 def test_broker_percent_leaders_zero(self): self.add_partitions(1, 1) assert self.broker.percent_leaders() == 0.0 def test_broker_percent_leaders_no_partitions(self): assert self.broker.percent_leaders() == 0.0 def test_broker_equality(self): broker2 = Broker(1, 'brokerhost1.example.com') assert self.broker == broker2 def test_broker_inequality_hostname(self): broker2 = Broker(1, 'brokerhost2.example.com') assert self.broker != broker2 def test_broker_inequality_id(self): broker2 = Broker(2, 'brokerhost1.example.com') assert self.broker != broker2 def test_broker_equality_typeerror(self): self.assertRaises(TypeError, self.broker.__eq__, None)
def setUp(self): self.broker = Broker(1, 'brokerhost1.example.com')
def add_brokers_from_zk(cluster, zk): for b in zk.get_children("/brokers/ids"): broker_data, bstat = zk.get("/brokers/ids/{0}".format(b)) cluster.add_broker(Broker.create_from_json(int(b), broker_data)) if cluster.num_brokers() == 0: raise ZookeeperException("The cluster specified does not have any brokers")
def test_broker_inequality_id(self): broker2 = Broker(2, 'brokerhost1.example.com') assert self.broker != broker2
def test_broker_inequality_hostname(self): broker2 = Broker(1, 'brokerhost2.example.com') assert self.broker != broker2
def test_broker_equality(self): broker2 = Broker(1, 'brokerhost1.example.com') assert self.broker == broker2
def add_brokers(self, num): for i in range(1, num + 1): broker = Broker(i, "brokerhost{0}.example.com".format(i)) self.cluster.add_broker(broker)
def test_partition_dict_for_reassignment_with_replicas(self): broker = Broker(1, 'testhost1') self.topic.partitions[0].replicas = [broker] expected = {"topic": 'testTopic', "partition": 0, "replicas": [1]} assert self.topic.partitions[0].dict_for_reassignment() == expected
def test_partition_equality_with_different_replicas(self): partition2 = Partition(self.topic, 0) broker = Broker(1, 'testhost1') partition2.replicas = [broker] assert self.topic.partitions[0] == partition2