Skip to content

Commit f3a90e0

Browse files
committed
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: #970
1 parent 9788f12 commit f3a90e0

File tree

10 files changed

+100
-42
lines changed

10 files changed

+100
-42
lines changed

changelog/970.feature

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
`--dist` option allows custom scheduler names now.

src/xdist/dsession.py

Lines changed: 8 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,7 @@
1414

1515
from xdist.remote import Producer
1616
from xdist.remote import WorkerInfo
17-
from xdist.scheduler import EachScheduling
18-
from xdist.scheduler import LoadFileScheduling
19-
from xdist.scheduler import LoadGroupScheduling
20-
from xdist.scheduler import LoadScheduling
21-
from xdist.scheduler import LoadScopeScheduling
2217
from xdist.scheduler import Scheduling
23-
from xdist.scheduler import WorkStealingScheduling
2418
from xdist.workermanage import NodeManager
2519
from xdist.workermanage import WorkerController
2620

@@ -81,11 +75,18 @@ def report_line(self, line: str) -> None:
8175

8276
@pytest.hookimpl(trylast=True)
8377
def pytest_sessionstart(self, session: pytest.Session) -> None:
84-
"""Creates and starts the nodes.
78+
"""Initializes the scheduler, creates and starts the nodes.
8579
8680
The nodes are setup to put their events onto self.queue. As
8781
soon as nodes start they will emit the worker_workerready event.
8882
"""
83+
self.sched = self.config.hook.pytest_xdist_make_scheduler(
84+
config=self.config, log=self.log
85+
)
86+
if self.sched is None:
87+
dist = self.config.getoption("dist")
88+
raise pytest.UsageError(f"pytest-xdist: scheduler {dist!r} not found")
89+
8990
self.nodemanager = NodeManager(self.config)
9091
nodes = self.nodemanager.setup_nodes(putevent=self.queue.put)
9192
self._active_nodes.update(nodes)
@@ -104,34 +105,8 @@ def pytest_collection(self) -> bool:
104105
# prohibit collection of test items in controller process
105106
return True
106107

107-
@pytest.hookimpl(trylast=True)
108-
def pytest_xdist_make_scheduler(
109-
self,
110-
config: pytest.Config,
111-
log: Producer,
112-
) -> Scheduling | None:
113-
dist = config.getvalue("dist")
114-
if dist == "each":
115-
return EachScheduling(config, log)
116-
if dist == "load":
117-
return LoadScheduling(config, log)
118-
if dist == "loadscope":
119-
return LoadScopeScheduling(config, log)
120-
if dist == "loadfile":
121-
return LoadFileScheduling(config, log)
122-
if dist == "loadgroup":
123-
return LoadGroupScheduling(config, log)
124-
if dist == "worksteal":
125-
return WorkStealingScheduling(config, log)
126-
return None
127-
128108
@pytest.hookimpl
129109
def pytest_runtestloop(self) -> bool:
130-
self.sched = self.config.hook.pytest_xdist_make_scheduler(
131-
config=self.config, log=self.log
132-
)
133-
assert self.sched is not None
134-
135110
self.shouldstop = False
136111
pending_exception = None
137112
while not self.session_finished:

src/xdist/plugin.py

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -101,15 +101,6 @@ def pytest_addoption(parser: pytest.Parser) -> None:
101101
"--dist",
102102
metavar="distmode",
103103
action="store",
104-
choices=[
105-
"each",
106-
"load",
107-
"loadscope",
108-
"loadfile",
109-
"loadgroup",
110-
"worksteal",
111-
"no",
112-
],
113104
dest="dist",
114105
default="no",
115106
help=(
@@ -235,6 +226,13 @@ def pytest_configure(config: pytest.Config) -> None:
235226
# Create the distributed session in case we have a valid distribution
236227
# mode and test environments.
237228
if _is_distribution_mode(config):
229+
config.pluginmanager.import_plugin("xdist.scheduler.each")
230+
config.pluginmanager.import_plugin("xdist.scheduler.load")
231+
config.pluginmanager.import_plugin("xdist.scheduler.loadfile")
232+
config.pluginmanager.import_plugin("xdist.scheduler.loadgroup")
233+
config.pluginmanager.import_plugin("xdist.scheduler.loadscope")
234+
config.pluginmanager.import_plugin("xdist.scheduler.worksteal")
235+
238236
from xdist.dsession import DSession
239237

240238
session = DSession(config)

src/xdist/scheduler/each.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from xdist.remote import Producer
88
from xdist.report import report_collection_diff
9+
from xdist.scheduler.protocol import Scheduling
910
from xdist.workermanage import parse_spec_config
1011
from xdist.workermanage import WorkerController
1112

@@ -150,3 +151,14 @@ def schedule(self) -> None:
150151
else:
151152
node.send_runtest_some(pending)
152153
self._started.append(node)
154+
155+
156+
@pytest.hookimpl(trylast=True)
157+
def pytest_xdist_make_scheduler(
158+
config: pytest.Config,
159+
log: Producer,
160+
) -> Scheduling | None:
161+
if config.getoption("dist") == "each":
162+
return EachScheduling(config, log)
163+
else:
164+
return None

src/xdist/scheduler/load.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from xdist.remote import Producer
99
from xdist.report import report_collection_diff
10+
from xdist.scheduler.protocol import Scheduling
1011
from xdist.workermanage import parse_spec_config
1112
from xdist.workermanage import WorkerController
1213

@@ -333,3 +334,14 @@ def _check_nodes_have_same_collection(self) -> bool:
333334
self.config.hook.pytest_collectreport(report=rep)
334335

335336
return same_collection
337+
338+
339+
@pytest.hookimpl(trylast=True)
340+
def pytest_xdist_make_scheduler(
341+
config: pytest.Config,
342+
log: Producer,
343+
) -> Scheduling | None:
344+
if config.getoption("dist") == "load":
345+
return LoadScheduling(config, log)
346+
else:
347+
return None

src/xdist/scheduler/loadfile.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import pytest
44

55
from xdist.remote import Producer
6+
from xdist.scheduler.protocol import Scheduling
67

78
from .loadscope import LoadScopeScheduling
89

@@ -58,3 +59,14 @@ def _split_scope(self, nodeid: str) -> str:
5859
example/loadsuite/epsilon/__init__.py
5960
"""
6061
return nodeid.split("::", 1)[0]
62+
63+
64+
@pytest.hookimpl(trylast=True)
65+
def pytest_xdist_make_scheduler(
66+
config: pytest.Config,
67+
log: Producer,
68+
) -> Scheduling | None:
69+
if config.getoption("dist") == "loadfile":
70+
return LoadFileScheduling(config, log)
71+
else:
72+
return None

src/xdist/scheduler/loadgroup.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import pytest
44

55
from xdist.remote import Producer
6+
from xdist.scheduler.protocol import Scheduling
67

78
from .loadscope import LoadScopeScheduling
89

@@ -57,3 +58,14 @@ def _split_scope(self, nodeid: str) -> str:
5758
return nodeid.split("@")[-1]
5859
else:
5960
return nodeid
61+
62+
63+
@pytest.hookimpl(trylast=True)
64+
def pytest_xdist_make_scheduler(
65+
config: pytest.Config,
66+
log: Producer,
67+
) -> Scheduling | None:
68+
if config.getoption("dist") == "loadgroup":
69+
return LoadGroupScheduling(config, log)
70+
else:
71+
return None

src/xdist/scheduler/loadscope.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
from xdist.remote import Producer
1010
from xdist.report import report_collection_diff
11+
from xdist.scheduler.protocol import Scheduling
1112
from xdist.workermanage import parse_spec_config
1213
from xdist.workermanage import WorkerController
1314

@@ -432,3 +433,14 @@ def _check_nodes_have_same_collection(self) -> bool:
432433
self.config.hook.pytest_collectreport(report=rep)
433434

434435
return same_collection
436+
437+
438+
@pytest.hookimpl(trylast=True)
439+
def pytest_xdist_make_scheduler(
440+
config: pytest.Config,
441+
log: Producer,
442+
) -> Scheduling | None:
443+
if config.getoption("dist") == "loadscope":
444+
return LoadScopeScheduling(config, log)
445+
else:
446+
return None

src/xdist/scheduler/worksteal.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
from xdist.remote import Producer
99
from xdist.report import report_collection_diff
10+
from xdist.scheduler.protocol import Scheduling
1011
from xdist.workermanage import parse_spec_config
1112
from xdist.workermanage import WorkerController
1213

@@ -343,3 +344,14 @@ def _check_nodes_have_same_collection(self) -> bool:
343344
self.config.hook.pytest_collectreport(report=rep)
344345

345346
return same_collection
347+
348+
349+
@pytest.hookimpl(trylast=True)
350+
def pytest_xdist_make_scheduler(
351+
config: pytest.Config,
352+
log: Producer,
353+
) -> Scheduling | None:
354+
if config.getoption("dist") == "worksteal":
355+
return WorkStealingScheduling(config, log)
356+
else:
357+
return None

testing/acceptance_test.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1641,3 +1641,15 @@ def test():
16411641
)
16421642
result = pytester.runpytest()
16431643
assert result.ret == 0
1644+
1645+
1646+
def test_dist_validation(pytester: pytest.Pytester) -> None:
1647+
"""Should exit early if incorrect --dist value is specified."""
1648+
f = pytester.makepyfile(
1649+
"""
1650+
assert 0
1651+
"""
1652+
)
1653+
result = pytester.runpytest(f, "-n1", "--dist=invalid")
1654+
assert result.ret == pytest.ExitCode.USAGE_ERROR
1655+
result.stderr.fnmatch_lines(["ERROR: pytest-xdist: scheduler 'invalid' not found"])

0 commit comments

Comments
 (0)