Add support for extra env variables in go tests (Cherry-pick of #16013)

src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py

@@ -35,7 +35,7 @@
     third_party_pkg,
 )
 from pants.build_graph.address import Address
-from pants.core.goals.test import TestResult
+from pants.core.goals.test import TestResult, get_filtered_environment
 from pants.core.util_rules import config_files, source_files, stripped_source_files
 from pants.core.util_rules.external_tool import rules as external_tool_rules
 from pants.engine.fs import Digest, DigestContents
@@ -70,6 +70,7 @@ def rule_runner() -> RuleRunner:
             *target_type_rules.rules(),
             *tests_analysis.rules(),
             *third_party_pkg.rules(),
+            get_filtered_environment,
             QueryRule(HydratedSources, [HydrateSourcesRequest]),
             QueryRule(GeneratedSources, [GenerateGoFromProtobufRequest]),
             QueryRule(DigestContents, (Digest,)),

src/python/pants/backend/go/goals/test.py

@@ -10,6 +10,7 @@
 from pants.backend.go.subsystems.gotest import GoTestSubsystem
 from pants.backend.go.target_types import (
     GoPackageSourcesField,
+    GoTestExtraEnvVarsField,
     GoTestTimeoutField,
     SkipGoTestsField,
 )
@@ -31,12 +32,14 @@
 from pants.core.goals.test import (
     TestDebugAdapterRequest,
     TestDebugRequest,
+    TestExtraEnv,
     TestFieldSet,
     TestResult,
     TestSubsystem,
 )
 from pants.core.target_types import FileSourceField
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
+from pants.engine.environment import Environment, EnvironmentRequest
 from pants.engine.fs import EMPTY_FILE_DIGEST, AddPrefix, Digest, MergeDigests
 from pants.engine.process import FallibleProcessResult, Process, ProcessCacheScope
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
@@ -85,6 +88,7 @@ class GoTestFieldSet(TestFieldSet):
     sources: GoPackageSourcesField
     dependencies: Dependencies
     timeout: GoTestTimeoutField
+    extra_env_vars: GoTestExtraEnvVarsField
 
     @classmethod
     def opt_out(cls, tgt: Target) -> bool:
@@ -148,7 +152,10 @@ def transform_test_args(args: Sequence[str], timeout_field_value: int | None) ->
 
 @rule(desc="Test with Go", level=LogLevel.DEBUG)
 async def run_go_tests(
-    field_set: GoTestFieldSet, test_subsystem: TestSubsystem, go_test_subsystem: GoTestSubsystem
+    field_set: GoTestFieldSet,
+    test_subsystem: TestSubsystem,
+    go_test_subsystem: GoTestSubsystem,
+    test_extra_env: TestExtraEnv,
 ) -> TestResult:
     maybe_pkg_analysis, maybe_pkg_digest, dependencies = await MultiGet(
         Get(FallibleFirstPartyPkgAnalysis, FirstPartyPkgAnalysisRequest(field_set.address)),
@@ -289,7 +296,10 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
     # This allows tests to open dependencies on `file` targets regardless of where they are
     # located. See https://dave.cheney.net/2016/05/10/test-fixtures-in-go.
     working_dir = field_set.address.spec_path
-    binary_with_prefix, files_sources = await MultiGet(
+    field_set_extra_env_get = Get(
+        Environment, EnvironmentRequest(field_set.extra_env_vars.value or ())
+    )
+    binary_with_prefix, files_sources, field_set_extra_env = await MultiGet(
         Get(Digest, AddPrefix(binary.digest, working_dir)),
         Get(
             SourceFiles,
@@ -299,11 +309,19 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
                 enable_codegen=True,
             ),
         ),
+        field_set_extra_env_get,
     )
     test_input_digest = await Get(
         Digest, MergeDigests((binary_with_prefix, files_sources.snapshot.digest))
     )
 
+    extra_env = {
+        **test_extra_env.env,
+        # NOTE: field_set_extra_env intentionally after `test_extra_env` to allow overriding within
+        # `go_package`.
+        **field_set_extra_env,
+    }
+
     cache_scope = (
         ProcessCacheScope.PER_SESSION if test_subsystem.force else ProcessCacheScope.SUCCESSFUL
     )
@@ -315,6 +333,7 @@ def compilation_failure(exit_code: int, stdout: str | None, stderr: str | None)
                 "./test_runner",
                 *transform_test_args(go_test_subsystem.args, field_set.timeout.value),
             ],
+            env=extra_env,
             input_digest=test_input_digest,
             description=f"Run Go tests: {field_set.address}",
             cache_scope=cache_scope,

src/python/pants/backend/go/goals/test_test.py

@@ -24,7 +24,7 @@
     third_party_pkg,
 )
 from pants.backend.go.util_rules.sdk import GoSdkProcess
-from pants.core.goals.test import TestResult
+from pants.core.goals.test import TestResult, get_filtered_environment
 from pants.core.target_types import FileTarget
 from pants.core.util_rules import source_files
 from pants.engine.addresses import Address
@@ -48,6 +48,7 @@ def rule_runner() -> RuleRunner:
             *tests_analysis.rules(),
             *third_party_pkg.rules(),
             *source_files.rules(),
+            get_filtered_environment,
             QueryRule(TestResult, [GoTestFieldSet]),
             QueryRule(ProcessResult, [GoSdkProcess]),
         ],
@@ -543,6 +544,74 @@ def test_fuzz_target_supported(rule_runner: RuleRunner) -> None:
     assert "PASS: FuzzFoo" in result.stdout
 
 
+def test_extra_env_vars(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "foo/BUILD": textwrap.dedent(
+                """
+                go_mod(name='mod')
+                go_package(
+                    test_extra_env_vars=(
+                        "GO_PACKAGE_VAR_WITHOUT_VALUE",
+                        "GO_PACKAGE_VAR_WITH_VALUE=go_package_var_with_value",
+                        "GO_PACKAGE_OVERRIDE_WITH_VALUE_VAR=go_package_override_with_value_var_override",
+                    )
+                )
+                """
+            ),
+            "foo/go.mod": "module foo",
+            "foo/add.go": textwrap.dedent(
+                """
+                package foo
+                import "os"
+                func envIs(e, v string) bool {
+                  return (os.Getenv(e) == v)
+                }
+                """
+            ),
+            "foo/add_test.go": textwrap.dedent(
+                """
+                package foo
+                import "testing"
+                func TestEnvs(t *testing.T) {
+                  if !envIs("ARG_WITH_VALUE_VAR", "arg_with_value_var") {
+                      t.Fail()
+                  }
+                  if !envIs("ARG_WITHOUT_VALUE_VAR", "arg_without_value_var") {
+                      t.Fail()
+                  }
+                  if !envIs("GO_PACKAGE_VAR_WITH_VALUE", "go_package_var_with_value") {
+                      t.Fail()
+                  }
+                  if !envIs("GO_PACKAGE_VAR_WITHOUT_VALUE", "go_package_var_without_value") {
+                      t.Fail()
+                  }
+                  if !envIs("GO_PACKAGE_OVERRIDE_WITH_VALUE_VAR", "go_package_override_with_value_var_override") {
+                      t.Fail()
+                  }
+                }
+                """
+            ),
+        }
+    )
+    tgt = rule_runner.get_target(Address("foo"))
+    rule_runner.set_options(
+        args=[
+            "--go-test-args=-v -bench=.",
+            '--test-extra-env-vars=["ARG_WITH_VALUE_VAR=arg_with_value_var", "ARG_WITHOUT_VALUE_VAR", "GO_PACKAGE_OVERRIDE_ARG_WITH_VALUE_VAR"]',
+        ],
+        env={
+            "ARG_WITHOUT_VALUE_VAR": "arg_without_value_var",
+            "GO_PACKAGE_VAR_WITHOUT_VALUE": "go_package_var_without_value",
+            "GO_PACKAGE_OVERRIDE_WITH_VALUE_VAR": "go_package_override_with_value_var",
+        },
+        env_inherit={"PATH"},
+    )
+    result = rule_runner.request(TestResult, [GoTestFieldSet.create(tgt)])
+    assert result.exit_code == 0
+    assert "PASS: TestEnvs" in result.stdout
+
+
 def test_skip_tests(rule_runner: RuleRunner) -> None:
     rule_runner.write_files(
         {

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

@@ -21,6 +21,7 @@
     InvalidTargetException,
     MultipleSourcesField,
     StringField,
+    StringSequenceField,
     Target,
     TargetGenerator,
     ValidNumbers,
@@ -179,6 +180,18 @@ class SkipGoTestsField(BoolField):
     help = "If true, don't run this package's tests."
 
 
+class GoTestExtraEnvVarsField(StringSequenceField):
+    alias = "test_extra_env_vars"
+    help = softwrap(
+        """
+         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 GoTestTimeoutField(IntField):
     alias = "test_timeout"
     help = softwrap(
@@ -197,6 +210,7 @@ class GoPackageTarget(Target):
         *COMMON_TARGET_FIELDS,
         GoPackageDependenciesField,
         GoPackageSourcesField,
+        GoTestExtraEnvVarsField,
         GoTestTimeoutField,
         SkipGoTestsField,
     )

src/python/pants/backend/go/util_rules/embed_integration_test.py

@@ -27,7 +27,7 @@
 )
 from pants.backend.go.util_rules.embedcfg import EmbedConfig
 from pants.build_graph.address import Address
-from pants.core.goals.test import TestResult
+from pants.core.goals.test import TestResult, get_filtered_environment
 from pants.core.target_types import ResourceTarget
 from pants.core.util_rules import source_files
 from pants.testutil.rule_runner import QueryRule, RuleRunner
@@ -49,6 +49,7 @@ def rule_runner() -> RuleRunner:
             *tests_analysis.rules(),
             *third_party_pkg.rules(),
             *source_files.rules(),
+            get_filtered_environment,
             QueryRule(TestResult, [GoTestFieldSet]),
         ],
         target_types=[GoModTarget, GoPackageTarget, ResourceTarget],