From f3a90e0d004b0f7d0f788d900388618122ed213b Mon Sep 17 00:00:00 2001 From: Aleksandr Mezin Date: Fri, 1 Nov 2024 01:37:49 +0200 Subject: [PATCH] Allow custom scheduler names in `--dist` command line argument There seems to be no way to fix the argument validation, at least without complicating things unnecessarily. So simply remove it. Instead, check `pytest_xdist_make_scheduler` return value. Also, convert every builtin scheduler to a separate plugin. Fixes: https://github.com/pytest-dev/pytest-xdist/issues/970 --- changelog/970.feature | 1 + src/xdist/dsession.py | 41 +++++++------------------------- src/xdist/plugin.py | 16 ++++++------- src/xdist/scheduler/each.py | 12 ++++++++++ src/xdist/scheduler/load.py | 12 ++++++++++ src/xdist/scheduler/loadfile.py | 12 ++++++++++ src/xdist/scheduler/loadgroup.py | 12 ++++++++++ src/xdist/scheduler/loadscope.py | 12 ++++++++++ src/xdist/scheduler/worksteal.py | 12 ++++++++++ testing/acceptance_test.py | 12 ++++++++++ 10 files changed, 100 insertions(+), 42 deletions(-) create mode 100644 changelog/970.feature diff --git a/changelog/970.feature b/changelog/970.feature new file mode 100644 index 00000000..f8333e74 --- /dev/null +++ b/changelog/970.feature @@ -0,0 +1 @@ +`--dist` option allows custom scheduler names now. diff --git a/src/xdist/dsession.py b/src/xdist/dsession.py index 62079a28..ef91cb32 100644 --- a/src/xdist/dsession.py +++ b/src/xdist/dsession.py @@ -14,13 +14,7 @@ from xdist.remote import Producer from xdist.remote import WorkerInfo -from xdist.scheduler import EachScheduling -from xdist.scheduler import LoadFileScheduling -from xdist.scheduler import LoadGroupScheduling -from xdist.scheduler import LoadScheduling -from xdist.scheduler import LoadScopeScheduling from xdist.scheduler import Scheduling -from xdist.scheduler import WorkStealingScheduling from xdist.workermanage import NodeManager from xdist.workermanage import WorkerController @@ -81,11 +75,18 @@ def report_line(self, line: str) -> None: @pytest.hookimpl(trylast=True) def pytest_sessionstart(self, session: pytest.Session) -> None: - """Creates and starts the nodes. + """Initializes the scheduler, creates and starts the nodes. The nodes are setup to put their events onto self.queue. As soon as nodes start they will emit the worker_workerready event. """ + self.sched = self.config.hook.pytest_xdist_make_scheduler( + config=self.config, log=self.log + ) + if self.sched is None: + dist = self.config.getoption("dist") + raise pytest.UsageError(f"pytest-xdist: scheduler {dist!r} not found") + self.nodemanager = NodeManager(self.config) nodes = self.nodemanager.setup_nodes(putevent=self.queue.put) self._active_nodes.update(nodes) @@ -104,34 +105,8 @@ def pytest_collection(self) -> bool: # prohibit collection of test items in controller process return True - @pytest.hookimpl(trylast=True) - def pytest_xdist_make_scheduler( - self, - config: pytest.Config, - log: Producer, - ) -> Scheduling | None: - dist = config.getvalue("dist") - if dist == "each": - return EachScheduling(config, log) - if dist == "load": - return LoadScheduling(config, log) - if dist == "loadscope": - return LoadScopeScheduling(config, log) - if dist == "loadfile": - return LoadFileScheduling(config, log) - if dist == "loadgroup": - return LoadGroupScheduling(config, log) - if dist == "worksteal": - return WorkStealingScheduling(config, log) - return None - @pytest.hookimpl def pytest_runtestloop(self) -> bool: - self.sched = self.config.hook.pytest_xdist_make_scheduler( - config=self.config, log=self.log - ) - assert self.sched is not None - self.shouldstop = False pending_exception = None while not self.session_finished: diff --git a/src/xdist/plugin.py b/src/xdist/plugin.py index f670d9de..1953cc42 100644 --- a/src/xdist/plugin.py +++ b/src/xdist/plugin.py @@ -101,15 +101,6 @@ def pytest_addoption(parser: pytest.Parser) -> None: "--dist", metavar="distmode", action="store", - choices=[ - "each", - "load", - "loadscope", - "loadfile", - "loadgroup", - "worksteal", - "no", - ], dest="dist", default="no", help=( @@ -235,6 +226,13 @@ def pytest_configure(config: pytest.Config) -> None: # Create the distributed session in case we have a valid distribution # mode and test environments. if _is_distribution_mode(config): + config.pluginmanager.import_plugin("xdist.scheduler.each") + config.pluginmanager.import_plugin("xdist.scheduler.load") + config.pluginmanager.import_plugin("xdist.scheduler.loadfile") + config.pluginmanager.import_plugin("xdist.scheduler.loadgroup") + config.pluginmanager.import_plugin("xdist.scheduler.loadscope") + config.pluginmanager.import_plugin("xdist.scheduler.worksteal") + from xdist.dsession import DSession session = DSession(config) diff --git a/src/xdist/scheduler/each.py b/src/xdist/scheduler/each.py index aa4f7ba1..882d6c17 100644 --- a/src/xdist/scheduler/each.py +++ b/src/xdist/scheduler/each.py @@ -6,6 +6,7 @@ from xdist.remote import Producer from xdist.report import report_collection_diff +from xdist.scheduler.protocol import Scheduling from xdist.workermanage import parse_spec_config from xdist.workermanage import WorkerController @@ -150,3 +151,14 @@ def schedule(self) -> None: else: node.send_runtest_some(pending) self._started.append(node) + + +@pytest.hookimpl(trylast=True) +def pytest_xdist_make_scheduler( + config: pytest.Config, + log: Producer, +) -> Scheduling | None: + if config.getoption("dist") == "each": + return EachScheduling(config, log) + else: + return None diff --git a/src/xdist/scheduler/load.py b/src/xdist/scheduler/load.py index 9d153bb9..2d2c73a9 100644 --- a/src/xdist/scheduler/load.py +++ b/src/xdist/scheduler/load.py @@ -7,6 +7,7 @@ from xdist.remote import Producer from xdist.report import report_collection_diff +from xdist.scheduler.protocol import Scheduling from xdist.workermanage import parse_spec_config from xdist.workermanage import WorkerController @@ -333,3 +334,14 @@ def _check_nodes_have_same_collection(self) -> bool: self.config.hook.pytest_collectreport(report=rep) return same_collection + + +@pytest.hookimpl(trylast=True) +def pytest_xdist_make_scheduler( + config: pytest.Config, + log: Producer, +) -> Scheduling | None: + if config.getoption("dist") == "load": + return LoadScheduling(config, log) + else: + return None diff --git a/src/xdist/scheduler/loadfile.py b/src/xdist/scheduler/loadfile.py index fb6f027f..bca315e2 100644 --- a/src/xdist/scheduler/loadfile.py +++ b/src/xdist/scheduler/loadfile.py @@ -3,6 +3,7 @@ import pytest from xdist.remote import Producer +from xdist.scheduler.protocol import Scheduling from .loadscope import LoadScopeScheduling @@ -58,3 +59,14 @@ def _split_scope(self, nodeid: str) -> str: example/loadsuite/epsilon/__init__.py """ return nodeid.split("::", 1)[0] + + +@pytest.hookimpl(trylast=True) +def pytest_xdist_make_scheduler( + config: pytest.Config, + log: Producer, +) -> Scheduling | None: + if config.getoption("dist") == "loadfile": + return LoadFileScheduling(config, log) + else: + return None diff --git a/src/xdist/scheduler/loadgroup.py b/src/xdist/scheduler/loadgroup.py index 798c7128..65892707 100644 --- a/src/xdist/scheduler/loadgroup.py +++ b/src/xdist/scheduler/loadgroup.py @@ -3,6 +3,7 @@ import pytest from xdist.remote import Producer +from xdist.scheduler.protocol import Scheduling from .loadscope import LoadScopeScheduling @@ -57,3 +58,14 @@ def _split_scope(self, nodeid: str) -> str: return nodeid.split("@")[-1] else: return nodeid + + +@pytest.hookimpl(trylast=True) +def pytest_xdist_make_scheduler( + config: pytest.Config, + log: Producer, +) -> Scheduling | None: + if config.getoption("dist") == "loadgroup": + return LoadGroupScheduling(config, log) + else: + return None diff --git a/src/xdist/scheduler/loadscope.py b/src/xdist/scheduler/loadscope.py index a4d63b29..164dcd85 100644 --- a/src/xdist/scheduler/loadscope.py +++ b/src/xdist/scheduler/loadscope.py @@ -8,6 +8,7 @@ from xdist.remote import Producer from xdist.report import report_collection_diff +from xdist.scheduler.protocol import Scheduling from xdist.workermanage import parse_spec_config from xdist.workermanage import WorkerController @@ -432,3 +433,14 @@ def _check_nodes_have_same_collection(self) -> bool: self.config.hook.pytest_collectreport(report=rep) return same_collection + + +@pytest.hookimpl(trylast=True) +def pytest_xdist_make_scheduler( + config: pytest.Config, + log: Producer, +) -> Scheduling | None: + if config.getoption("dist") == "loadscope": + return LoadScopeScheduling(config, log) + else: + return None diff --git a/src/xdist/scheduler/worksteal.py b/src/xdist/scheduler/worksteal.py index fd208486..48df9d05 100644 --- a/src/xdist/scheduler/worksteal.py +++ b/src/xdist/scheduler/worksteal.py @@ -7,6 +7,7 @@ from xdist.remote import Producer from xdist.report import report_collection_diff +from xdist.scheduler.protocol import Scheduling from xdist.workermanage import parse_spec_config from xdist.workermanage import WorkerController @@ -343,3 +344,14 @@ def _check_nodes_have_same_collection(self) -> bool: self.config.hook.pytest_collectreport(report=rep) return same_collection + + +@pytest.hookimpl(trylast=True) +def pytest_xdist_make_scheduler( + config: pytest.Config, + log: Producer, +) -> Scheduling | None: + if config.getoption("dist") == "worksteal": + return WorkStealingScheduling(config, log) + else: + return None diff --git a/testing/acceptance_test.py b/testing/acceptance_test.py index 3ef10cc9..51acaa31 100644 --- a/testing/acceptance_test.py +++ b/testing/acceptance_test.py @@ -1641,3 +1641,15 @@ def test(): ) result = pytester.runpytest() assert result.ret == 0 + + +def test_dist_validation(pytester: pytest.Pytester) -> None: + """Should exit early if incorrect --dist value is specified.""" + f = pytester.makepyfile( + """ + assert 0 + """ + ) + result = pytester.runpytest(f, "-n1", "--dist=invalid") + assert result.ret == pytest.ExitCode.USAGE_ERROR + result.stderr.fnmatch_lines(["ERROR: pytest-xdist: scheduler 'invalid' not found"])