Skip to content

Commit

Permalink
Let Starlark tests inherit env variables
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
fmeum committed Feb 17, 2022
1 parent b3baae9 commit bb08f5b
Show file tree
Hide file tree
Showing 9 changed files with 176 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down Expand Up @@ -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<>();
}

Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -442,7 +451,7 @@ private TestParams createTestAction(int shards)
coverageArtifact,
coverageDirectory,
testProperties,
runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
runfilesSupport.getActionEnvironment().addVariables(extraTestEnv, extraInheritedEnv),
executionSettings,
shard,
run,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand All @@ -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
Expand All @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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")));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();
}
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand Down Expand Up @@ -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;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
72 changes: 72 additions & 0 deletions src/test/shell/bazel/bazel_rules_test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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"

0 comments on commit bb08f5b

Please sign in to comment.