Skip to content

Commit

Permalink
Allow custom scheduler names in --dist command line argument
Browse files Browse the repository at this point in the history
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: pytest-dev#970
  • Loading branch information
amezin committed Nov 1, 2024
1 parent 9788f12 commit f3a90e0
Show file tree
Hide file tree
Showing 10 changed files with 100 additions and 42 deletions.
1 change: 1 addition & 0 deletions changelog/970.feature
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
`--dist` option allows custom scheduler names now.
41 changes: 8 additions & 33 deletions src/xdist/dsession.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand Down
16 changes: 7 additions & 9 deletions src/xdist/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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=(
Expand Down Expand Up @@ -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)
Expand Down
12 changes: 12 additions & 0 deletions src/xdist/scheduler/each.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
12 changes: 12 additions & 0 deletions src/xdist/scheduler/load.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
12 changes: 12 additions & 0 deletions src/xdist/scheduler/loadfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import pytest

from xdist.remote import Producer
from xdist.scheduler.protocol import Scheduling

from .loadscope import LoadScopeScheduling

Expand Down Expand Up @@ -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
12 changes: 12 additions & 0 deletions src/xdist/scheduler/loadgroup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import pytest

from xdist.remote import Producer
from xdist.scheduler.protocol import Scheduling

from .loadscope import LoadScopeScheduling

Expand Down Expand Up @@ -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
12 changes: 12 additions & 0 deletions src/xdist/scheduler/loadscope.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
12 changes: 12 additions & 0 deletions src/xdist/scheduler/worksteal.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
12 changes: 12 additions & 0 deletions testing/acceptance_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])

0 comments on commit f3a90e0

Please sign in to comment.