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 init_first_report():
report = get_first_report()
message = Message.create(MessageType.INFO, body=report)
transfer = get_current_transfer()
result = get_io_loop().run_sync(lambda: transfer.send(message))
if not result:
logger.error('Upload report failed')
sys.exit(-1)
def eventloop(self):
while True:
try:
params = {'cursor': self._cursor}
url = utils.update_query_params(self._base_url, params)
timeout = (config['message_fetch_timeout'] +
MessageFetchScheduler.TIMEOUT_DELAY)
response = yield self._http_client.fetch(
url,
method='GET',
headers={'Accept': 'application/json'},
connect_timeout=timeout,
request_timeout=timeout,
validate_cert=False,
raise_error=True,
)
self._retry_times = 0
json_data = escape.json_decode(response.body)
self._cursor = json_data['cursor']
self.set_cursor(self._cursor)
messages = json_data['messages']
logger.info('Received %s messages', len(messages))
for message in messages:
msg = Message.create(message)
reactor.feed(msg)
except Exception as exc:
logger.error('Error while fetching message: %s',
exc,
exc_info=True)
if hasattr(exc, 'response') and exc.response:
logger.error(exc.response.body)
self._retry_times += 1
logger.info('Current message cursor: %s', self._cursor)
if self._retry_times == 1:
wait_unit = 0
else:
wait_unit = (self._retry_times % 12 + 1)
wait_time = wait_unit * self._retry_interval
logger.info('Wait %s seconds to re-fetch', wait_time)
yield gen.sleep(wait_time)
def handle(self, msg):
body = msg.body
task_id = body.get('task_id')
type = body.get('type', TaskType.ACTION)
if type == TaskType.ACTION:
task = ActionTask(task_id=task_id,
action=body['action'],
args=body.get('args'))
elif type == TaskType.PYTHON:
task = PythonTask(task_id=task_id, stream=body['script'])
elif type == TaskType.COMMAND:
task = CommandTask(task_id=task_id,
args=body['shell_args'],
shell=body.get('shell', True))
else:
raise TypeError('Invalid task type: ' + str(type))
task.start()
running_tasks[task_id] = task
message = Message.create(MessageType.TASK_START, {'task_id': task_id})
result = yield self._transfer.send(message)
if not result:
yield cache.put_message(message)
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)
import os