def test_update_brokers_from_metadata_update_rack(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'wrongrack' self.client.cluster.add_broker(broker1) self.client._update_brokers_from_metadata(self.metadata_response) assert_cluster_has_brokers(self.client.cluster, self.metadata_response)
def test_update_topics_from_metadata_create(self): # Don't want to test the broker update code here self.client.cluster.add_broker(Broker('host1.example.com', id=1, port=8031)) self.client.cluster.add_broker(Broker('host2.example.com', id=101, port=8032)) self.client._update_topics_from_metadata(self.metadata_response) assert_cluster_has_topics(self.client.cluster, self.metadata_response)
def setUp(self): self.cluster = Cluster() self.cluster.add_broker(Broker("brokerhost1.example.com", id=1)) self.cluster.add_broker(Broker("brokerhost2.example.com", id=2)) self.cluster.add_broker(Broker("brokerhost3.example.com", id=3)) self.cluster.add_topic(Topic("testTopic1", 2)) self.cluster.add_topic(Topic("testTopic2", 2))
def test_update_brokers_from_metadata_update_host(self): broker1 = Broker('wronghost.example.com', id=1, port=8031) self.client.cluster.add_broker(broker1) broker1.close = MagicMock() self.client._update_brokers_from_metadata(self.metadata_response) assert_cluster_has_brokers(self.client.cluster, self.metadata_response) broker1.close.assert_called_once()
def setUp(self): self.configuration = ClientConfiguration(num_retries=3, retry_backoff=0.5) self.mock_sock = MagicMock() self.broker = Broker('brokerhost1.example.com', id=1, sock=self.mock_sock, configuration=self.configuration)
def test_add_or_update_replica_swap(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker2 = Broker('host2.example.com', id=101, port=8032) partition = Partition('topic1', 0) partition.add_replica(broker1) partition.add_replica(broker2) partition.add_or_update_replica(1, broker1) assert partition.replicas[1] == broker1
def test_assure_topic_has_partitions_add_one(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker2 = Broker('host2.example.com', id=101, port=8032) topic = Topic('topic1', 1) topic.partitions[0].add_replica(broker1) topic.partitions[0].add_replica(broker2) assert len(topic.partitions) == 1 topic.assure_has_partitions(2) assert len(topic.partitions) == 2 assert topic.partitions[1].num == 1
def test_delete_replicas_from_partition_all(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker2 = Broker('host2.example.com', id=101, port=8032) topic = Topic('topic1', 1) topic.partitions[0].add_replica(broker1) topic.partitions[0].add_replica(broker2) partition = topic.partitions[0] assert len(partition.replicas) == 2 partition.delete_replicas(0) assert len(partition.replicas) == 0
def test_check_partition_swappable_racks_ok(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker('brokerhost3.example.com', id=3) b4 = Broker('brokerhost4.example.com', id=4) 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('brokerhost3.example.com', id=3) 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 test_send_group_aware_request(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') self.client.cluster.add_broker(broker1) self.client._send_any_broker = MagicMock() self.client._send_any_broker.return_value = self.group_coordinator self.client._send_group_aware_request('testgroup', 'fakerequest') assert 'testgroup' in self.client.cluster.groups assert self.client.cluster.groups['testgroup'].coordinator == broker1 broker1.send.assert_called_once_with('fakerequest')
def test_send_group_aware_request(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') self.client.cluster.add_broker(broker1) self.client._send_any_broker = MagicMock() self.client._send_any_broker.return_value = self.group_coordinator self.client._send_group_aware_request('testgroup', 'fakerequest') assert 'testgroup' in self.client.cluster.groups assert self.client.cluster.groups['testgroup'].coordinator == broker1 broker1.send.assert_called_once_with('fakerequest')
def test_try_pick_new_broker(self): b1 = self.cluster.brokers[1] b2 = self.cluster.brokers[2] b3 = Broker('brokerhost3.example.com', id=3) 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("brokerhost1.example.com", id=1)) cluster.add_broker(Broker("brokerhost2.example.com", id=2)) cluster.add_broker(Broker("brokerhost3.example.com", id=3)) cluster.add_broker(Broker("brokerhost4.example.com", id=4)) 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 setUp(self): # Dummy client for testing - we're not going to connect that bootstrap broker self.client = Client() self.client._connected = True # Two brokers for the client broker = Broker('host1.example.com', id=1, port=8031) broker.rack = 'rack1' self.client.cluster.add_broker(broker) broker = Broker('host2.example.com', id=101, port=8032) broker.rack = 'rack1' self.client.cluster.add_broker(broker) self.metadata = topic_metadata() self.metadata_error = topic_metadata_error()
def test_process_cluster_empty_one(self): self.cluster.add_broker(Broker('brokerhost3.example.com', id=3)) 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 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 setUp(self): self.broker = Broker('testhost', port=3945) self.request_classes = {'testklass': {1: TestKlass1}, 'testklassy': {1: TestKlassy1}} self.request_cmds = {'testklass': TestKlass1, 'testklassv1': TestKlass1, 'testklassy': TestKlassy1, 'testklassyv1': TestKlassy1}
def setUp(self): self.topic = Topic('testTopic', 10) self.broker = Broker('brokerhost1.example.com', id=1) 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 main(): # Set up and parse all CLI arguments args = set_up_arguments() # Connect to the specified Kafka broker broker = Broker(args.broker, port=args.port) broker.connect() # Loop on reading a command _cli_loop(broker) # Disconnect broker.close() # Nothing to see here, move along return 0
def test_process_cluster_empty_broker(self): self.cluster.add_broker(Broker('brokerhost3.example.com', id=3)) 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 test_send_group_aware_request_new_broker(self, mock_broker_class): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') mock_broker_class.return_value = broker1 self.client._send_any_broker = MagicMock() self.client._send_any_broker.return_value = self.group_coordinator self.client._send_group_aware_request('testgroup', 'fakerequest') mock_broker_class.assert_called_once_with('host1.example.com', id=1, port=8031, configuration=self.client.configuration) assert 1 in self.client.cluster.brokers assert self.client.cluster.brokers[1] == broker1 assert self.client.cluster.groups['testgroup'].coordinator == broker1 broker1.send.assert_called_once_with('fakerequest')
def test_process_cluster_zero_partitions(self): self.cluster.add_broker(Broker("brokerhost3.example.com", id=3)) self.cluster.brokers[1].rack = "c" self.cluster.add_topic(Topic("testTopic3", 1)) partition = self.cluster.topics['testTopic3'].partitions[0] partition.add_replica(self.cluster.brokers[1], 0) partition.add_replica(self.cluster.brokers[2], 1) partition.add_replica(self.cluster.brokers[3], 2) self.cluster.topics['testTopic3'].partitions[0].set_size(2000) self.args.exclude_topics = ['testTopic3'] action = ActionBalanceRate(self.args, self.cluster) action.process_cluster() assert sum( [p.scaled_size for p in self.cluster.brokers[1].partitions[0]], 0) == 3000 assert sum( [p.scaled_size for p in self.cluster.brokers[1].partitions[1]], 0) == 2000 assert sum( [p.scaled_size for p in self.cluster.brokers[2].partitions[0]], 0) == 3000 assert sum( [p.scaled_size for p in self.cluster.brokers[2].partitions[1]], 0) == 4000 assert sum( [p.scaled_size for p in self.cluster.brokers[3].partitions[0]], 0) == 2000 assert sum( [p.scaled_size for p in self.cluster.brokers[3].partitions[1]], 0) == 2000 assert sum( [p.scaled_size for p in self.cluster.brokers[3].partitions[2]], 0) == 2000
def test_partition_copy_with_replicas(self): broker = Broker('testhost1', id=1) 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_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('brokerhost3.example.com', id=3) b4 = Broker('brokerhost4.example.com', id=4) 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 test_update_topics_from_metadata_update_replicas(self): # Don't want to test the broker update code here broker1 = Broker('host1.example.com', id=1, port=8031) broker2 = Broker('host2.example.com', id=101, port=8032) broker3 = Broker('host3.example.com', id=304, port=8033) topic = Topic('topic1', 2) self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) self.client.cluster.add_broker(broker3) self.client.cluster.add_topic(topic) topic.partitions[0].add_replica(broker3) topic.partitions[0].add_replica(broker1) topic.partitions[1].add_replica(broker2) topic.partitions[1].add_replica(broker1) self.client._update_topics_from_metadata(self.metadata_response) assert_cluster_has_topics(self.client.cluster, self.metadata_response)
def test_update_topics_from_metadata_missing_broker(self): # Don't want to test the broker update code here self.client.cluster.add_broker(Broker('host1.example.com', id=1, port=8031)) self.client._update_topics_from_metadata(self.metadata_response) assert_cluster_has_topics(self.client.cluster, self.metadata_response) assert 101 in self.client.cluster.brokers assert self.client.cluster.brokers[101].endpoint.hostname is None
def test_broker_get_endpoint(self): jsonstr = ('{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092", "SSL://10.0.0.10:2834"],' '"host":"10.0.0.10","version":3,"port":9092}') broker = Broker.create_from_json(1, jsonstr) endpoint = broker.get_endpoint("SSL") assert endpoint.protocol == "SSL" assert endpoint.hostname == "10.0.0.10" assert endpoint.port == 2834
def test_add_or_update_group(self): broker = Broker('host1.example.com', id=1, port=8031) self.client.cluster.add_broker(broker) list_group = list_groups() self.client._add_or_update_group(list_group['groups'][0], 1) assert 'group1' in self.client.cluster.groups assert self.client.cluster.groups['group1'].coordinator == broker assert self.client.cluster.groups['group1'].protocol_type == 'protocol1'
def test_broker_get_endpoint(self): jsonstr = ( '{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092", "SSL://10.0.0.10:2834"],' '"host":"10.0.0.10","version":3,"port":9092}') broker = Broker.create_from_json(1, jsonstr) endpoint = broker.get_endpoint("SSL") assert endpoint.protocol == "SSL" assert endpoint.hostname == "10.0.0.10" assert endpoint.port == 2834
def _update_or_add_partition(self, partition_metadata, partition): partition.leader = self.cluster.brokers[partition_metadata['leader']] for i, replica in enumerate(partition_metadata['replicas']): if replica not in self.cluster.brokers: # We have a replica ID that is not a known broker. This can happen if a broker is offline, or # if the partition is otherwise assigned to a non-existent broker ID. In this case, we need to # create a broker object for this with no endpoint information as a placeholder. self.cluster.add_broker(Broker(hostname=None, id=replica)) partition.add_or_update_replica(i, self.cluster.brokers[replica]) partition.delete_replicas(len(partition_metadata['replicas']))
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(None, id=replica)) newtopic.partitions[int(partition)].add_replica(cluster.brokers[replica], i) cluster.add_topic(newtopic)
def test_maybe_delete_topics_not_in_metadata(self): # Don't want to test the broker update code here broker1 = Broker('host1.example.com', id=1, port=8031) broker2 = Broker('host2.example.com', id=101, port=8032) topic = Topic('topic1', 1) self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) self.client.cluster.add_topic(topic) topic.partitions[0].add_replica(broker2) topic.partitions[0].add_replica(broker1) topic = Topic('topic2', 1) self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) self.client.cluster.add_topic(topic) topic.partitions[0].add_replica(broker2) topic.partitions[0].add_replica(broker1) self.client._maybe_delete_topics_not_in_metadata(self.metadata_response, delete=True) assert 'topic2' not in self.client.cluster.topics assert 'topic1' in self.client.cluster.topics
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 len(broker2.endpoints) == 1 assert broker2.endpoints[0].protocol == 'PLAINTEXT' assert broker2.endpoints[0].hostname == '10.0.0.10' assert broker2.endpoints[0].port == 9092 assert broker2.version == 3 assert broker2.port == 9092
def test_send_group_aware_request_new_broker(self, mock_broker_class): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') mock_broker_class.return_value = broker1 self.client._send_any_broker = MagicMock() self.client._send_any_broker.return_value = self.group_coordinator self.client._send_group_aware_request('testgroup', 'fakerequest') mock_broker_class.assert_called_once_with( 'host1.example.com', id=1, port=8031, configuration=self.client.configuration) assert 1 in self.client.cluster.brokers assert self.client.cluster.brokers[1] == broker1 assert self.client.cluster.groups['testgroup'].coordinator == broker1 broker1.send.assert_called_once_with('fakerequest')
def set_up_cluster(): cluster = Cluster() cluster.retention = 100000 cluster.add_broker(Broker("brokerhost1.example.com", id=1)) cluster.add_broker(Broker("brokerhost2.example.com", id=2)) 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_send_any_broker(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') broker2 = Broker('host2.example.com', id=101, port=8032) broker2.rack = 'rack1' broker2.send = MagicMock() broker2.send.return_value = (2, 'otherresponse') self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) val = self.client._send_any_broker('fakerequest') broker2.send.assert_called_once_with('fakerequest') assert val == 'otherresponse'
def test_close(self): # Two brokers for the client broker = Broker('host1.example.com', id=1, port=8031) broker.rack = 'rack1' broker.close = MagicMock() self.client.cluster.add_broker(broker) broker = Broker('host2.example.com', id=101, port=8032) broker.rack = 'rack1' broker.close = MagicMock() self.client.cluster.add_broker(broker) self.client.close() for broker_id in self.client.cluster.brokers: self.client.cluster.brokers[broker_id].close.assert_called_once()
def test_send_any_broker(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') broker2 = Broker('host2.example.com', id=101, port=8032) broker2.rack = 'rack1' broker2.send = MagicMock() broker2.send.return_value = (2, 'otherresponse') self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) val = self.client._send_any_broker('fakerequest') broker2.send.assert_called_once_with('fakerequest') assert val == 'otherresponse'
def test_send_all_error(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.side_effect = ConnectionError broker2 = Broker('host2.example.com', id=101, port=8032) broker2.rack = 'rack1' broker2.send = MagicMock() broker2.send.side_effect = ConnectionError self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) self.assertRaises(ConnectionError, self.client._send_any_broker, 'fakerequest') broker1.send.assert_called_once_with('fakerequest') broker2.send.assert_called_once_with('fakerequest')
def test_close(self): # Two brokers for the client broker = Broker('host1.example.com', id=1, port=8031) broker.rack = 'rack1' broker.close = MagicMock() self.client.cluster.add_broker(broker) broker = Broker('host2.example.com', id=101, port=8032) broker.rack = 'rack1' broker.close = MagicMock() self.client.cluster.add_broker(broker) self.client.close() for broker_id in self.client.cluster.brokers: self.client.cluster.brokers[broker_id].close.assert_called_once()
def test_send_all_brokers_error(self): broker1 = Broker('host1.example.com', id=1, port=8031) broker1.rack = 'rack1' broker1.send = MagicMock() broker1.send.return_value = (1, 'fakeresponse') broker2 = Broker('host2.example.com', id=101, port=8032) broker2.rack = 'rack1' broker2.send = MagicMock() broker2.send.side_effect = ConnectionError self.client.cluster.add_broker(broker1) self.client.cluster.add_broker(broker2) val = self.client._send_all_brokers('fakerequest') broker1.send.assert_called_once_with('fakerequest') broker2.send.assert_called_once_with('fakerequest') assert val[1] == 'fakeresponse' assert val[101] is None
def _maybe_bootstrap_cluster(self, broker_port): """Attempt to bootstrap the cluster information using the given broker""" broker = Broker(broker_port[0], port=broker_port[1], configuration=self.configuration) try: broker.connect() except ConnectionError: # Just skip to the next bootstrap broker return False # Fetch topic metadata for all topics/brokers req = TopicMetadataV1Request({'topics': None}) correlation_id, metadata = broker.send(req) broker.close() # Add brokers and topics to cluster self._controller_id = metadata['controller_id'] self._update_from_metadata(metadata) self._last_full_metadata = time.time() return True
class BrokerTests(unittest.TestCase): def setUp(self): self.configuration = ClientConfiguration(num_retries=3, retry_backoff=0.5) self.mock_sock = MagicMock() self.broker = Broker('brokerhost1.example.com', id=1, sock=self.mock_sock, configuration=self.configuration) 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 == {} assert isinstance(self.broker._sock, MagicMock) 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_get_endpoint(self): jsonstr = ('{"jmx_port":-1,"timestamp":"1466985807242","endpoints":["PLAINTEXT://10.0.0.10:9092", "SSL://10.0.0.10:2834"],' '"host":"10.0.0.10","version":3,"port":9092}') broker = Broker.create_from_json(1, jsonstr) endpoint = broker.get_endpoint("SSL") assert endpoint.protocol == "SSL" assert endpoint.hostname == "10.0.0.10" assert endpoint.port == 2834 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 len(broker2.endpoints) == 1 assert broker2.endpoints[0].protocol == 'PLAINTEXT' assert broker2.endpoints[0].hostname == '10.0.0.10' assert broker2.endpoints[0].port == 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('brokerhost1.example.com', id=1) assert self.broker == broker2 def test_broker_inequality_hostname(self): broker2 = Broker('brokerhost2.example.com', id=1) assert self.broker != broker2 def test_broker_inequality_id(self): broker2 = Broker('brokerhost1.example.com', id=2) assert self.broker != broker2 def test_broker_equality_typeerror(self): self.assertRaises(TypeError, self.broker.__eq__, None) def test_broker_get_socket(self): sock = self.broker._get_socket(None) assert isinstance(sock, socket.socket) @patch('kafka.tools.models.broker.socket.socket') def test_broker_get_socket_wrapped(self, mock_socket): mock_context = MagicMock() mock_context.wrap_socket.return_value = 'fakewrappedsocket' mock_socket.return_value = 'fakesocket' sock = self.broker._get_socket(mock_context) mock_context.wrap_socket.assert_called_once_with('fakesocket', server_hostname='brokerhost1.example.com') print(sock) assert sock == 'fakewrappedsocket' def test_broker_connect(self): self.broker.connect() self.mock_sock.connect.assert_called_once_with(('brokerhost1.example.com', 9092)) def test_broker_connect_error(self): self.mock_sock.connect.side_effect = socket.error self.assertRaises(ConnectionError, self.broker.connect) def test_broker_close(self): self.broker.close() self.mock_sock.shutdown.assert_called_once() self.mock_sock.close.assert_called_once() def test_broker_close_not_open(self): self.mock_sock.shutdown.side_effect = OSError() self.broker.close() self.mock_sock.close.assert_called_once() @patch.object(Broker, '_read_bytes') def test_broker_single_send(self, mock_read_bytes): # recv is called to get the response size (first 4 bytes) self.mock_sock.recv.return_value = b'\x00\x00\x00\x10' # read_bytes returns the response payload, minus the response size # correlation_id (4), error (2), array of (api_key (2), min_version (2), max_version (2)) mock_read_bytes.return_value = b'\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x01\x01\x01\x02\x02' request = ApiVersionsV0Request({}) (correlation_id, response) = self.broker._single_send(request) # Check that the request was encoded properly self.mock_sock.sendall.assert_called_once_with(bytearray(b'\x00\x00\x00\x15\x00\x12\x00\x00\x00\x00\x00\x01\x00\x0bkafka-tools')) self.mock_sock.recv.assert_called_once_with(4) mock_read_bytes.assert_called_once_with(16) # Check that the response is what was expected to be decoded assert isinstance(response, ApiVersionsV0Response) assert response.correlation_id == 1 assert response['error'] == 0 assert isinstance(response['api_versions'], collections.Sequence) assert len(response['api_versions']) == 1 assert len(response['api_versions'][0]) == 3 assert response['api_versions'][0]['api_key'] == 1 assert response['api_versions'][0]['min_version'] == 257 assert response['api_versions'][0]['max_version'] == 514 # Correlation ID must be incremented after each request assert self.broker._correlation_id == 2 @patch.object(Broker, '_read_bytes') def test_broker_single_send_error(self, mock_read_bytes): self.mock_sock.recv.side_effect = socket.error mock_read_bytes.return_value = b'\x00\x00\x00\x01\x00\x00\x00\x00\x00\x01\x00\x01\x01\x01\x02\x02' request = ApiVersionsV0Request({}) self.assertRaises(ConnectionError, self.broker._single_send, request) @patch.object(ByteBuffer, 'getInt32') def test_broker_single_send_short(self, mock_getint): self.mock_sock.recv.return_value = b'\x00\x00\x00\x10' mock_getint.side_effect = EOFError request = ApiVersionsV0Request({}) self.assertRaises(ConnectionError, self.broker._single_send, request) @patch.object(Broker, '_single_send') @patch.object(Broker, 'connect') @patch.object(Broker, 'close') def test_broker_send(self, mock_close, mock_connect, mock_send): mock_send.return_value = 'fakeresponse' self.broker.send('fakerequest') mock_connect.assert_not_called() mock_close.assert_not_called() mock_send.assert_called_once_with('fakerequest') @patch.object(Broker, '_single_send') @patch.object(Broker, 'connect') @patch.object(Broker, 'close') def test_broker_send_retry(self, mock_close, mock_connect, mock_send): def close_broker(): self.broker._sock = None mock_send.side_effect = [ConnectionError, 'fakeresponse'] mock_close.side_effect = close_broker self.broker.send('fakerequest') mock_close.assert_called_once_with() mock_connect.assert_called_once_with() mock_send.assert_has_calls([call('fakerequest'), call('fakerequest')]) @patch.object(Broker, '_single_send') @patch.object(Broker, 'connect') @patch.object(Broker, 'close') def test_broker_send_exhaust_retry(self, mock_close, mock_connect, mock_send): def close_broker(): self.broker._sock = None mock_send.side_effect = ConnectionError mock_close.side_effect = close_broker self.assertRaises(ConnectionError, self.broker.send, 'fakerequest') mock_close.assert_has_calls([call(), call()]) mock_connect.assert_has_calls([call(), call()]) mock_send.assert_has_calls([call('fakerequest'), call('fakerequest'), call('fakerequest')]) def test_broker_read_bytes(self): self.mock_sock.recv.return_value = b'\x01\x02\x03\x04' response = self.broker._read_bytes(4) self.mock_sock.recv.assert_called_once_with(4) assert response == b'\x01\x02\x03\x04' def test_broker_read_bytes_concat(self): self.mock_sock.recv.side_effect = [b'\x01\x02', b'\x03\x04'] response = self.broker._read_bytes(4) self.mock_sock.recv.assert_has_calls([call(4), call(2)]) assert response == b'\x01\x02\x03\x04' def test_broker_read_bytes_nodata(self): self.mock_sock.recv.return_value = b'' self.assertRaises(socket.error, self.broker._read_bytes, 4) def test_broker_read_bytes_error(self): self.mock_sock.recv.side_effect = socket.error self.assertRaises(socket.error, self.broker._read_bytes, 4)
def setUp(self): self.configuration = ClientConfiguration(num_retries=3, retry_backoff=0.5) self.mock_sock = MagicMock() self.broker = Broker('brokerhost1.example.com', id=1, sock=self.mock_sock, configuration=self.configuration)
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'