Changeset View
Changeset View
Standalone View
Standalone View
swh/scheduler/tests/test_cli_task_type.py
- This file was added.
| # Copyright (C) 2019 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 pytest | |||||
| import traceback | |||||
| import yaml | |||||
| import pkg_resources | |||||
| from click.testing import CliRunner | |||||
| from swh.scheduler import get_scheduler | |||||
| from swh.scheduler.cli import cli | |||||
| FAKE_MODULE_ENTRY_POINTS = { | |||||
| 'lister.gnu=swh.lister.gnu:register', | |||||
| 'lister.pypi=swh.lister.pypi:register', | |||||
| } | |||||
| @pytest.fixture | |||||
| def mock_pkg_resources(monkeypatch): | |||||
| """Monkey patch swh.scheduler's mock_pkg_resources.iter_entry_point call | |||||
| """ | |||||
| def fake_iter_entry_points(*args, **kwargs): | |||||
| """Substitute fake function to return a fixed set of entrypoints | |||||
| """ | |||||
| from pkg_resources import EntryPoint, Distribution | |||||
| d = Distribution() | |||||
| return [EntryPoint.parse(entry, dist=d) | |||||
| for entry in FAKE_MODULE_ENTRY_POINTS] | |||||
| original_method = pkg_resources.iter_entry_points | |||||
| monkeypatch.setattr( | |||||
| pkg_resources, "iter_entry_points", fake_iter_entry_points) | |||||
| yield | |||||
| # reset monkeypatch: is that needed? | |||||
| monkeypatch.setattr( | |||||
| pkg_resources, "iter_entry_points", original_method) | |||||
| @pytest.fixture | |||||
| def local_sched_config(swh_scheduler_config): | |||||
| """Expose the local scheduler configuration | |||||
| """ | |||||
| return { | |||||
| 'scheduler': { | |||||
| 'cls': 'local', | |||||
| 'args': swh_scheduler_config | |||||
| } | |||||
| } | |||||
| @pytest.fixture | |||||
| def local_sched_configfile(local_sched_config, tmp_path): | |||||
| """Write in temporary location the local scheduler configuration | |||||
| """ | |||||
| configfile = tmp_path / 'config.yml' | |||||
| configfile.write_text(yaml.dump(local_sched_config)) | |||||
| return configfile.as_posix() | |||||
| def test_register_ttypes_all( | |||||
| mock_pkg_resources, local_sched_config, local_sched_configfile): | |||||
| """Registering all task types""" | |||||
| for command in [ | |||||
| ['--config-file', local_sched_configfile, 'task-type', 'register'], | |||||
| ['--config-file', local_sched_configfile, 'task-type', 'register', | |||||
| '-p', 'all'], | |||||
| ['--config-file', local_sched_configfile, 'task-type', 'register', | |||||
| '-p', 'lister.gnu', | |||||
| '-p', 'lister.pypi'], | |||||
| ]: | |||||
| result = CliRunner().invoke(cli, command) | |||||
| assert result.exit_code == 0, traceback.print_exception( | |||||
| *result.exc_info) | |||||
| scheduler = get_scheduler(**local_sched_config['scheduler']) | |||||
| all_tasks = [ | |||||
| 'list-gnu-full', | |||||
| 'list-pypi', | |||||
| ] | |||||
| for task in all_tasks: | |||||
| task_type_desc = scheduler.get_task_type(task) | |||||
| assert task_type_desc | |||||
| assert task_type_desc['type'] == task | |||||
| assert task_type_desc['backoff_factor'] == 1 | |||||
| def test_register_ttypes_filter( | |||||
| mock_pkg_resources, local_sched_config, local_sched_configfile): | |||||
| """Filtering on one worker should only register its associated task type | |||||
| """ | |||||
| result = CliRunner().invoke(cli, [ | |||||
| '--config-file', local_sched_configfile, | |||||
| 'task-type', 'register', '--plugins', 'lister.gnu' | |||||
| ]) | |||||
| assert result.exit_code == 0, traceback.print_exception(*result.exc_info) | |||||
| scheduler = get_scheduler(**local_sched_config['scheduler']) | |||||
| all_tasks = [ | |||||
| 'list-gnu-full', | |||||
| ] | |||||
| for task in all_tasks: | |||||
| task_type_desc = scheduler.get_task_type(task) | |||||
| assert task_type_desc | |||||
| assert task_type_desc['type'] == task | |||||
| assert task_type_desc['backoff_factor'] == 1 | |||||