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 @@ -171,9 +171,14 @@ might want the disabled ones.""", ) @click.option( - "--lister-uuid", + "--lister-name", default=None, - help="Limit origins to those listed from such lister", + help="Limit origins to those listed from lister with provided name", +) +@click.option( + "--lister-instance-name", + default=None, + help="Limit origins to those listed from lister with instance name", ) @click.argument("type", type=str) @click.pass_context @@ -184,7 +189,8 @@ tablesample: Optional[float], type: str, enabled: bool, - lister_uuid: Optional[str] = None, + lister_name: Optional[str] = None, + lister_instance_name: Optional[str] = None, ): """Send the next origin visits of the TYPE loader to celery, filling the queue.""" from kombu.utils.uuid import uuid @@ -201,6 +207,13 @@ num_tasks = get_available_slots(app, queue_name, task_type["max_queue_length"]) click.echo(f"{num_tasks} slots available in celery queue") + + lister_uuid: Optional[str] = None + if lister_name and lister_instance_name: + lister = scheduler.get_lister(lister_name, lister_instance_name) + if lister: + lister_uuid = lister.id + origins = scheduler.grab_next_visits( type, num_tasks, 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 @@ -119,7 +119,7 @@ LISTERS = ( - {"name": "github"}, + {"name": "github", "instance_name": "github"}, {"name": "gitlab", "instance_name": "gitlab"}, {"name": "gitlab", "instance_name": "freedesktop"}, {"name": "npm"}, diff --git a/swh/scheduler/tests/test_cli_origin.py b/swh/scheduler/tests/test_cli_origin.py --- a/swh/scheduler/tests/test_cli_origin.py +++ b/swh/scheduler/tests/test_cli_origin.py @@ -114,11 +114,16 @@ assert scheduled_tasks <= all_possible_tasks +@pytest.mark.parametrize( + "extra_cmd_args", + [[], ["--lister-name", "github", "--lister-instance-name", "github"]], +) def test_send_to_celery( mocker, swh_scheduler, swh_scheduler_celery_app, listed_origins_by_type, + extra_cmd_args, ): for task_type in TASK_TYPES.values(): swh_scheduler.create_task_type(task_type) @@ -136,7 +141,9 @@ send_task = mocker.patch.object(swh_scheduler_celery_app, "send_task") send_task.return_value = None - result = invoke(swh_scheduler, args=("send-to-celery", visit_type)) + cmd_args = ["send-to-celery", visit_type] + extra_cmd_args + + result = invoke(swh_scheduler, args=tuple(cmd_args)) assert result.exit_code == 0 scheduled_tasks = { 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 @@ -162,7 +162,7 @@ "kwargs": { "url": "http://example.com/1", "lister_name": LISTERS[0]["name"], - "lister_instance_name": None, + "lister_instance_name": LISTERS[0]["instance_name"], }, }, }, @@ -173,7 +173,7 @@ "kwargs": { "url": "http://example.com/2", "lister_name": LISTERS[0]["name"], - "lister_instance_name": None, + "lister_instance_name": LISTERS[0]["instance_name"], }, }, },