diff --git a/swh/scheduler/celery_backend/recurrent_visits.py b/swh/scheduler/celery_backend/recurrent_visits.py --- a/swh/scheduler/celery_backend/recurrent_visits.py +++ b/swh/scheduler/celery_backend/recurrent_visits.py @@ -1,4 +1,4 @@ -# Copyright (C) 2021 The Software Heritage developers +# Copyright (C) 2021-2022 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 @@ -24,6 +24,7 @@ from kombu.utils.uuid import uuid from swh.scheduler.celery_backend.config import get_available_slots +from swh.scheduler.utils import create_origin_task_dict if TYPE_CHECKING: from ..interface import SchedulerInterface @@ -233,7 +234,7 @@ random.shuffle(origins) for origin in origins: - task_dict = origin.as_task_dict() + task_dict = create_origin_task_dict(origin) app.send_task( queue_name, task_id=uuid(), diff --git a/swh/scheduler/cli/origin.py b/swh/scheduler/cli/origin.py --- a/swh/scheduler/cli/origin.py +++ b/swh/scheduler/cli/origin.py @@ -10,6 +10,7 @@ import click from . import cli +from ..utils import create_origin_task_dict if TYPE_CHECKING: from uuid import UUID @@ -129,7 +130,7 @@ created = scheduler.create_tasks( [ { - **origin.as_task_dict(), + **create_origin_task_dict(origin), "policy": "oneshot", "next_run": utcnow(), "retries_left": 1, @@ -211,7 +212,7 @@ click.echo(f"{len(origins)} visits to send to celery") for origin in origins: - task_dict = origin.as_task_dict() + task_dict = create_origin_task_dict(origin) app.send_task( task_name, task_id=uuid(), diff --git a/swh/scheduler/model.py b/swh/scheduler/model.py --- a/swh/scheduler/model.py +++ b/swh/scheduler/model.py @@ -170,15 +170,6 @@ metadata={"auto_now": True}, ) - def as_task_dict(self): - return { - "type": f"load-{self.visit_type}", - "arguments": { - "args": [], - "kwargs": {"url": self.url, **self.extra_loader_arguments}, - }, - } - class LastVisitStatus(Enum): successful = "successful" diff --git a/swh/scheduler/simulator/__init__.py b/swh/scheduler/simulator/__init__.py --- a/swh/scheduler/simulator/__init__.py +++ b/swh/scheduler/simulator/__init__.py @@ -17,6 +17,7 @@ from simpy import Event from swh.scheduler.interface import SchedulerInterface +from swh.scheduler.utils import create_origin_task_dict from . import origin_scheduler, task_scheduler from .common import Environment, Queue, SimulationReport, Task @@ -121,7 +122,7 @@ scheduler.create_tasks( [ { - **origin.as_task_dict(), + **create_origin_task_dict(origin), "policy": "recurring", "next_run": origin.last_update, "interval": timedelta(days=64), diff --git a/swh/scheduler/tests/test_model.py b/swh/scheduler/tests/test_model.py --- a/swh/scheduler/tests/test_model.py +++ b/swh/scheduler/tests/test_model.py @@ -4,7 +4,6 @@ # See top-level LICENSE file for more information import datetime -import uuid import attr @@ -93,35 +92,3 @@ test1 = attr.ib(type=str) assert TestModel2.primary_key_columns() == ("col1", "col2") - - -def test_listed_origin_as_task_dict(): - origin = model.ListedOrigin( - lister_id=uuid.uuid4(), - url="http://example.com/", - visit_type="git", - ) - - task = origin.as_task_dict() - assert task == { - "type": "load-git", - "arguments": {"args": [], "kwargs": {"url": "http://example.com/"}}, - } - - loader_args = {"foo": "bar", "baz": {"foo": "bar"}} - - origin_w_args = model.ListedOrigin( - lister_id=uuid.uuid4(), - url="http://example.com/svn/", - visit_type="svn", - extra_loader_arguments=loader_args, - ) - - task_w_args = origin_w_args.as_task_dict() - assert task_w_args == { - "type": "load-svn", - "arguments": { - "args": [], - "kwargs": {"url": "http://example.com/svn/", **loader_args}, - }, - } diff --git a/swh/scheduler/tests/test_utils.py b/swh/scheduler/tests/test_utils.py --- a/swh/scheduler/tests/test_utils.py +++ b/swh/scheduler/tests/test_utils.py @@ -5,8 +5,9 @@ from datetime import timezone from unittest.mock import patch +import uuid -from swh.scheduler import utils +from swh.scheduler import model, utils @patch("swh.scheduler.utils.datetime") @@ -80,3 +81,35 @@ assert actual_task == expected_task mock_datetime.now.assert_called_once_with(tz=timezone.utc) + + +def test_create_origin_task_dict(): + origin = model.ListedOrigin( + lister_id=uuid.uuid4(), + url="http://example.com/", + visit_type="git", + ) + + task = utils.create_origin_task_dict(origin) + assert task == { + "type": "load-git", + "arguments": {"args": [], "kwargs": {"url": "http://example.com/"}}, + } + + loader_args = {"foo": "bar", "baz": {"foo": "bar"}} + + origin_w_args = model.ListedOrigin( + lister_id=uuid.uuid4(), + url="http://example.com/svn/", + visit_type="svn", + extra_loader_arguments=loader_args, + ) + + task_w_args = utils.create_origin_task_dict(origin_w_args) + assert task_w_args == { + "type": "load-svn", + "arguments": { + "args": [], + "kwargs": {"url": "http://example.com/svn/", **loader_args}, + }, + } diff --git a/swh/scheduler/utils.py b/swh/scheduler/utils.py --- a/swh/scheduler/utils.py +++ b/swh/scheduler/utils.py @@ -1,10 +1,13 @@ -# Copyright (C) 2017-2018 The Software Heritage developers +# Copyright (C) 2017-2022 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 +from typing import Any, Dict + +from .model import ListedOrigin def utcnow(): @@ -64,6 +67,16 @@ return task +def create_origin_task_dict(origin: ListedOrigin) -> Dict[str, Any]: + return { + "type": f"load-{origin.visit_type}", + "arguments": { + "args": [], + "kwargs": {"url": origin.url, **origin.extra_loader_arguments}, + }, + } + + def create_oneshot_task_dict(type, *args, **kwargs): """Create a oneshot task scheduled for as soon as possible.