def test_get_count_running_tasks_on_slave():
with mock.patch(
'paasta_tools.mesos_tools.get_mesos_master',
autospec=True,
) as mock_get_master, mock.patch(
'paasta_tools.mesos_tools.get_mesos_task_count_by_slave',
autospec=True,
) as mock_get_mesos_task_count_by_slave:
mock_master = mock.Mock()
mock_mesos_state = mock.Mock()
mock_master.state_summary.return_value = mock_mesos_state
mock_get_master.return_value = mock_master
mock_slave_counts = [
{
'task_counts': mock.Mock(count=3, slave={'hostname': 'host1'})
},
{
'task_counts': mock.Mock(count=0, slave={'hostname': 'host2'})
},
]
mock_get_mesos_task_count_by_slave.return_value = mock_slave_counts
assert mesos_tools.get_count_running_tasks_on_slave('host1') == 3
assert mesos_tools.get_count_running_tasks_on_slave('host2') == 0
assert mesos_tools.get_count_running_tasks_on_slave('host3') == 0
assert mock_master.state_summary.called
mock_get_mesos_task_count_by_slave.assert_called_with(mock_mesos_state)
def test_get_count_running_tasks_on_slave():
with asynctest.patch(
"paasta_tools.mesos_tools.get_mesos_master",
autospec=True) as mock_get_master, asynctest.patch(
"paasta_tools.mesos_tools.get_mesos_task_count_by_slave",
autospec=True) as mock_get_mesos_task_count_by_slave:
mock_master = mock.Mock()
mock_mesos_state = mock.Mock()
mock_master.state_summary = asynctest.CoroutineMock(
func=asynctest.CoroutineMock(
), # https://github.com/notion/a_sync/pull/40
return_value=mock_mesos_state,
)
mock_get_master.return_value = mock_master
mock_slave_counts = [
{
"task_counts": mock.Mock(count=3, slave={"hostname": "host1"})
},
{
"task_counts": mock.Mock(count=0, slave={"hostname": "host2"})
},
]
mock_get_mesos_task_count_by_slave.return_value = mock_slave_counts
assert mesos_tools.get_count_running_tasks_on_slave("host1") == 3
assert mesos_tools.get_count_running_tasks_on_slave("host2") == 0
assert mesos_tools.get_count_running_tasks_on_slave("host3") == 0
assert mock_master.state_summary.called
mock_get_mesos_task_count_by_slave.assert_called_with(mock_mesos_state)
def is_host_drained(hostname):
"""Checks if a host has drained successfully by confirming it is
draining and currently running 0 tasks
:param hostname: hostname to check
:returns: True or False
"""
return is_host_draining(hostname=hostname) and get_count_running_tasks_on_slave(hostname) == 0
def test_get_count_running_tasks_on_slave():
with contextlib.nested(
mock.patch('paasta_tools.mesos_tools.get_mesos_master', autospec=True),
mock.patch('paasta_tools.mesos_tools.get_mesos_task_count_by_slave', autospec=True),
) as (
mock_get_master,
mock_get_mesos_task_count_by_slave
):
mock_master = mock.Mock()
mock_mesos_state = mock.Mock()
mock_master.state_summary.return_value = mock_mesos_state
mock_get_master.return_value = mock_master
mock_slave_counts = [{'task_counts': mock.Mock(count=3, slave={'hostname': 'host1'})},
{'task_counts': mock.Mock(count=0, slave={'hostname': 'host2'})}]
mock_get_mesos_task_count_by_slave.return_value = mock_slave_counts
assert mesos_tools.get_count_running_tasks_on_slave('host1') == 3
assert mesos_tools.get_count_running_tasks_on_slave('host2') == 0
assert mesos_tools.get_count_running_tasks_on_slave('host3') == 0
assert mock_master.state_summary.called
mock_get_mesos_task_count_by_slave.assert_called_with(mock_mesos_state)
