コード例 #1
0
    def test_runs_task(self, mock_sleep):

        class RequestTasksMock(object):
            def __init__(self):
                self.called = False

            def __call__(self, *args, **kwargs):
                if self.called:
                    raise FatalError('Exit agent')
                self.called = True
                return {
                    'return_time': timeutil.format(timeutil.now()),
                    'tasks': [{
                        'task_type': 'test',
                        'task_version': 1,
                        'task_id': uuid.uuid4(),
                        'task_data': {}
                    }]
                }

        with patch.object(Communicator, 'request_tasks', side_effect=RequestTasksMock()):
            mock_handler = Mock(return_value={})
            TaskHandler.register(mock_handler, 'test', 1)

            with self.assertRaises(FatalError):
                Agent('https://localhost', default_wait=0).run()

        self.assertEqual(
            mock_handler.call_count, 1,
            'Agent should have executed task once'
        )
コード例 #2
0
 def test_dynamic_handler_registration(self):
     builtin_count = len(TaskHandler.list_all())
     self.assertGreaterEqual(
         builtin_count, 1,
         'There should be at least single built in tasks defined.'
     )
     register_1()
     register_2()
     self.assertEqual(
         len(TaskHandler.list_all()), builtin_count + 2,
         'There should be more than two new tasks defined.'
     )
コード例 #3
0
    def run(self):
        """The "main function" of the agent, looping the claim & execute tasks flow."""
        with Executor(self.max_tasks) as executor:
            while self._run:

                wait_time = self.default_wait
                min_wait_time = 1

                # request for tasks
                try:
                    task_response = self.afm.request_tasks(
                        agent_id=self.uuid,
                        agent_name=self.name,
                        agent_time=timeutil.format(timeutil.now()),
                        agent_capabilities=TaskHandler.list_all(),
                        max_tasks=executor.available_executors(),
                    )
                    if "tasks" in task_response:
                        for task_data in task_response["tasks"]:
                            executor.submit_task(
                                task_data, lambda *args, **kwargs: self.afm.post_result(*args, **kwargs)
                            )
                    if "return_time" in task_response:
                        return_time = timeutil.parse(task_response["return_time"])
                        wait_time = max(min_wait_time, (return_time - timeutil.now()).total_seconds())
                except TemporaryError as e:
                    logging.getLogger("Agent").error("An error occurred while claiming tasks: %s", e)

                time.sleep(wait_time)
コード例 #4
0
    def _run_task(self, task, callback):
        try:
            self._logger.debug('Creating task handler for task %s', task['task_id'])
            handler = TaskHandler.create(task['task_type'], task['task_version'])

            self._logger.debug('Executing task handler for task %s', task['task_id'])
            result = handler(task['task_data'])

            self._logger.info('Task executed successfully: %s', task['task_id'])
            callback(task['task_id'], task_data=result)

        except Exception as e:
            self._logger.warning('Execution of task %s raised an exception: %s', task['task_id'], str(e))
            callback(task['task_id'], task_error=str(e))

        finally:
            with self._lock:
                self._active_executors -= 1