diff --git a/src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py b/src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py index ee3d7f2ace4..1f51b6b0250 100644 --- a/src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py +++ b/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,)), diff --git a/src/python/pants/backend/go/goals/test.py b/src/python/pants/backend/go/goals/test.py index 93180241ff2..8e92212b48a 100644 --- a/src/python/pants/backend/go/goals/test.py +++ b/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, diff --git a/src/python/pants/backend/go/goals/test_test.py b/src/python/pants/backend/go/goals/test_test.py index 47c4fa00923..91bbee1294f 100644 --- a/src/python/pants/backend/go/goals/test_test.py +++ b/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( { diff --git a/src/python/pants/backend/go/target_types.py b/src/python/pants/backend/go/target_types.py index 07f4048e0dc..c1ae84c6790 100644 --- a/src/python/pants/backend/go/target_types.py +++ b/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, ) diff --git a/src/python/pants/backend/go/util_rules/embed_integration_test.py b/src/python/pants/backend/go/util_rules/embed_integration_test.py index 12370ef50b8..7872b7cd186 100644 --- a/src/python/pants/backend/go/util_rules/embed_integration_test.py +++ b/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],