Skip to content

Commit

Permalink
Rust test targets now create a test launcher allow for setting env va…
Browse files Browse the repository at this point in the history
…rs (#579)

I was mistaken when I originally implemented #577 because I didn't realize that [std::env](https://doc.rust-lang.org/std/macro.env.html) loaded environment variables at compile time. The goal of this PR was to define environment variables at runtime. This appeared to be more complicated than I had thought. It seems there needs to be a launcher/runner that is produced by your rule and invokes your test executable to be able to guarantee the environment is configured for the test.

This PR creates a small shell script that wraps `rust_test` targets to enable the use of the `env` attribute at runtime.
  • Loading branch information
UebelAndre committed Feb 12, 2021
1 parent 6267c26 commit ab86e53
Show file tree
Hide file tree
Showing 8 changed files with 283 additions and 13 deletions.
117 changes: 114 additions & 3 deletions rust/private/rust.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
# buildifier: disable=module-docstring
load("//rust/private:common.bzl", "rust_common")
load("//rust/private:rustc.bzl", "rustc_compile_action")
load("//rust/private:utils.bzl", "determine_output_hash", "find_toolchain")
load("//rust/private:utils.bzl", "determine_output_hash", "expand_locations", "find_toolchain")

# TODO(marco): Separate each rule into its own file.

Expand Down Expand Up @@ -219,6 +219,99 @@ def _rust_binary_impl(ctx):
),
)

def _create_test_launcher(ctx, toolchain, output, providers):
"""Create a process wrapper to ensure runtime environment variables are defined for the test binary
Args:
ctx (ctx): The rule's context object
toolchain (rust_toolchain): The current rust toolchain
output (File): The output File that will be produced, depends on crate type.
providers (list): Providers from a rust compile action. See `rustc_compile_action`
Returns:
list: A list of providers similar to `rustc_compile_action` but with modified default info
"""

args = ctx.actions.args()

# TODO: It's unclear if the toolchain is in the same configuration as the `_launcher` attribute
# This should be investigated but for now, we generally assume if the target environment is windows,
# the execution environment is windows.
if toolchain.os == "windows":
launcher = ctx.actions.declare_file(name_to_crate_name(ctx.label.name + ".launcher.exe"))
# Because the windows target is a batch file, it expects native windows paths (with backslashes)
args.add_all([
ctx.executable._launcher.path.replace("/", "\\"),
launcher.path.replace("/", "\\"),
])
else:
launcher = ctx.actions.declare_file(name_to_crate_name(ctx.label.name + ".launcher"))
args.add_all([
ctx.executable._launcher,
launcher,
])

# Because returned executables must be created from the same rule, the
# launcher target is simply copied and exposed.
ctx.actions.run(
outputs = [launcher],
tools = [ctx.executable._launcher],
mnemonic = "GeneratingLauncher",
executable = ctx.executable._launcher_installer,
arguments = [args],
)

# Get data attribute
data = getattr(ctx.attr, "data", [])

# Expand the environment variables and write them to a file
environ_file = ctx.actions.declare_file(launcher.basename + ".launchfiles/env")
environ = expand_locations(
ctx,
getattr(ctx.attr, "env", {}),
data,
)

# Convert the environment variables into a list to be written into a file.
environ_list = []
for key, value in sorted(environ.items()):
environ_list.extend([key, value])

ctx.actions.write(
output = environ_file,
content = "\n".join(environ_list)
)

launcher_files = [environ_file]

# Replace the `DefaultInfo` provider in the returned list
default_info = None
for i in range(len(providers)):
if type(providers[i]) == "DefaultInfo":
default_info = providers[i]
providers.pop(i)
break

if not default_info:
fail("No DefaultInfo provider returned from `rustc_compile_action`")

providers.extend([
DefaultInfo(
files = default_info.files,
runfiles = default_info.default_runfiles.merge(
# The output is now also considered a runfile
ctx.runfiles(files = launcher_files + [output]),
),
executable = launcher,
),
OutputGroupInfo(
launcher_files = depset(launcher_files),
output = depset([output]),
),
])

return providers

def _rust_test_common(ctx, toolchain, output):
"""Builds a Rust test binary.
Expand Down Expand Up @@ -267,15 +360,16 @@ def _rust_test_common(ctx, toolchain, output):
is_test = True,
)

return rustc_compile_action(
providers = rustc_compile_action(
ctx = ctx,
toolchain = toolchain,
crate_type = crate_type,
crate_info = target,
rust_flags = ["--test"],
environ = ctx.attr.env,
)

return _create_test_launcher(ctx, toolchain, output, providers)

def _rust_test_impl(ctx):
"""The implementation of the `rust_test` rule
Expand Down Expand Up @@ -511,6 +605,23 @@ _rust_test_attrs = {
["Make variable"](https://docs.bazel.build/versions/master/be/make-variables.html) substitution.
"""),
),
"_launcher": attr.label(
executable = True,
default = Label("//util/launcher:launcher"),
cfg = "exec",
doc = _tidy("""
A launcher executable for loading environment and argument files passed in via the `env` attribute
and ensuring the variables are set for the underlying test executable.
"""),
),
"_launcher_installer": attr.label(
executable = True,
default = Label("//util/launcher:installer"),
cfg = "exec",
doc = _tidy("""
A helper script for creating an installer within the test rule.
"""),
),
}

rust_library = rule(
Expand Down
7 changes: 0 additions & 7 deletions rust/private/rustc.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -586,13 +586,6 @@ def rustc_compile_action(
build_flags_files,
)

# Make the user defined enviroment variables available to the action
expanded_env = dict()
data = getattr(ctx.attr, "data", [])
for key in environ:
expanded_env[key] = ctx.expand_location(environ[key], data)
env.update(expanded_env)

if hasattr(ctx.attr, "version") and ctx.attr.version != "0.0.0":
formatted_version = " v{}".format(ctx.attr.version)
else:
Expand Down
6 changes: 3 additions & 3 deletions test/test_env/tests/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ fn run() {
let output = std::process::Command::new(path).output().expect("Failed to run process");
assert_eq!(&b"Hello world\n"[..], output.stdout.as_slice());

// Test the `env` attribute of `rust_test`
assert_eq!(env!("FERRIS_SAYS"), "Hello fellow Rustaceans!");
assert_eq!(env!("HELLO_WORLD_BIN"), "test/test_env/hello-world");
// Test the `env` attribute of `rust_test` at run time
assert_eq!(std::env::var("FERRIS_SAYS").unwrap(), "Hello fellow Rustaceans!");
assert_eq!(std::env::var("HELLO_WORLD_BIN").unwrap(), "test/test_env/hello-world");
}
18 changes: 18 additions & 0 deletions util/launcher/BUILD.bazel
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
load("//rust:rust.bzl", "rust_binary")
load(":installer.bzl", "installer")

package(default_visibility = ["//visibility:public"])

rust_binary(
name = "launcher",
edition = "2018",
srcs = ["launcher_main.rs"]
)

installer(
name = "installer",
src = select({
"//rust/platform:windows": "launcher_installer.bat",
"//conditions:default": "launcher_installer.sh",
}),
)
36 changes: 36 additions & 0 deletions util/launcher/installer.bzl
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
"""A module defining the installer rule for the rules_rust test launcher"""

def _installer_impl(ctx):
"""The `installer` rule's implementation
Args:
ctx (ctx): The rule's context object
Returns:
list: A list a DefaultInfo provider
"""

installer = ctx.actions.declare_file(ctx.file.src.basename)

ctx.actions.expand_template(
template = ctx.file.src,
output = installer,
substitutions = {},
is_executable = True,
)

return [DefaultInfo(
files = depset([installer]),
executable = installer,
)]

installer = rule(
doc = "A rule which makes a native executable script available to other rules",
implementation = _installer_impl,
attrs = {
"src": attr.label(
allow_single_file = [".sh", ".bat"],
),
},
executable = True,
)
5 changes: 5 additions & 0 deletions util/launcher/launcher_installer.bat
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
@ECHO OFF
@REM A native windows script for creating a `run` action output of the
@REM the rules_rust launcher binary

copy /v /y /b "%1" "%2"
5 changes: 5 additions & 0 deletions util/launcher/launcher_installer.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
#!/bin/bash
# A simple script for creating a `run` action output of the
# the rules_rust launcher binary

cp "$1" "$2"
102 changes: 102 additions & 0 deletions util/launcher/launcher_main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
use std::collections::BTreeMap;
use std::fs::File;
use std::io::{BufReader, BufRead};
use std::path::PathBuf;
use std::process::Command;
use std::vec::Vec;

#[cfg(target_family = "unix")]
use std::os::unix::process::CommandExt;

/// This string must match the one found in `_create_test_launcher`
const LAUNCHFILES_ENV_PATH: &'static str = ".launchfiles/env";

/// Load environment variables from a uniquly formatted
fn environ() -> BTreeMap<String, String> {
let mut environ = BTreeMap::new();

let mut key: Option<String> = None;

// Load the environment file into a map
let env_path = std::env::args().nth(0).expect("arg 0 was not set") + LAUNCHFILES_ENV_PATH;
let file = File::open(env_path).expect("Failed to load the environment file");

// Find all environment variables by reading pairs of lines as key/value pairs
for line in BufReader::new(file).lines() {
if key.is_none() {
key = Some(line.expect("Failed to read line"));
continue;
}

environ.insert(
key.expect("Key is not set"),
line.expect("Failed to read line"),
);

key = None;
}

environ
}

/// Locate the executable based on the name of the launcher executable
fn executable() -> PathBuf {
let mut exec_path = std::env::args().nth(0).expect("arg 0 was not set");
let stem_index = exec_path.rfind(".launcher").expect("This executable should always contain `.launcher`");

// Remove the substring from the exec path
for _char in ".launcher".chars() {
exec_path.remove(stem_index);
}

PathBuf::from(exec_path)
}

/// Parse the command line arguments but skip the first element which
/// is the path to the test runner executable.
fn args() -> Vec<String> {
std::env::args().skip(1).collect()
}

/// Simply replace the current process with our test
#[cfg(target_family = "unix")]
fn exec(environ: BTreeMap<String, String>, executable: PathBuf, args: Vec<String>) {
let error = Command::new(&executable)
.envs(environ.iter())
.args(args)
.exec();

panic!("Process failed to start: {:?} with {:?}", executable, error)
}

/// On windows, there is no way to replace the current process
/// so instead we allow the command to run in a subprocess.
#[cfg(target_family = "windows")]
fn exec(environ: BTreeMap<String, String>, executable: PathBuf, args: Vec<String>) {
let output = Command::new(executable)
.envs(environ.iter())
.args(args)
.output()
.expect("Failed to run process");

std::process::exit(output.status.code().unwrap_or(1));
}

/// Main entrypoint
fn main() {
// Gather environment variables
let environ = environ();

// Gather arguments
let args = args();

// Find executable
let executable = executable();

// Replace the current process with the test target
exec(environ, executable, args);

// The call to exec should have exited the application.
// This code should be unreachable.
panic!("Process did not exit");
}

0 comments on commit ab86e53

Please sign in to comment.