def test_handle_command_task(self):
with patch('framework.message.handlers.task.CommandTask') as mock:
mocked_transfer = MagicMock()
future = Future()
future.set_result(True)
instance = mock.return_value
mocked_transfer.send.return_value = future
test_new_task_message = Message.create({
'id': 'test_msg_id',
'type': MessageType.NEW_TASK,
'body': {
'task_id': 1,
'type': TaskType.COMMAND,
'shell_args': ['python', '-m', 'this'],
'shell': True
}
})
with mock_transfer(mocked_transfer):
handler = NewTaskHandler()
yield handler.handle(test_new_task_message)
mock.assert_called_once_with(task_id=1,
args=['python', '-m', 'this'],
shell=True)
self.assertEqual(mocked_transfer.send.call_count, 1)
test_new_task_message = mocked_transfer.send.call_args[0][0]
self.assertEqual(test_new_task_message.type,
MessageType.TASK_START)
self.assertEqual(test_new_task_message.body['task_id'], 1)
mocked_transfer.send.assert_called_once_with(test_new_task_message)
self.assertEqual(instance.start.call_count, 1)
self.assertEqual(len(running_tasks), 1)
def test_mock_transfer():
class MockedTransfer(object):
pass
mocked_transfer = MockedTransfer()
with mock_transfer(mocked_transfer):
transfer = get_current_transfer()
assert transfer is mocked_transfer
def test_heartbeat():
io_loop = IOLoop()
class MockedTransfer(object):
@gen.coroutine
def send(self, msg):
assert len(msg['id']) == 32
assert msg['type'] == MessageType.HEARTBEAT
assert 'id' in msg['body']
assert 'rss' in msg['body']
assert 'vms' in msg['body']
assert 'cpu' in msg['body']
io_loop.stop()
transfer = MockedTransfer()
scheduler = HeartbeatScheduler()
with config.mock({'heartbeat': 0.01}):
with mock_transfer(transfer):
io_loop.spawn_callback(scheduler.eventloop)
io_loop.start()
def test_handle_action_task(self):
with patch('framework.message.handlers.task.ActionTask') as mock:
mocked_transfer = MagicMock()
future = Future()
future.set_result(True)
instance = mock.return_value
mocked_transfer.send.return_value = future
test_new_task_message = Message.create({
'id': 'test_msg_id',
'type': MessageType.NEW_TASK,
'body': {
'task_id': 1,
'type': TaskType.ACTION,
'action': 'test_action',
'args': {
'a': 1,
'b': 2
}
}
})
with mock_transfer(mocked_transfer):
handler = NewTaskHandler()
yield handler.handle(test_new_task_message)
mock.assert_called_once_with(task_id=1,
action='test_action',
args={
'a': 1,
'b': 2
})
self.assertEqual(mocked_transfer.send.call_count, 1)
test_new_task_message = mocked_transfer.send.call_args[0][0]
self.assertEqual(test_new_task_message.type,
MessageType.TASK_START)
self.assertEqual(test_new_task_message.body['task_id'], 1)
mocked_transfer.send.assert_called_once_with(test_new_task_message)
self.assertEqual(instance.start.call_count, 1)
self.assertEqual(len(running_tasks), 1)