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