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)