Allow for automatic use of pytest-xdist (pantsbuild#16499)

Closes pantsbuild#15026.

[ci skip-rust]
[ci skip-build-wheels]

Co-authored-by: Stu Hood <[email protected]>

src/python/pants/backend/helm/goals/BUILD

@@ -3,4 +3,4 @@
 
 python_sources()
 
-python_tests(name="tests", overrides={"deploy_test.py": {"timeout": 120}})
+python_tests(name="tests", overrides={"deploy_test.py": {"timeout": 180}})

src/python/pants/backend/python/goals/pytest_runner.py

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import logging
+import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Optional, Tuple
@@ -44,6 +45,7 @@
     EMPTY_DIGEST,
     CreateDigest,
     Digest,
+    DigestContents,
     DigestSubset,
     Directory,
     MergeDigests,
@@ -164,6 +166,13 @@ class TestSetup:
     __test__ = False
 
 
+_TEST_PATTERN = re.compile(b"def\\s+test_")
+
+
+def _count_pytest_tests(contents: DigestContents) -> int:
+    return sum(len(_TEST_PATTERN.findall(file.content)) for file in contents)
+
+
 @rule(level=LogLevel.DEBUG)
 async def setup_pytest_for_target(
     request: TestSetupRequest,
@@ -325,12 +334,22 @@ async def setup_pytest_for_target(
     cache_scope = (
         ProcessCacheScope.PER_SESSION if test_subsystem.force else ProcessCacheScope.SUCCESSFUL
     )
+
+    xdist_concurrency = 0
+    if pytest.xdist_enabled and not request.is_debug:
+        concurrency = request.field_set.xdist_concurrency.value
+        if concurrency is None:
+            contents = await Get(DigestContents, Digest, field_set_source_files.snapshot.digest)
+            concurrency = _count_pytest_tests(contents)
+        xdist_concurrency = concurrency
+
     process = await Get(
         Process,
         VenvPexProcess(
             pytest_runner_pex,
             argv=(
                 *(("-c", pytest.config) if pytest.config else ()),
+                *(("-n", "{pants_concurrency}") if xdist_concurrency else ()),
                 *request.prepend_argv,
                 *pytest.args,
                 *coverage_args,
@@ -344,6 +363,7 @@ async def setup_pytest_for_target(
                 test_subsystem, pytest
             ),
             execution_slot_variable=pytest.execution_slot_var,
+            concurrency_available=xdist_concurrency,
             description=f"Run Pytest for {request.field_set.address}",
             level=LogLevel.DEBUG,
             cache_scope=cache_scope,

src/python/pants/backend/python/goals/pytest_runner_integration_test.py

@@ -39,6 +39,7 @@
 from pants.core.util_rules import config_files, distdir
 from pants.engine.addresses import Address
 from pants.engine.fs import CreateDigest, Digest, DigestContents, FileContent
+from pants.engine.process import InteractiveProcessResult
 from pants.engine.rules import Get, rule
 from pants.engine.target import Target
 from pants.engine.unions import UnionRule
@@ -94,19 +95,28 @@ def test():
 )
 
 
-def run_pytest(
+def _configure_pytest_runner(
     rule_runner: RuleRunner,
-    test_target: Target,
     *,
     extra_args: list[str] | None = None,
     env: dict[str, str] | None = None,
-) -> TestResult:
+) -> None:
     args = [
         "--backend-packages=pants.backend.python",
         f"--source-root-patterns={SOURCE_ROOT}",
         *(extra_args or ()),
     ]
     rule_runner.set_options(args, env=env, env_inherit={"PATH", "PYENV_ROOT", "HOME"})
+
+
+def run_pytest(
+    rule_runner: RuleRunner,
+    test_target: Target,
+    *,
+    extra_args: list[str] | None = None,
+    env: dict[str, str] | None = None,
+) -> TestResult:
+    _configure_pytest_runner(rule_runner, extra_args=extra_args, env=env)
     inputs = [PythonTestFieldSet.create(test_target)]
     test_result = rule_runner.request(TestResult, inputs)
     debug_request = rule_runner.request(TestDebugRequest, inputs)
@@ -117,6 +127,32 @@ def run_pytest(
     return test_result
 
 
+def run_pytest_noninteractive(
+    rule_runner: RuleRunner,
+    test_target: Target,
+    *,
+    extra_args: list[str] | None = None,
+    env: dict[str, str] | None = None,
+) -> TestResult:
+    _configure_pytest_runner(rule_runner, extra_args=extra_args, env=env)
+    inputs = [PythonTestFieldSet.create(test_target)]
+    return rule_runner.request(TestResult, inputs)
+
+
+def run_pytest_interactive(
+    rule_runner: RuleRunner,
+    test_target: Target,
+    *,
+    extra_args: list[str] | None = None,
+    env: dict[str, str] | None = None,
+) -> InteractiveProcessResult:
+    _configure_pytest_runner(rule_runner, extra_args=extra_args, env=env)
+    inputs = [PythonTestFieldSet.create(test_target)]
+    debug_request = rule_runner.request(TestDebugRequest, inputs)
+    with mock_console(rule_runner.options_bootstrapper):
+        return rule_runner.run_interactive_process(debug_request.process)
+
+
 @pytest.mark.platform_specific_behavior
 @pytest.mark.parametrize(
     "major_minor_interpreter",
@@ -231,7 +267,11 @@ def test() -> None:
             ),
         }
     )
-    extra_args = ["--pytest-version=pytest>=4.6.6,<4.7", "--pytest-lockfile=<none>"]
+    extra_args = [
+        "--pytest-version=pytest>=4.6.6,<4.7",
+        "--pytest-extra-requirements=[]",
+        "--pytest-lockfile=<none>",
+    ]
 
     py2_tgt = rule_runner.get_target(
         Address(PACKAGE, target_name="py2", relative_file_path="tests.py")
@@ -270,6 +310,72 @@ def test_ignore_me():
     assert "collected 2 items / 1 deselected / 1 selected" in result.stdout
 
 
+def test_xdist_enabled_noninteractive(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/tests.py": dedent(
+                """\
+                import os
+
+                def test_worker_id_set():
+                  assert "PYTEST_XDIST_WORKER" in os.environ
+
+                def test_worker_count_set():
+                  assert "PYTEST_XDIST_WORKER_COUNT" in os.environ
+                """
+            ),
+            f"{PACKAGE}/BUILD": "python_tests(xdist_concurrency=2)",
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="tests.py"))
+    result = run_pytest_noninteractive(rule_runner, tgt, extra_args=["--pytest-xdist-enabled"])
+    assert result.exit_code == 0
+
+
+def test_xdist_enabled_but_disabled_for_target(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/tests.py": dedent(
+                """\
+                import os
+
+                def test_worker_id_not_set():
+                  assert "PYTEST_XDIST_WORKER" not in os.environ
+
+                def test_worker_count_not_set():
+                  assert "PYTEST_XDIST_WORKER_COUNT" not in os.environ
+                """
+            ),
+            f"{PACKAGE}/BUILD": "python_tests(xdist_concurrency=0)",
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="tests.py"))
+    result = run_pytest_noninteractive(rule_runner, tgt, extra_args=["--pytest-xdist-enabled"])
+    assert result.exit_code == 0
+
+
+def test_xdist_enabled_interactive(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/tests.py": dedent(
+                """\
+                import os
+
+                def test_worker_id_not_set():
+                  assert "PYTEST_XDIST_WORKER" not in os.environ
+
+                def test_worker_count_not_set():
+                  assert "PYTEST_XDIST_WORKER_COUNT" not in os.environ
+                """
+            ),
+            f"{PACKAGE}/BUILD": "python_tests(xdist_concurrency=2)",
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="tests.py"))
+    result = run_pytest_interactive(rule_runner, tgt, extra_args=["--pytest-xdist-enabled"])
+    assert result.exit_code == 0
+
+
 @pytest.mark.parametrize(
     "config_path,extra_args",
     (["pytest.ini", []], ["custom_config.ini", ["--pytest-config=custom_config.ini"]]),

src/python/pants/backend/python/goals/pytest_runner_test.py

@@ -0,0 +1,56 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from pants.backend.python.goals.pytest_runner import _count_pytest_tests
+from pants.engine.fs import DigestContents, FileContent
+
+EXAMPLE_TEST1 = b"""
+def test_foo():
+    pass
+
+def test_bar():
+    pass
+"""
+
+EXAMPLE_TEST2 = b"""
+class TestStuff(TestCase):
+    def test_baz():
+        pass
+
+    def testHelper():
+        pass
+"""
+
+
+def test_count_pytest_tests_empty() -> None:
+    digest_contents = DigestContents([FileContent(path="tests/test_empty.py", content=b"")])
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 0
+
+
+def test_count_pytest_tests_methods() -> None:
+    digest_contents = DigestContents(
+        [FileContent(path="tests/test_example1.py", content=EXAMPLE_TEST1)]
+    )
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 2
+
+
+def test_count_pytest_tests_in_class() -> None:
+    digest_contents = DigestContents(
+        [FileContent(path="tests/test_example1.py", content=EXAMPLE_TEST2)]
+    )
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 1
+
+
+def test_count_pytest_tests_multiple() -> None:
+    digest_contents = DigestContents(
+        [
+            FileContent(path="tests/test_empty.py", content=b""),
+            FileContent(path="tests/test_example1.py", content=EXAMPLE_TEST1),
+            FileContent(path="tests/test_example2.py", content=EXAMPLE_TEST2),
+        ]
+    )
+    test_count = _count_pytest_tests(digest_contents)
+    assert test_count == 3