Repository: pantsbuild/pants

Issue:
Add support for setting per-test environment variables
Our support for setting environment variables in tests is currently global (see https://www.pantsbuild.org/docs/python-test-goal#setting-environment-variables).

We should add support for per-`python_tests` overrides of the global value, to allow cache keys to be smaller in the common case where different tests have different requirements.


Discussion:
Some thoughts on how to implement this:

## Step 1: define the new field

Start by reading the Target API docs https://www.pantsbuild.org/v2.5/docs/target-api-concepts.

Define a new `StringSequenceField` like `PythonTestsExtraEnvVars` in `backend/python/target_types.py` in the section for `PythonTests`. In the `help`, explain the syntax, which as the same as the below for `[test].extra_env_vars`. Also explain that this will be combined with `[test].extra_env_vars`, and that the field will override that option if there are any conflicts. 

https://github.com/pantsbuild/pants/blob/714518abba57f58354ddc7831d485a98fd9dd723/src/python/pants/core/goals/test.py#L328-L338

Register the field in `PythonTests.core_fields` and verify it shows up with `./pants help python_tests`.

## Step 2: wire up

Here is how we currently create the `env`:

https://github.com/pantsbuild/pants/blob/714518abba57f58354ddc7831d485a98fd9dd723/src/python/pants/backend/python/goals/pytest_runner.py#L224-L228

To hook up your new field, you will:

1. Add `extra_env_vars: PythonTestsExtraEnvVars` to the "FieldSet", which is how we extract out the relevant fields from the target

https://github.com/pantsbuild/pants/blob/714518abba57f58354ddc7831d485a98fd9dd723/src/python/pants/backend/python/goals/pytest_runner.py#L75-L81

2. Add `complete_env: CompleteEnvironment` to 

https://github.com/pantsbuild/pants/blob/714518abba57f58354ddc7831d485a98fd9dd723/src/python/pants/backend/python/goals/pytest_runner.py#L107-L117

3. In that line where we create the `env`, add `**complete_env.get_subset(request.field_set. or ())` to the end. Note that we need it after `**test_extra_env.env` so that the field overrides the option - that is probably worth a comment.

You can manually verify it works by doing something like adding `raise ValueError(os.environ)` to a `util/strutil_test.py` test and playing with the `python_tests` target's field.

## Step 3: expand tests

Update this test

https://github.com/pantsbuild/pants/blob/714518abba57f58354ddc7831d485a98fd9dd723/src/python/pants/backend/python/goals/pytest_runner_integration_test.py#L379-L401

We want to continue testing that those two env vars work, although feel free to rename them. We also want to test that both those styles of `VAR` and `VAR=val` work for the field. Finally, we want to confirm that the field overrides the option. So, there should be at least 5, maybe 6 env vars used in that test.

Reminder you can do `./pants test path/to/test.py -- -k test_extra_env_vars` to only run the test you care about.

Resolving patch:
diff --git a/src/python/pants/backend/python/goals/pytest_runner.py b/src/python/pants/backend/python/goals/pytest_runner.py
--- a/src/python/pants/backend/python/goals/pytest_runner.py
+++ b/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`.
diff --git a/src/python/pants/backend/python/target_types.py b/src/python/pants/backend/python/target_types.py
--- a/src/python/pants/backend/python/target_types.py
+++ b/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, "


Test patch:
diff --git a/src/python/pants/backend/python/goals/pytest_runner_integration_test.py b/src/python/pants/backend/python/goals/pytest_runner_integration_test.py
--- a/src/python/pants/backend/python/goals/pytest_runner_integration_test.py
+++ b/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