diff --git a/swh/scheduler/tests/common.py b/swh/scheduler/tests/common.py --- a/swh/scheduler/tests/common.py +++ b/swh/scheduler/tests/common.py @@ -5,7 +5,7 @@ import copy import datetime -from typing import Dict, List +from typing import Dict, List, Optional TEMPLATES = { "git": { @@ -51,7 +51,11 @@ def _task_from_template( - template: Dict, next_run: datetime.datetime, priority: str, *args, **kwargs + template: Dict, + next_run: datetime.datetime, + priority: Optional[str], + *args, + **kwargs, ) -> Dict: ret = copy.deepcopy(template) ret["next_run"] = next_run @@ -67,40 +71,30 @@ def tasks_from_template( template: Dict, max_timestamp: datetime.datetime, - num: int, - num_priority: int = 0, - priorities: Dict = {}, + num: Optional[int] = None, + priority: Optional[str] = None, + num_priorities: Dict[Optional[str], int] = {}, ) -> List[Dict]: - """Build tasks from template + """Build ``num`` tasks from template """ - - def _pop_priority(priorities): - if not priorities: - return None - for priority, remains in priorities.items(): - if remains > 0: - priorities[priority] = remains - 1 - return priority - return None - - if num_priority and priorities: - priorities = { - priority: ratio * num_priority for priority, ratio in priorities.items() - } - - tasks = [] - for i in range(num + num_priority): - priority = _pop_priority(priorities) - tasks.append( - _task_from_template( - template, - max_timestamp - datetime.timedelta(microseconds=i), - priority, - "argument-%03d" % i, - **{"kwarg%03d" % i: "bogus-kwarg"}, + assert bool(num) != bool(num_priorities), "mutually exclusive" + if not num_priorities: + assert num is not None # to please mypy + num_priorities = {None: num} + tasks: List[Dict] = [] + for (priority, num) in num_priorities.items(): + for _ in range(num): + i = len(tasks) + tasks.append( + _task_from_template( + template, + max_timestamp - datetime.timedelta(microseconds=i), + priority, + "argument-%03d" % i, + **{"kwarg%03d" % i: "bogus-kwarg"}, + ) ) - ) return tasks diff --git a/swh/scheduler/tests/test_common.py b/swh/scheduler/tests/test_common.py --- a/swh/scheduler/tests/test_common.py +++ b/swh/scheduler/tests/test_common.py @@ -27,27 +27,20 @@ def test_tasks_from_template_priority(): - nb_tasks_no_priority = 3 - nb_tasks_priority = 10 template = TEMPLATES["hg"] - priorities = { - "high": 0.5, - "normal": 0.3, - "low": 0.2, + num_priorities = { + None: 3, + "high": 5, + "normal": 3, + "low": 2, } next_run = datetime.datetime.utcnow() - tasks = tasks_from_template( - template, - next_run, - nb_tasks_no_priority, - num_priority=nb_tasks_priority, - priorities=priorities, - ) + tasks = tasks_from_template(template, next_run, num_priorities=num_priorities,) - assert len(tasks) == nb_tasks_no_priority + nb_tasks_priority + assert len(tasks) == sum(num_priorities.values()) - repartition_priority = {k: 0 for k in priorities.keys()} + repartition_priority = {k: 0 for k in num_priorities} for i, t in enumerate(tasks): assert t["type"] == template["type"] assert t["arguments"] is not None @@ -56,10 +49,7 @@ assert len(t["arguments"]["kwargs"].keys()) == 1 assert t["next_run"] == next_run - datetime.timedelta(microseconds=i) priority = t.get("priority") - if priority: - assert priority in priorities - repartition_priority[priority] += 1 + assert priority in num_priorities + repartition_priority[priority] += 1 - assert repartition_priority == { - k: v * nb_tasks_priority for k, v in priorities.items() - } + assert repartition_priority == num_priorities diff --git a/swh/scheduler/tests/test_scheduler.py b/swh/scheduler/tests/test_scheduler.py --- a/swh/scheduler/tests/test_scheduler.py +++ b/swh/scheduler/tests/test_scheduler.py @@ -31,8 +31,7 @@ ONEDAY = datetime.timedelta(days=1) -# for compatibility purpose with existing test code -PRIORITY_RATIO = {"high": 0.6, "normal": 0.3, "low": 0.2} +NUM_PRIORITY_TASKS = {None: 100, "high": 60, "normal": 30, "low": 20} def subdict(d, keys=None, excl=()): @@ -105,19 +104,15 @@ def test_create_tasks(self, swh_scheduler): self._create_task_types(swh_scheduler) - num_tasks_priority = 100 - tasks_1 = tasks_from_template(TEMPLATES["git"], utcnow(), 100) + num_git = 100 + tasks_1 = tasks_from_template(TEMPLATES["git"], utcnow(), num_git) tasks_2 = tasks_from_template( - TEMPLATES["hg"], - utcnow(), - 100, - num_tasks_priority, - priorities=PRIORITY_RATIO, + TEMPLATES["hg"], utcnow(), num_priorities=NUM_PRIORITY_TASKS ) tasks = tasks_1 + tasks_2 # tasks are returned only once with their ids - ret1 = swh_scheduler.create_tasks(tasks + tasks_1 + tasks_2) + ret1 = swh_scheduler.create_tasks(tasks + tasks) set_ret1 = set([t["id"] for t in ret1]) # creating the same set result in the same ids @@ -139,8 +134,7 @@ assert task["current_interval"] == task_type["default_interval"] assert task["policy"] == orig_task.get("policy", "recurring") priority = task.get("priority") - if priority: - actual_priorities[priority] += 1 + actual_priorities[priority] += 1 assert task["retries_left"] == (task_type["num_retries"] or 0) ids.add(task["id"]) @@ -154,10 +148,9 @@ del task["priority"] assert task == orig_task - assert dict(actual_priorities) == { - priority: int(ratio * num_tasks_priority) - for priority, ratio in PRIORITY_RATIO.items() - } + expected_priorities = NUM_PRIORITY_TASKS.copy() + expected_priorities[None] += num_git + assert dict(actual_priorities) == expected_priorities def test_peek_ready_tasks_no_priority(self, swh_scheduler): self._create_task_types(swh_scheduler) @@ -205,15 +198,9 @@ self._create_task_types(swh_scheduler) t = utcnow() task_type = TEMPLATES["git"]["type"] - num_tasks_priority = 100 - num_tasks_no_priority = 100 # Create tasks with and without priorities tasks = tasks_from_template( - TEMPLATES["git"], - t, - num=num_tasks_no_priority, - num_priority=num_tasks_priority, - priorities=PRIORITY_RATIO, + TEMPLATES["git"], t, num_priorities=NUM_PRIORITY_TASKS, ) count_priority = 0 @@ -238,15 +225,9 @@ self._create_task_types(swh_scheduler) t = utcnow() task_type = TEMPLATES["git"]["type"] - num_tasks_priority = 100 - num_tasks_no_priority = 100 # Create tasks with and without priorities tasks = tasks_from_template( - TEMPLATES["git"], - t, - num=num_tasks_no_priority, - num_priority=num_tasks_priority, - priorities=PRIORITY_RATIO, + TEMPLATES["git"], t, num_priorities=NUM_PRIORITY_TASKS ) random.shuffle(tasks) swh_scheduler.create_tasks(tasks)