diff --git a/swh/scheduler/tests/test_utils.py b/swh/scheduler/tests/test_utils.py index 192be61..2e34570 100644 --- a/swh/scheduler/tests/test_utils.py +++ b/swh/scheduler/tests/test_utils.py @@ -1,56 +1,79 @@ # Copyright (C) 2017-2018 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information import unittest from datetime import timezone from unittest.mock import patch from swh.scheduler import utils class UtilsTest(unittest.TestCase): @patch('swh.scheduler.utils.datetime') def test_create_oneshot_task_dict_simple(self, mock_datetime): mock_datetime.now.return_value = 'some-date' actual_task = utils.create_oneshot_task_dict('some-task-type') expected_task = { 'policy': 'oneshot', 'type': 'some-task-type', 'next_run': 'some-date', 'arguments': { 'args': [], 'kwargs': {}, }, - 'priority': None, } self.assertEqual(actual_task, expected_task) mock_datetime.now.assert_called_once_with(tz=timezone.utc) @patch('swh.scheduler.utils.datetime') def test_create_oneshot_task_dict_other_call(self, mock_datetime): mock_datetime.now.return_value = 'some-other-date' actual_task = utils.create_oneshot_task_dict( 'some-task-type', 'arg0', 'arg1', priority='high', other_stuff='normal' ) expected_task = { 'policy': 'oneshot', 'type': 'some-task-type', 'next_run': 'some-other-date', 'arguments': { 'args': ('arg0', 'arg1'), 'kwargs': {'other_stuff': 'normal'}, }, 'priority': 'high', } self.assertEqual(actual_task, expected_task) mock_datetime.now.assert_called_once_with(tz=timezone.utc) + + @patch('swh.scheduler.utils.datetime') + def test_create_task_dict(self, mock_datetime): + mock_datetime.now.return_value = 'date' + + actual_task = utils.create_task_dict( + 'task-type', 'recurring', 'arg0', 'arg1', + priority='low', other_stuff='normal', retries_left=3 + ) + + expected_task = { + 'policy': 'recurring', + 'type': 'task-type', + 'next_run': 'date', + 'arguments': { + 'args': ('arg0', 'arg1'), + 'kwargs': {'other_stuff': 'normal'}, + }, + 'priority': 'low', + 'retries_left': 3, + } + + self.assertEqual(actual_task, expected_task) + mock_datetime.now.assert_called_once_with(tz=timezone.utc) diff --git a/swh/scheduler/utils.py b/swh/scheduler/utils.py index bf848c2..cf9c502 100644 --- a/swh/scheduler/utils.py +++ b/swh/scheduler/utils.py @@ -1,71 +1,74 @@ # Copyright (C) 2017-2018 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information from datetime import datetime, timezone def get_task(task_name): """Retrieve task object in our application instance by its fully qualified python name. Args: task_name (str): task's name (e.g swh.loader.git.tasks.LoadDiskGitRepository) Returns: Instance of task """ from swh.scheduler.celery_backend.config import app for module in app.conf.CELERY_IMPORTS: __import__(module) return app.tasks[task_name] def create_task_dict(type, policy, *args, **kwargs): """Create a task with type and policy, scheduled for as soon as possible. Args: type (str): Type of oneshot task as per swh-scheduler's db table task_type's column (Ex: origin-update-git, swh-deposit-archive-checks) policy (str): oneshot or recurring policy Returns: Expected dictionary for the one-shot task scheduling api (swh.scheduler.backend.create_tasks) """ - priority = None - if 'priority' in kwargs: - priority = kwargs.pop('priority') - return { + task = { 'policy': policy, 'type': type, 'next_run': datetime.now(tz=timezone.utc), 'arguments': { 'args': args if args else [], 'kwargs': kwargs if kwargs else {}, }, - 'priority': priority, } + for extra_key in ['priority', 'retries_left']: + if extra_key in kwargs: + extra_val = kwargs.pop(extra_key) + task[extra_key] = extra_val + + return task + def create_oneshot_task_dict(type, *args, **kwargs): """Create a oneshot task scheduled for as soon as possible. Args: type (str): Type of oneshot task as per swh-scheduler's db table task_type's column (Ex: origin-update-git, swh-deposit-archive-checks) Returns: Expected dictionary for the one-shot task scheduling api (swh.scheduler.backend.create_tasks) """ return create_task_dict(type, 'oneshot', *args, **kwargs)