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