From a9a303f6888fbbd034613e35e75f80a47711d5b8 Mon Sep 17 00:00:00 2001 From: Alon Livne <2005alonlivne@gmail.com> Date: Sat, 11 Jan 2025 17:02:13 +0200 Subject: [PATCH] add proxy gateways with --px arg Proxy gateways do not run workers, and are meant to be passed with the `via` attribute to additional gateways. They are useful for running multiple workers on remote machines. Example usage: ``` pytest -sv --dist=load --px "id=my_proxy//socket=IP:PORT" --tx "5*popen//via=my_proxy" ``` Proxy gateways do not run workers, anda re meant to be passed --- src/xdist/plugin.py | 11 +++++++++++ src/xdist/scheduler/each.py | 4 ++-- src/xdist/scheduler/load.py | 4 ++-- src/xdist/scheduler/loadscope.py | 4 ++-- src/xdist/scheduler/worksteal.py | 4 ++-- src/xdist/workermanage.py | 18 ++++++++++++++---- testing/test_plugin.py | 4 ++-- 7 files changed, 35 insertions(+), 14 deletions(-) diff --git a/src/xdist/plugin.py b/src/xdist/plugin.py index f670d9d..0cf90f8 100644 --- a/src/xdist/plugin.py +++ b/src/xdist/plugin.py @@ -139,6 +139,17 @@ def pytest_addoption(parser: pytest.Parser) -> None: "--tx ssh=user@codespeak.net//chdir=testcache" ), ) + group.addoption( + "--px", + dest="px", + action="append", + default=[], + metavar="xspec", + help=( + "Add a proxy gateway to pass to test execution environments using `via`. Example:\n" + "--px id=my_proxy//socket=192.168.1.102:8888 --tx 5*popen//via=my_proxy" + ), + ) group._addoption( "-d", action="store_true", diff --git a/src/xdist/scheduler/each.py b/src/xdist/scheduler/each.py index 8552b53..bf91451 100644 --- a/src/xdist/scheduler/each.py +++ b/src/xdist/scheduler/each.py @@ -6,7 +6,7 @@ import pytest from xdist.remote import Producer from xdist.report import report_collection_diff -from xdist.workermanage import parse_spec_config +from xdist.workermanage import parse_tx_spec_config from xdist.workermanage import WorkerController @@ -26,7 +26,7 @@ class EachScheduling: def __init__(self, config: pytest.Config, log: Producer | None = None) -> None: self.config = config - self.numnodes = len(parse_spec_config(config)) + self.numnodes = len(parse_tx_spec_config(config)) self.node2collection: dict[WorkerController, list[str]] = {} self.node2pending: dict[WorkerController, list[int]] = {} self._started: list[WorkerController] = [] diff --git a/src/xdist/scheduler/load.py b/src/xdist/scheduler/load.py index ac011e5..56f6092 100644 --- a/src/xdist/scheduler/load.py +++ b/src/xdist/scheduler/load.py @@ -7,7 +7,7 @@ import pytest from xdist.remote import Producer from xdist.report import report_collection_diff -from xdist.workermanage import parse_spec_config +from xdist.workermanage import parse_tx_spec_config from xdist.workermanage import WorkerController @@ -58,7 +58,7 @@ class LoadScheduling: """ def __init__(self, config: pytest.Config, log: Producer | None = None) -> None: - self.numnodes = len(parse_spec_config(config)) + self.numnodes = len(parse_tx_spec_config(config)) self.node2collection: dict[WorkerController, list[str]] = {} self.node2pending: dict[WorkerController, list[int]] = {} self.pending: list[int] = [] diff --git a/src/xdist/scheduler/loadscope.py b/src/xdist/scheduler/loadscope.py index ee3f2fd..114561b 100644 --- a/src/xdist/scheduler/loadscope.py +++ b/src/xdist/scheduler/loadscope.py @@ -8,7 +8,7 @@ import pytest from xdist.remote import Producer from xdist.report import report_collection_diff -from xdist.workermanage import parse_spec_config +from xdist.workermanage import parse_tx_spec_config from xdist.workermanage import WorkerController @@ -91,7 +91,7 @@ class LoadScopeScheduling: """ def __init__(self, config: pytest.Config, log: Producer | None = None) -> None: - self.numnodes = len(parse_spec_config(config)) + self.numnodes = len(parse_tx_spec_config(config)) self.collection: list[str] | None = None self.workqueue: OrderedDict[str, dict[str, bool]] = OrderedDict() diff --git a/src/xdist/scheduler/worksteal.py b/src/xdist/scheduler/worksteal.py index 28708fc..550372e 100644 --- a/src/xdist/scheduler/worksteal.py +++ b/src/xdist/scheduler/worksteal.py @@ -7,7 +7,7 @@ import pytest from xdist.remote import Producer from xdist.report import report_collection_diff -from xdist.workermanage import parse_spec_config +from xdist.workermanage import parse_tx_spec_config from xdist.workermanage import WorkerController @@ -65,7 +65,7 @@ class WorkStealingScheduling: """ def __init__(self, config: pytest.Config, log: Producer | None = None) -> None: - self.numnodes = len(parse_spec_config(config)) + self.numnodes = len(parse_tx_spec_config(config)) self.node2collection: dict[WorkerController, list[str]] = {} self.node2pending: dict[WorkerController, list[int]] = {} self.pending: list[int] = [] diff --git a/src/xdist/workermanage.py b/src/xdist/workermanage.py index 08ba243..ca87381 100644 --- a/src/xdist/workermanage.py +++ b/src/xdist/workermanage.py @@ -23,7 +23,7 @@ from xdist.remote import Producer from xdist.remote import WorkerInfo -def parse_spec_config(config: pytest.Config) -> list[str]: +def parse_tx_spec_config(config: pytest.Config) -> list[str]: xspeclist = [] tx: list[str] = config.getvalue("tx") for xspec in tx: @@ -57,8 +57,15 @@ class NodeManager: if self.testrunuid is None: self.testrunuid = uuid.uuid4().hex self.group = execnet.Group(execmodel="main_thread_only") + for proxy_spec in self._getpxspecs(): + # Proxy gateways do not run workers, and are meant to be passed with the `via` attribute + # to additional gateways. + # They are useful for running multiple workers on remote machines. + if getattr(proxy_spec, "id", None) is None: + raise pytest.UsageError(f"Proxy gateway {proxy_spec} must include an id") + self.group.makegateway(proxy_spec) if specs is None: - specs = self._getxspecs() + specs = self._gettxspecs() self.specs: list[execnet.XSpec] = [] for spec in specs: if not isinstance(spec, execnet.XSpec): @@ -107,8 +114,11 @@ class NodeManager: def teardown_nodes(self) -> None: self.group.terminate(self.EXIT_TIMEOUT) - def _getxspecs(self) -> list[execnet.XSpec]: - return [execnet.XSpec(x) for x in parse_spec_config(self.config)] + def _gettxspecs(self) -> list[execnet.XSpec]: + return [execnet.XSpec(x) for x in parse_tx_spec_config(self.config)] + + def _getpxspecs(self) -> list[execnet.XSpec]: + return [execnet.XSpec(x) for x in self.config.getoption("px")] def _getrsyncdirs(self) -> list[Path]: for spec in self.specs: diff --git a/testing/test_plugin.py b/testing/test_plugin.py index f352667..81d9e87 100644 --- a/testing/test_plugin.py +++ b/testing/test_plugin.py @@ -295,7 +295,7 @@ class TestDistOptions: def test_getxspecs(self, pytester: pytest.Pytester) -> None: config = pytester.parseconfigure("--tx=popen", "--tx", "ssh=xyz") nodemanager = NodeManager(config) - xspecs = nodemanager._getxspecs() + xspecs = nodemanager._gettxspecs() assert len(xspecs) == 2 print(xspecs) assert xspecs[0].popen @@ -303,7 +303,7 @@ class TestDistOptions: def test_xspecs_multiplied(self, pytester: pytest.Pytester) -> None: config = pytester.parseconfigure("--tx=3*popen") - xspecs = NodeManager(config)._getxspecs() + xspecs = NodeManager(config)._gettxspecs() assert len(xspecs) == 3 assert xspecs[1].popen