Let Starlark tests inherit env variables

Adds an inherited_environment field to testing.TestEnvironment to allow Starlark rules to implement the equivalent of the native rules' `env_inherit` attribute. RELNOTES: Starlark test rules can use the new inherited_environment parameter of testing.TestEnvironment to specify environment variables whose values should be inherited from the shell environment.
bazelbuild · Feb 17, 2022 · bb08f5b · bb08f5b
1 parent b3baae9
commit bb08f5b
Show file tree

Hide file tree

Showing 9 changed files with 176 additions and 9 deletions.
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java b/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java
@@ -186,6 +186,23 @@ public ActionEnvironment addFixedVariables(Map<String, String> vars) {
     return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), inheritedEnv);
   }
 
+  /**
+   * Returns a copy of the environment with the given fixed and inherited variables added to it,
+   * <em>overwriting any existing occurrences of those variables</em>.
+   */
+  public ActionEnvironment addVariables(Map<String, String> vars, Set<String> inheritedVars) {
+    if (inheritedVars.isEmpty()) {
+      return addFixedVariables(vars);
+    } else {
+      // TODO: inheritedEnv is currently not optimized for allocation avoidance in the same way as
+      //  fixedEnv.
+      ImmutableSet<String> newInheritedEnv = ImmutableSet.<String>builder().addAll(inheritedEnv)
+          .addAll(inheritedVars).build();
+      return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv),
+          newInheritedEnv);
+    }
+  }
+
   /** Returns the combined size of the fixed and inherited environments. */
   public int size() {
     return fixedEnv.size() + inheritedEnv.size();

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java b/src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java
@@ -476,6 +476,7 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
             providersBuilder.getProvider(TestEnvironmentInfo.PROVIDER.getKey());
     if (environmentProvider != null) {
       testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
+      testActionBuilder.addExtraInheritedEnv(environmentProvider.getInheritedEnvironment());
     }
 
     TestParams testParams =

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestActionBuilder.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestActionBuilder.java
@@ -55,7 +55,9 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.TreeMap;
+import java.util.TreeSet;
 import javax.annotation.Nullable;
 
 /**
@@ -96,10 +98,12 @@ public NestedSet<PackageGroupContents> getPackageSpecifications() {
   private InstrumentedFilesInfo instrumentedFiles;
   private int explicitShardCount;
   private final Map<String, String> extraEnv;
+  private final Set<String> extraInheritedEnv;
 
   public TestActionBuilder(RuleContext ruleContext) {
     this.ruleContext = ruleContext;
     this.extraEnv = new TreeMap<>();
+    this.extraInheritedEnv = new TreeSet<>();
     this.additionalTools = new ImmutableList.Builder<>();
   }
 
@@ -154,6 +158,11 @@ public TestActionBuilder addExtraEnv(Map<String, String> extraEnv) {
     return this;
   }
 
+  public TestActionBuilder addExtraInheritedEnv(List<String> extraInheritedEnv) {
+    this.extraInheritedEnv.addAll(extraInheritedEnv);
+    return this;
+  }
+
   /**
    * Set the explicit shard count. Note that this may be overridden by the sharding strategy.
    */
@@ -442,7 +451,7 @@ private TestParams createTestAction(int shards)
                 coverageArtifact,
                 coverageDirectory,
                 testProperties,
-                runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
+                runfilesSupport.getActionEnvironment().addVariables(extraTestEnv, extraInheritedEnv),
                 executionSettings,
                 shard,
                 run,

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestEnvironmentInfo.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestEnvironmentInfo.java
@@ -19,6 +19,7 @@
 import com.google.devtools.build.lib.packages.BuiltinProvider;
 import com.google.devtools.build.lib.packages.NativeInfo;
 import com.google.devtools.build.lib.starlarkbuildapi.test.TestEnvironmentInfoApi;
+import java.util.List;
 import java.util.Map;
 
 /** Provider containing any additional environment variables for use in the test action. */
@@ -30,10 +31,12 @@ public final class TestEnvironmentInfo extends NativeInfo implements TestEnviron
       new BuiltinProvider<TestEnvironmentInfo>("TestEnvironment", TestEnvironmentInfo.class) {};
 
   private final Map<String, String> environment;
+  private final List<String> inheritedEnvironment;
 
-  /** Constructs a new provider with the given variable name to variable value mapping. */
-  public TestEnvironmentInfo(Map<String, String> environment) {
+  /** Constructs a new provider with the given fixed and inherited environment variables. */
+  public TestEnvironmentInfo(Map<String, String> environment, List<String> inheritedEnvironment) {
     this.environment = Preconditions.checkNotNull(environment);
+    this.inheritedEnvironment = Preconditions.checkNotNull(inheritedEnvironment);
   }
 
   @Override
@@ -48,4 +51,12 @@ public BuiltinProvider<TestEnvironmentInfo> getProvider() {
   public Map<String, String> getEnvironment() {
     return environment;
   }
+
+  /**
+   * Returns names of environment variables whose value should be obtained from the environment.
+   */
+  @Override
+  public List<String> getInheritedEnvironment() {
+    return inheritedEnvironment;
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java b/src/main/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModule.java
@@ -18,6 +18,8 @@
 import com.google.devtools.build.lib.starlarkbuildapi.test.TestingModuleApi;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
+import net.starlark.java.eval.StarlarkList;
 
 /** A class that exposes testing infrastructure to Starlark. */
 public class StarlarkTestingModule implements TestingModuleApi {
@@ -29,9 +31,12 @@ public ExecutionInfo executionInfo(Dict<?, ?> requirements /* <String, String> *
   }
 
   @Override
-  public TestEnvironmentInfo testEnvironment(Dict<?, ?> environment /* <String, String> */)
+  public TestEnvironmentInfo testEnvironment(Dict<?, ?> environment /* <String, String> */,
+      Sequence<?> inheritedEnvironment /* <String> */)
       throws EvalException {
     return new TestEnvironmentInfo(
-        Dict.cast(environment, String.class, String.class, "environment"));
+        Dict.cast(environment, String.class, String.class, "environment"),
+        StarlarkList.immutableCopyOf(
+            Sequence.cast(inheritedEnvironment, String.class, "inherited_environment")));
   }
 }
diff --git a/...main/java/com/google/devtools/build/lib/starlarkbuildapi/test/TestEnvironmentInfoApi.java b/...main/java/com/google/devtools/build/lib/starlarkbuildapi/test/TestEnvironmentInfoApi.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.starlarkbuildapi.test;
 
 import com.google.devtools.build.lib.starlarkbuildapi.core.StructApi;
+import java.util.List;
 import java.util.Map;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
@@ -28,4 +29,10 @@ public interface TestEnvironmentInfoApi extends StructApi {
       doc = "A dict containing environment variables which should be set on the test action.",
       structField = true)
   Map<String, String> getEnvironment();
+
+  @StarlarkMethod(
+      name = "inherited_environment",
+      doc = "A list of variables that the test action should inherit from the environment.",
+      structField = true)
+  List<String> getInheritedEnvironment();
 }
diff --git a/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/test/TestingModuleApi.java b/src/main/java/com/google/devtools/build/lib/starlarkbuildapi/test/TestingModuleApi.java
@@ -15,10 +15,12 @@
 package com.google.devtools.build.lib.starlarkbuildapi.test;
 
 import net.starlark.java.annot.Param;
+import net.starlark.java.annot.ParamType;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.Dict;
 import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.StarlarkValue;
 
 /** Helper module for accessing test infrastructure. */
@@ -56,12 +58,22 @@ ExecutionInfoApi executionInfo(Dict<?, ?> requirements // <String, String> expec
       parameters = {
         @Param(
             name = "environment",
-            named = false,
+            named = true,
             positional = true,
             doc =
                 "A map of string keys and values that represent environment variables and their"
-                    + " values. These will be made available during the test execution.")
+                    + " values. These will be made available during the test execution."),
+          @Param(
+              name = "inherited_environment",
+              allowedTypes = {@ParamType(type = Sequence.class, generic1 = String.class)},
+              defaultValue = "[]",
+              named = true,
+              positional = true,
+              doc =
+                  "A sequence of string keys that represent environment variables. These variables"
+                      + "will be made available during the test execution with their current value"
+                      + " taken from the environment.")
       })
-  TestEnvironmentInfoApi testEnvironment(Dict<?, ?> environment // <String, String> expected
-      ) throws EvalException;
+  TestEnvironmentInfoApi testEnvironment(Dict<?, ?> environment, // <String, String> expected
+      Sequence<?> inheritedEnvironment /* <String> expected */) throws EvalException;
 }
diff --git a/src/test/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModuleTest.java b/src/test/java/com/google/devtools/build/lib/rules/test/StarlarkTestingModuleTest.java
@@ -79,6 +79,39 @@ public void testStarlarkRulePropagatesTestEnvironmentProvider() throws Exception
     assertThat(provider.getEnvironment().get("XCODE_VERSION_OVERRIDE")).isEqualTo("7.3.1");
   }
 
+  @Test
+  public void testStarlarkRulePropagatesTestEnvironmentProviderWithInheritedEnv() throws Exception {
+    scratch.file("examples/rule/BUILD");
+    scratch.file(
+        "examples/rule/apple_rules.bzl",
+        "def my_rule_impl(ctx):",
+        "  test_env = testing.TestEnvironment(",
+        "    {'XCODE_VERSION_OVERRIDE': '7.3.1'},",
+        "    [",
+        "      'DEVELOPER_DIR',",
+        "      'XCODE_VERSION_OVERRIDE',",
+        "    ]",
+        ")",
+        "  return [test_env]",
+        "my_rule = rule(implementation = my_rule_impl,",
+        "  attrs = {},",
+        ")");
+    scratch.file(
+        "examples/apple_starlark/BUILD",
+        "package(default_visibility = ['//visibility:public'])",
+        "load('//examples/rule:apple_rules.bzl', 'my_rule')",
+        "my_rule(",
+        "    name = 'my_target',",
+        ")");
+
+    ConfiguredTarget starlarkTarget = getConfiguredTarget("//examples/apple_starlark:my_target");
+    TestEnvironmentInfo provider = starlarkTarget.get(TestEnvironmentInfo.PROVIDER);
+
+    assertThat(provider.getEnvironment().get("XCODE_VERSION_OVERRIDE")).isEqualTo("7.3.1");
+    assertThat(provider.getInheritedEnvironment()).contains("DEVELOPER_DIR");
+    assertThat(provider.getInheritedEnvironment()).contains("XCODE_VERSION_OVERRIDE");
+  }
+
   @Test
   public void testExecutionInfoProviderCanMarkTestAsLocal() throws Exception {
     scratch.file("examples/rule/BUILD");

diff --git a/src/test/shell/bazel/bazel_rules_test.sh b/src/test/shell/bazel/bazel_rules_test.sh
@@ -593,4 +593,76 @@ EOF
   assert_contains "hello world" bazel-bin/pkg/hello_gen.txt
 }
 
+function test_starlark_test_with_test_environment() {
+  mkdir pkg
+  cat >pkg/BUILD <<'EOF'
+load(":rules.bzl", "my_test")
+my_test(
+  name = "my_test",
+)
+EOF
+
+  # On Windows this file needs to be acceptable by CreateProcessW(), rather
+  # than a Bourne script.
+  if "$is_windows"; then
+    cat >pkg/rules.bzl <<'EOF'
+_SCRIPT_EXT = ".bat"
+_SCRIPT_CONTENT = """@ECHO OFF
+if "%FIXED_ONLY%" != "fixed" {
+  exit /B 1
+}
+if "%FIXED_AND_INHERITED%" != "fixed" {
+  exit /B 1
+}
+if "%INHERITED_ONLY%" != "inherited" {
+  exit /B 1
+}
+"""
+EOF
+  else
+    cat >pkg/rules.bzl <<'EOF'
+_SCRIPT_EXT = ".sh"
+_SCRIPT_CONTENT = """#!/bin/bash
+env
+[[ "$FIXED_ONLY" == "fixed" && "$FIXED_AND_INHERITED" = "fixed" && "$INHERITED_ONLY" == "inherited" ]]
+"""
+EOF
+  fi
+
+  cat >>pkg/rules.bzl <<'EOF'
+def _my_test_impl(ctx):
+    test_sh = ctx.actions.declare_file(ctx.attr.name + _SCRIPT_EXT)
+    ctx.actions.write(
+        output = test_sh,
+        content = _SCRIPT_CONTENT,
+        is_executable = True,
+    )
+    test_env = testing.TestEnvironment(
+      {
+        "FIXED_AND_INHERITED": "fixed",
+        "FIXED_ONLY": "fixed",
+      },
+      [
+        "FIXED_AND_INHERITED",
+        "INHERITED_ONLY",
+      ]
+    )
+    return [
+        DefaultInfo(
+            executable = test_sh,
+        ),
+        test_env
+    ]
+
+my_test = rule(
+    implementation = _my_test_impl,
+    attrs = {},
+    test = True,
+)
+EOF
+
+  FIXED_AND_INHERITED=inherited INHERITED_ONLY=inherited \
+    bazel test //pkg:my_test --test_output=streamed || fail "Test should pass"
+}
+
 run_suite "rules test"