def test_get_mesos_task_count_by_slave(): with contextlib.nested( mock.patch('paasta_tools.mesos_tools.get_running_tasks_from_active_frameworks', autospec=True), ) as ( mock_get_running_tasks_from_active_frameworks, ): mock_chronos = mock.Mock() mock_chronos.name = 'chronos' mock_marathon = mock.Mock() mock_marathon.name = 'marathon' mock_task1 = mock.Mock() mock_task1.slave = {'id': 'slave1'} mock_task1.framework = mock_chronos mock_task2 = mock.Mock() mock_task2.slave = {'id': 'slave1'} mock_task2.framework = mock_marathon mock_task3 = mock.Mock() mock_task3.slave = {'id': 'slave2'} mock_task3.framework = mock_marathon mock_task4 = mock.Mock() mock_task4.slave = {'id': 'slave2'} mock_task4.framework = mock_marathon mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_running_tasks_from_active_frameworks.return_value = mock_tasks mock_slave_1 = {'id': 'slave1', 'attributes': {'pool': 'default'}, 'hostname': 'host1'} mock_slave_2 = {'id': 'slave2', 'attributes': {'pool': 'default'}, 'hostname': 'host2'} mock_slave_3 = {'id': 'slave3', 'attributes': {'pool': 'another'}, 'hostname': 'host3'} mock_mesos_state = {'slaves': [mock_slave_1, mock_slave_2, mock_slave_3]} ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool='default') mock_get_running_tasks_from_active_frameworks.assert_called_with('') expected = [{'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=1, slave=mock_slave_1)}, {'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2)}] assert len(ret) == len(expected) and sorted(ret) == sorted(expected) ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool=None) mock_get_running_tasks_from_active_frameworks.assert_called_with('') expected = [{'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=1, slave=mock_slave_1)}, {'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2)}, {'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3)}] assert len(ret) == len(expected) and sorted(ret) == sorted(expected) # test slaves_list override mock_task2 = mock.Mock() mock_task2.slave = {'id': 'slave2'} mock_task2.framework = mock_marathon mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_running_tasks_from_active_frameworks.return_value = mock_tasks ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, slaves_list=[mock_slave_1, mock_slave_2, mock_slave_3]) expected = [{'id': 'slave1', 'attributes': {'pool': 'default'}, 'hostname': 'host1', 'task_counts': mesos_tools.SlaveTaskCount(count=1, chronos_count=1, slave=mock_slave_1)}, {'id': 'slave2', 'attributes': {'pool': 'default'}, 'hostname': 'host2', 'task_counts': mesos_tools.SlaveTaskCount(count=3, chronos_count=0, slave=mock_slave_2)}, {'id': 'slave3', 'attributes': {'pool': 'another'}, 'hostname': 'host3', 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3)}] assert len(ret) == len(expected) and sorted(ret) == sorted(expected)
def test_get_mesos_task_count_by_slave(): with contextlib.nested( mock.patch('paasta_tools.mesos_tools.get_running_tasks_from_active_frameworks'), ) as ( mock_get_running_tasks_from_active_frameworks, ): mock_chronos = mock.Mock() mock_chronos.name = 'chronos' mock_marathon = mock.Mock() mock_marathon.name = 'marathon' mock_task1 = mock.Mock() mock_task1.slave = {'id': 'slave1'} mock_task1.framework = mock_chronos mock_task2 = mock.Mock() mock_task2.slave = {'id': 'slave1'} mock_task2.framework = mock_marathon mock_task3 = mock.Mock() mock_task3.slave = {'id': 'slave2'} mock_task3.framework = mock_marathon mock_task4 = mock.Mock() mock_task4.slave = {'id': 'slave2'} mock_task4.framework = mock_marathon mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_running_tasks_from_active_frameworks.return_value = mock_tasks mock_slave_1 = {'id': 'slave1', 'attributes': {'pool': 'default'}, 'hostname': 'host1'} mock_slave_2 = {'id': 'slave2', 'attributes': {'pool': 'default'}, 'hostname': 'host2'} mock_slave_3 = {'id': 'slave3', 'attributes': {'pool': 'another'}, 'hostname': 'host3'} mock_mesos_state = {'slaves': [mock_slave_1, mock_slave_2, mock_slave_3]} ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool='default') mock_get_running_tasks_from_active_frameworks.assert_called_with('') assert ret == {'host1': mesos_tools.SlaveTaskCount(count=2, chronos_count=1, slave=mock_slave_1), 'host2': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2)} ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool=None) mock_get_running_tasks_from_active_frameworks.assert_called_with('') assert ret == {'host1': mesos_tools.SlaveTaskCount(count=2, chronos_count=1, slave=mock_slave_1), 'host2': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2), 'host3': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3)}
def test_get_mesos_task_count_by_slave(): with mock.patch('paasta_tools.mesos_tools.get_all_running_tasks', autospec=True) as mock_get_all_running_tasks: mock_chronos = mock.Mock() mock_chronos.name = 'chronos' mock_marathon = mock.Mock() mock_marathon.name = 'marathon' mock_task1 = mock.Mock() mock_task1.slave = {'id': 'slave1'} mock_task1.framework = mock_chronos mock_task2 = mock.Mock() mock_task2.slave = {'id': 'slave1'} mock_task2.framework = mock_marathon mock_task3 = mock.Mock() mock_task3.slave = {'id': 'slave2'} mock_task3.framework = mock_marathon mock_task4 = mock.Mock() mock_task4.slave = {'id': 'slave2'} mock_task4.framework = mock_marathon mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_all_running_tasks.return_value = mock_tasks mock_slave_1 = { 'id': 'slave1', 'attributes': { 'pool': 'default' }, 'hostname': 'host1' } mock_slave_2 = { 'id': 'slave2', 'attributes': { 'pool': 'default' }, 'hostname': 'host2' } mock_slave_3 = { 'id': 'slave3', 'attributes': { 'pool': 'another' }, 'hostname': 'host3' } mock_mesos_state = { 'slaves': [mock_slave_1, mock_slave_2, mock_slave_3] } ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool='default') assert mock_get_all_running_tasks.called expected = [ { 'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=1, slave=mock_slave_1) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected) ret = mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool=None) assert mock_get_all_running_tasks.called expected = [ { 'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=1, slave=mock_slave_1) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected) # test slaves_list override mock_task2 = mock.Mock() mock_task2.slave = {'id': 'slave2'} mock_task2.framework = mock_marathon mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_all_running_tasks.return_value = mock_tasks mock_slaves_list = [ { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_1) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_2) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3) }, ] ret = mesos_tools.get_mesos_task_count_by_slave( mock_mesos_state, slaves_list=mock_slaves_list, ) expected = [ { 'task_counts': mesos_tools.SlaveTaskCount(count=1, chronos_count=1, slave=mock_slave_1) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=3, chronos_count=0, slave=mock_slave_2) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected) # test SlaveDoesNotExist exception handling mock_task2.__getitem__ = mock.Mock(side_effect="fakeid") mock_task2.slave = mock.Mock() mock_task2.slave.__getitem__ = mock.Mock() mock_task2.slave.__getitem__.side_effect = mesos.exceptions.SlaveDoesNotExist # we expect to handle this SlaveDoesNotExist exception gracefully, and continue on to handle other tasks mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_all_running_tasks.return_value = mock_tasks mock_slaves_list = [ { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_1) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_2) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3) }, ] ret = mesos_tools.get_mesos_task_count_by_slave( mock_mesos_state, slaves_list=mock_slaves_list, ) # we expect mock_slave_2 to only count 2 tasks, as one of them returned a SlaveDoesNotExist exception expected = [ { 'task_counts': mesos_tools.SlaveTaskCount(count=1, chronos_count=1, slave=mock_slave_1) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=2, chronos_count=0, slave=mock_slave_2) }, { 'task_counts': mesos_tools.SlaveTaskCount(count=0, chronos_count=0, slave=mock_slave_3) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected)
async def test_get_mesos_task_count_by_slave(): with asynctest.patch("paasta_tools.mesos_tools.get_all_running_tasks", autospec=True) as mock_get_all_running_tasks: mock_tron = mock.Mock() mock_tron.name = "tron" mock_marathon = mock.Mock() mock_marathon.name = "marathon" mock_task1 = mock.Mock() mock_task1.slave = asynctest.CoroutineMock( return_value={"id": "slave1"}) mock_task1.framework = asynctest.CoroutineMock(return_value=mock_tron) mock_task2 = mock.Mock() mock_task2.slave = asynctest.CoroutineMock( return_value={"id": "slave1"}) mock_task2.framework = asynctest.CoroutineMock( return_value=mock_marathon) mock_task3 = mock.Mock() mock_task3.slave = asynctest.CoroutineMock( return_value={"id": "slave2"}) mock_task3.framework = asynctest.CoroutineMock( return_value=mock_marathon) mock_task4 = mock.Mock() mock_task4.slave = asynctest.CoroutineMock( return_value={"id": "slave2"}) mock_task4.framework = asynctest.CoroutineMock( return_value=mock_marathon) mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_all_running_tasks.return_value = mock_tasks mock_slave_1 = { "id": "slave1", "attributes": { "pool": "default" }, "hostname": "host1", } mock_slave_2 = { "id": "slave2", "attributes": { "pool": "default" }, "hostname": "host2", } mock_slave_3 = { "id": "slave3", "attributes": { "pool": "another" }, "hostname": "host3", } mock_mesos_state = { "slaves": [mock_slave_1, mock_slave_2, mock_slave_3] } ret = await mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool="default") assert mock_get_all_running_tasks.called expected = [ { "task_counts": mesos_tools.SlaveTaskCount(count=2, batch_count=1, slave=mock_slave_1) }, { "task_counts": mesos_tools.SlaveTaskCount(count=2, batch_count=0, slave=mock_slave_2) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected) ret = await mesos_tools.get_mesos_task_count_by_slave(mock_mesos_state, pool=None) assert mock_get_all_running_tasks.called expected = [ { "task_counts": mesos_tools.SlaveTaskCount(count=2, batch_count=1, slave=mock_slave_1) }, { "task_counts": mesos_tools.SlaveTaskCount(count=2, batch_count=0, slave=mock_slave_2) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_3) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected) # test slaves_list override mock_task2 = mock.Mock() mock_task2.slave = asynctest.CoroutineMock( return_value={"id": "slave2"}) mock_task2.framework = asynctest.CoroutineMock( return_value=mock_marathon) mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_all_running_tasks.return_value = mock_tasks mock_slaves_list = [ { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_1) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_2) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_3) }, ] ret = await mesos_tools.get_mesos_task_count_by_slave( mock_mesos_state, slaves_list=mock_slaves_list) expected = [ { "task_counts": mesos_tools.SlaveTaskCount(count=1, batch_count=1, slave=mock_slave_1) }, { "task_counts": mesos_tools.SlaveTaskCount(count=3, batch_count=0, slave=mock_slave_2) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_3) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected) # test SlaveDoesNotExist exception handling mock_task2.__getitem__ = mock.Mock(side_effect="fakeid") mock_task2.slave = asynctest.CoroutineMock(return_value=mock.Mock( __getitem__=mock.Mock( side_effect=mesos.exceptions.SlaveDoesNotExist))) # we expect to handle this SlaveDoesNotExist exception gracefully, and continue on to handle other tasks mock_tasks = [mock_task1, mock_task2, mock_task3, mock_task4] mock_get_all_running_tasks.return_value = mock_tasks mock_slaves_list = [ { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_1) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_2) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_3) }, ] ret = await mesos_tools.get_mesos_task_count_by_slave( mock_mesos_state, slaves_list=mock_slaves_list) # we expect mock_slave_2 to only count 2 tasks, as one of them returned a SlaveDoesNotExist exception expected = [ { "task_counts": mesos_tools.SlaveTaskCount(count=1, batch_count=1, slave=mock_slave_1) }, { "task_counts": mesos_tools.SlaveTaskCount(count=2, batch_count=0, slave=mock_slave_2) }, { "task_counts": mesos_tools.SlaveTaskCount(count=0, batch_count=0, slave=mock_slave_3) }, ] assert len(ret) == len(expected) and utils.sort_dicts( ret) == utils.sort_dicts(expected)