Allowpython_tests() to set environment variables

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

@@ -13,7 +13,12 @@
     PytestCoverageData,
 )
 from pants.backend.python.subsystems.pytest import PyTest
-from pants.backend.python.target_types import ConsoleScript, PythonTestsSources, PythonTestsTimeout
+from pants.backend.python.target_types import (
+    ConsoleScript,
+    PythonTestsExtraEnvVars,
+    PythonTestsSources,
+    PythonTestsTimeout,
+)
 from pants.backend.python.util_rules.pex import (
     Pex,
     PexInterpreterConstraints,
@@ -39,6 +44,7 @@
 )
 from pants.core.util_rules.config_files import ConfigFiles, ConfigFilesRequest
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
+from pants.engine.environment import CompleteEnvironment
 from pants.engine.fs import (
     AddPrefix,
     CreateDigest,
@@ -79,6 +85,7 @@ class PythonTestFieldSet(TestFieldSet):
     sources: PythonTestsSources
     timeout: PythonTestsTimeout
     runtime_package_dependencies: RuntimePackageDependenciesField
+    extra_env_vars: PythonTestsExtraEnvVars
 
     def is_conftest_or_type_stub(self) -> bool:
         """We skip both `conftest.py` and `.pyi` stubs, even though though they often belong to a
@@ -114,6 +121,7 @@ async def setup_pytest_for_target(
     coverage_subsystem: CoverageSubsystem,
     test_extra_env: TestExtraEnv,
     global_options: GlobalOptions,
+    complete_env: CompleteEnvironment,
 ) -> TestSetup:
     transitive_targets = await Get(
         TransitiveTargets, TransitiveTargetsRequest([request.field_set.address])
@@ -225,6 +233,8 @@ async def setup_pytest_for_target(
         "PYTEST_ADDOPTS": " ".join(add_opts),
         "PEX_EXTRA_SYS_PATH": ":".join(prepared_sources.source_roots),
         **test_extra_env.env,
+        # NOTE: `complete_env` intentionally after `test_extra_env` to allow overriding within `python_tests`
+        **complete_env.get_subset(request.field_set.extra_env_vars.value or ()),
     }
 
     # Cache test runs only if they are successful, or not at all if `--test-force`.

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

@@ -384,19 +384,38 @@ def test_extra_env_vars(rule_runner: RuleRunner) -> None:
                 import os
 
                 def test_args():
-                    assert os.getenv("SOME_VAR") == "some_value"
-                    assert os.getenv("OTHER_VAR") == "other_value"
+                    assert os.getenv("ARG_WITH_VALUE_VAR") == "arg_with_value_var"
+                    assert os.getenv("ARG_WITHOUT_VALUE_VAR") == "arg_without_value_value"
+                    assert os.getenv("PYTHON_TESTS_VAR_WITH_VALUE") == "python_tests_var_with_value"
+                    assert os.getenv("PYTHON_TESTS_VAR_WITHOUT_VALUE") == "python_tests_var_without_value"
+                    assert os.getenv("PYTHON_TESTS_OVERRIDE_WITH_VALUE_VAR") == "python_tests_override_with_value_var_override"
                 """
             ),
-            f"{PACKAGE}/BUILD": "python_tests()",
+            f"{PACKAGE}/BUILD": dedent(
+                """\
+            python_tests(
+                extra_env_vars=(
+                    "PYTHON_TESTS_VAR_WITHOUT_VALUE",
+                    "PYTHON_TESTS_VAR_WITH_VALUE=python_tests_var_with_value",
+                    "PYTHON_TESTS_OVERRIDE_WITH_VALUE_VAR=python_tests_override_with_value_var_override",
+                )
+            )
+            """
+            ),
         }
     )
     tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="test_extra_env_vars.py"))
     result = run_pytest(
         rule_runner,
         tgt,
-        extra_args=['--test-extra-env-vars=["SOME_VAR=some_value", "OTHER_VAR"]'],
-        env={"OTHER_VAR": "other_value"},
+        extra_args=[
+            '--test-extra-env-vars=["ARG_WITH_VALUE_VAR=arg_with_value_var", "ARG_WITHOUT_VALUE_VAR", "PYTHON_TESTS_OVERRIDE_ARG_WITH_VALUE_VAR"]'
+        ],
+        env={
+            "ARG_WITHOUT_VALUE_VAR": "arg_without_value_value",
+            "PYTHON_TESTS_VAR_WITHOUT_VALUE": "python_tests_var_without_value",
+            "PYTHON_TESTS_OVERRIDE_WITH_VALUE_VAR": "python_tests_override_with_value_var",
+        },
     )
     assert result.exit_code == 0
 

src/python/pants/backend/python/target_types.py

@@ -428,6 +428,16 @@ def calculate_from_global_options(self, pytest: PyTest) -> Optional[int]:
         return result
 
 
+class PythonTestsExtraEnvVars(StringSequenceField):
+    alias = "extra_env_vars"
+    help = (
+        "Additional environment variables to include in test processes. "
+        "Entries are strings in the form `ENV_VAR=value` to use explicitly; or just "
+        "`ENV_VAR` to copy the value of a variable in Pants's own environment. "
+        "This will be merged with and override values from [test].extra_env_vars."
+    )
+
+
 class PythonTests(Target):
     alias = "python_tests"
     core_fields = (
@@ -437,6 +447,7 @@ class PythonTests(Target):
         PythonTestsDependencies,
         PythonTestsTimeout,
         RuntimePackageDependenciesField,
+        PythonTestsExtraEnvVars,
     )
     help = (
         "Python tests, written in either Pytest style or unittest style.\n\nAll test util code, "