def test_get_host_to_terminate(mock_draining_client):
with mock.patch(
'clusterman.draining.queue.json', autospec=True,
) as mock_json:
mock_draining_client.client.receive_message.return_value = {'Messages': []}
assert mock_draining_client.get_host_to_terminate() is None
mock_draining_client.client.receive_message.return_value = {'Messages': [{
'MessageAttributes': {'Sender': {'StringValue': 'clusterman'}},
'ReceiptHandle': 'receipt_id',
'Body': 'Helloworld',
}]}
mock_json.loads.return_value = {
'instance_id': 'i123',
'ip': '10.1.1.1',
'hostname': 'host123',
'group_id': 'sfr123',
}
assert mock_draining_client.get_host_to_terminate() == Host(
sender='clusterman',
receipt_handle='receipt_id',
instance_id='i123',
ip='10.1.1.1',
hostname='host123',
group_id='sfr123',
)
mock_json.loads.assert_called_with('Helloworld')
mock_draining_client.client.receive_message.assert_called_with(
QueueUrl=mock_draining_client.termination_queue_url,
MessageAttributeNames=['Sender'],
MaxNumberOfMessages=1,
)
def test_host_from_instance_id():
with mock.patch(
'clusterman.draining.queue.ec2_describe_instances',
autospec=True,
) as mock_ec2_describe, mock.patch(
'socket.gethostbyaddr',
autospec=True,
) as mock_gethostbyaddr:
mock_ec2_describe.return_value = []
assert host_from_instance_id(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
) is None
mock_ec2_describe.return_value = [{
'Tags': [{
'Key': 'thing',
'Value': 'bar'
}]
}]
assert host_from_instance_id(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
) is None
mock_ec2_describe.return_value = [{
'Tags': [{
'Key': 'aws:ec2spot:fleet-request-id',
'Value': 'sfr-123'
}]
}]
assert host_from_instance_id(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
) is None
mock_ec2_describe.return_value = [{
'PrivateIpAddress':
'10.1.1.1',
'Tags': [{
'Key': 'aws:ec2spot:fleet-request-id',
'Value': 'sfr-123'
}]
}]
assert host_from_instance_id(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
) == Host(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
hostname=mock_gethostbyaddr.return_value[0],
group_id='sfr-123',
ip='10.1.1.1',
)
mock_gethostbyaddr.side_effect = socket.error
assert host_from_instance_id(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
) is None
# instance has no tags, probably because it is new and tags have not
# yet propagated
mock_ec2_describe.return_value = [{'InstanceId': 'i-123'}]
assert host_from_instance_id(
sender='aws',
receipt_handle='rcpt',
instance_id='i-123',
) is None
def test_process_drain_queue(mock_draining_client):
with mock.patch(
'clusterman.draining.queue.drain',
autospec=True,
) as mock_drain, mock.patch(
'clusterman.draining.queue.DrainingClient.get_host_to_drain',
autospec=True,
) as mock_get_host_to_drain, mock.patch(
'clusterman.draining.queue.DrainingClient.delete_drain_messages',
autospec=True,
) as mock_delete_drain_messages, mock.patch(
'clusterman.draining.queue.DrainingClient.submit_host_for_termination',
autospec=True,
) as mock_submit_host_for_termination, mock.patch(
'clusterman.draining.queue.arrow',
autospec=False,
) as mock_arrow:
mock_arrow.now = mock.Mock(return_value=mock.Mock(timestamp=1))
mock_mesos_client = mock.Mock()
mock_kubernetes_client = mock.Mock()
mock_get_host_to_drain.return_value = None
mock_draining_client.process_drain_queue(mock_mesos_client,
mock_kubernetes_client)
assert mock_draining_client.get_host_to_drain.called
assert not mock_drain.called
assert not mock_submit_host_for_termination.called
mock_host = mock.Mock(hostname='')
mock_get_host_to_drain.return_value = mock_host
mock_draining_client.process_drain_queue(mock_mesos_client,
mock_kubernetes_client)
mock_submit_host_for_termination.assert_called_with(
mock_draining_client, mock_host, delay=0)
mock_delete_drain_messages.assert_called_with(mock_draining_client,
[mock_host])
assert not mock_drain.called
mock_host = Host(
hostname='host1',
ip='10.1.1.1',
group_id='sfr1',
instance_id='i123',
sender='mmb',
receipt_handle='aaaaa',
)
mock_get_host_to_drain.return_value = mock_host
mock_draining_client.process_drain_queue(mock_mesos_client,
mock_kubernetes_client)
assert mock_draining_client.get_host_to_drain.called
mock_drain.assert_called_with(
mock_mesos_client,
['host1|10.1.1.1'],
1000000000,
1000000000,
)
mock_submit_host_for_termination.assert_called_with(
mock_draining_client, mock_host)
mock_delete_drain_messages.assert_called_with(mock_draining_client,
[mock_host])
# test we can't submit same host twice
mock_host = Host(
hostname='host1',
ip='10.1.1.1',
group_id='sfr1',
instance_id='i123',
sender='mmb',
receipt_handle='bbb',
)
mock_drain.reset_mock()
mock_submit_host_for_termination.reset_mock()
mock_get_host_to_drain.return_value = mock_host
mock_draining_client.process_drain_queue(mock_mesos_client,
mock_kubernetes_client)
assert mock_draining_client.get_host_to_drain.called
assert not mock_drain.called
assert not mock_submit_host_for_termination.called
mock_delete_drain_messages.assert_called_with(mock_draining_client,
[mock_host])