Finish wiring up everything, add tests

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

@@ -63,6 +63,7 @@ class TestTargetSetup:
     timeout_seconds: Optional[int]
     xml_dir: Optional[str]
     junit_family: str
+    execution_slot_variable: str
 
     # Prevent this class from being detected by pytest as a test class.
     __test__ = False
@@ -201,6 +202,7 @@ async def setup_pytest_for_target(
         timeout_seconds=field_set.timeout.calculate_from_global_options(pytest),
         xml_dir=pytest.options.junit_xml_dir,
         junit_family=pytest.options.junit_family,
+        execution_slot_variable=pytest.options.execution_slot_var,
     )
 
 
@@ -247,6 +249,7 @@ async def run_python_test(
         description=f"Run Pytest for {field_set.address.reference()}",
         timeout_seconds=test_setup.timeout_seconds,
         env=env,
+        execution_slot_variable=test_setup.execution_slot_variable,
     )
     result = await Get(FallibleProcessResult, Process, process)
 

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

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import os
+import re
 from pathlib import Path, PurePath
 from textwrap import dedent
 from typing import List, Optional
@@ -140,6 +141,7 @@ def run_pytest(
         origin: Optional[OriginSpec] = None,
         junit_xml_dir: Optional[str] = None,
         use_coverage: bool = False,
+        execution_slot_var: Optional[str] = None,
     ) -> TestResult:
         args = [
             "--backend-packages=pants.backend.python",
@@ -154,6 +156,9 @@ def run_pytest(
             args.append(f"--pytest-junit-xml-dir={junit_xml_dir}")
         if use_coverage:
             args.append("--test-use-coverage")
+        if execution_slot_var:
+            args.append(f"--pytest-execution-slot-var={execution_slot_var}")
+
         options_bootstrapper = create_options_bootstrapper(args=args)
         address = Address(self.package, "target")
         if origin is None:
@@ -371,3 +376,22 @@ def test_coverage(self) -> None:
         assert result.status == Status.SUCCESS
         assert f"{self.package}/test_good.py ." in result.stdout
         assert result.coverage_data is not None
+
+    def test_execution_slot_variable(self) -> None:
+        source = FileContent(
+            path="test_concurrency_slot.py",
+            content=dedent(
+                """\
+                import os
+
+                def test_fail_printing_slot_env_var():
+                    slot = os.getenv("SLOT")
+                    print(f"Value of slot is {slot}")
+                    # Deliberately fail the test so the SLOT output gets printed to stdout
+                    assert 1 == 2
+                """
+            ).encode(),
+        )
+        self.create_python_test_target([source])
+        result = self.run_pytest(execution_slot_var="SLOT")
+        assert re.search(r"Value of slot is \d+", result.stdout)

src/python/pants/backend/python/subsystems/pytest.py

@@ -76,6 +76,13 @@ def register_options(cls, register):
             advanced=True,
             help="The format of the generated XML file. See https://docs.pytest.org/en/latest/reference.html#confval-junit_family.",
         )
+        register(
+            "--execution-slot-var",
+            type=str,
+            default="",
+            advanced=True,
+            help="If a non-empty string, the process execution slot id (an integer) will be exposed to tests under this environment variable name.",
+        )
 
     def get_requirement_strings(self) -> Tuple[str, ...]:
         """Returns a tuple of requirements-style strings for Pytest and Pytest plugins."""

src/python/pants/engine/process.py

@@ -33,7 +33,7 @@ class Process:
     timeout_seconds: Union[int, float]
     jdk_home: Optional[str]
     is_nailgunnable: bool
-    execution_slot_variable: str = ""
+    execution_slot_variable: str
 
     def __init__(
         self,
@@ -49,6 +49,7 @@ def __init__(
         timeout_seconds: Optional[Union[int, float]] = None,
         jdk_home: Optional[str] = None,
         is_nailgunnable: bool = False,
+        execution_slot_variable: str = "",
     ) -> None:
         """Request to run a subprocess, similar to subprocess.Popen.
 
@@ -87,6 +88,7 @@ def __init__(
         self.timeout_seconds = timeout_seconds if timeout_seconds and timeout_seconds > 0 else -1
         self.jdk_home = jdk_home
         self.is_nailgunnable = is_nailgunnable
+        self.execution_slot_variable = execution_slot_variable
 
 
 @frozen_after_init