[nextest-runner] set env vars from build script

chrjabs · sunshowers · commit eb1a16d7c479 · 2024-10-23T17:32:31.000-07:00
Set environment variables specified via `cargo::rustc-env` in the build
script when running tests
diff --git a/fixture-data/src/nextest_tests.rs b/fixture-data/src/nextest_tests.rs
@@ -166,6 +166,7 @@ pub static EXPECTED_TEST_SUITES: Lazy<BTreeMap<RustBinaryId, TestSuiteFixture>>
             "with-build-script",
             BuildPlatform::Target,
             vec![
+                TestCaseFixture::new("tests::test_build_script_vars_set", TestCaseFixtureStatus::Pass),
                 TestCaseFixture::new("tests::test_out_dir_present", TestCaseFixtureStatus::Pass),
             ],
         ),
diff --git a/fixtures/nextest-tests/with-build-script/build.rs b/fixtures/nextest-tests/with-build-script/build.rs
@@ -9,4 +9,28 @@ fn main() {
 
     // The presence of this file is checked by a test.
     std::fs::write(out_dir.join("this-is-a-test-file"), "test-contents").unwrap();
+
+    // Needed for 1.79+
+    println!("cargo:rustc-check-cfg=cfg(new_format)");
+
+    // The presence of these environment variables is checked by a test.
+    if rustc_minor_version().is_some_and(|minor| minor >= 77) {
+        println!("cargo::rustc-cfg=new_format");
+        println!("cargo::rustc-env=BUILD_SCRIPT_NEW_FMT=new_val");
+    }
+    println!("cargo:rustc-env=BUILD_SCRIPT_OLD_FMT=old_val");
+}
+
+fn rustc_minor_version() -> Option<u32> {
+    let rustc = std::env::var_os("RUSTC")?;
+    let output = std::process::Command::new(rustc)
+        .arg("--version")
+        .output()
+        .ok()?;
+    let version = std::str::from_utf8(&output.stdout).ok()?;
+    let mut pieces = version.split('.');
+    if pieces.next() != Some("rustc 1") {
+        return None;
+    }
+    pieces.next()?.parse().ok()
 }
diff --git a/fixtures/nextest-tests/with-build-script/src/lib.rs b/fixtures/nextest-tests/with-build-script/src/lib.rs
@@ -12,4 +12,17 @@ mod tests {
         let contents = std::fs::read(&path).expect("test file exists in OUT_DIR");
         assert_eq!(contents, b"test-contents");
     }
+
+    #[test]
+    fn test_build_script_vars_set() {
+        // Since the build script wrote `cargo::rustc-env` instructions, these variables are
+        // expected to be set by nextest
+        #[cfg(new_format)]
+        {
+            let val = std::env::var("BUILD_SCRIPT_NEW_FMT").expect("BUILD_SCRIPT_NEW_FMT is valid");
+            assert_eq!(val, "new_val");
+        }
+        let val = std::env::var("BUILD_SCRIPT_OLD_FMT").expect("BUILD_SCRIPT_OLD_FMT is valid");
+        assert_eq!(val, "old_val");
+    }
 }
diff --git a/integration-tests/tests/integration/snapshots/integration__list_with_default_set_basic.snap b/integration-tests/tests/integration/snapshots/integration__list_with_default_set_basic.snap
@@ -37,4 +37,5 @@ nextest-tests::example/nextest-tests:
 nextest-tests::example/other:
     tests::other_example_success
 with-build-script:
+    tests::test_build_script_vars_set
     tests::test_out_dir_present
diff --git a/integration-tests/tests/integration/snapshots/integration__list_with_default_set_bound_all.snap b/integration-tests/tests/integration/snapshots/integration__list_with_default_set_bound_all.snap
@@ -41,4 +41,5 @@ nextest-tests::example/nextest-tests:
 nextest-tests::example/other:
     tests::other_example_success
 with-build-script:
+    tests::test_build_script_vars_set
     tests::test_out_dir_present
diff --git a/integration-tests/tests/integration/snapshots/integration__list_with_default_set_expr_all.snap b/integration-tests/tests/integration/snapshots/integration__list_with_default_set_expr_all.snap
@@ -37,4 +37,5 @@ nextest-tests::example/nextest-tests:
 nextest-tests::example/other:
     tests::other_example_success
 with-build-script:
+    tests::test_build_script_vars_set
     tests::test_out_dir_present
diff --git a/integration-tests/tests/integration/snapshots/integration__list_with_default_set_expr_default.snap b/integration-tests/tests/integration/snapshots/integration__list_with_default_set_expr_default.snap
@@ -37,4 +37,5 @@ nextest-tests::example/nextest-tests:
 nextest-tests::example/other:
     tests::other_example_success
 with-build-script:
+    tests::test_build_script_vars_set
     tests::test_out_dir_present
diff --git a/integration-tests/tests/integration/snapshots/integration__show_config_test_groups-2.snap b/integration-tests/tests/integration/snapshots/integration__show_config_test_groups-2.snap
@@ -47,5 +47,6 @@ group: @global
       nextest-tests::example/other:
           tests::other_example_success
       with-build-script:
+          tests::test_build_script_vars_set
           tests::test_out_dir_present
 
diff --git a/integration-tests/tests/integration/snapshots/integration__show_config_test_groups-4.snap b/integration-tests/tests/integration/snapshots/integration__show_config_test_groups-4.snap
@@ -48,5 +48,6 @@ group: @global
       nextest-tests::example/other:
           tests::other_example_success
       with-build-script:
+          tests::test_build_script_vars_set
           tests::test_out_dir_present
 
diff --git a/integration-tests/tests/integration/snapshots/integration__show_config_test_groups-6.snap b/integration-tests/tests/integration/snapshots/integration__show_config_test_groups-6.snap
@@ -49,5 +49,6 @@ group: @global
       nextest-tests::example/other:
           tests::other_example_success
       with-build-script:
+          tests::test_build_script_vars_set
           tests::test_out_dir_present
 
diff --git a/nextest-runner/src/reuse_build/archiver.rs b/nextest-runner/src/reuse_build/archiver.rs
@@ -23,7 +23,7 @@ use std::{
     io::{self, BufWriter, Write},
     time::{Instant, SystemTime},
 };
-use tracing::{debug, trace};
+use tracing::{debug, trace, warn};
 use zstd::Encoder;
 
 /// Archive format.
@@ -394,6 +394,23 @@ impl<'a, W: Write> Archiver<'a, W> {
                 false,
                 callback,
             )?;
+
+            // Archive build script output in order to set environment variables from there
+            let Some(out_dir_parent) = build_script_out_dir.parent() else {
+                warn!("could not determine parent directory of output directory {build_script_out_dir}");
+                continue;
+            };
+            let out_file_path = out_dir_parent.join("output");
+            let src_path = self
+                .binary_list
+                .rust_build_meta
+                .target_directory
+                .join(&out_file_path);
+
+            let rel_path = Utf8Path::new("target").join(out_file_path);
+            let rel_path = convert_rel_path_to_forward_slash(&rel_path);
+
+            self.append_file(ArchiveStep::BuildScriptOutDirs, &src_path, &rel_path)?;
         }
 
         // Write linked paths to the archive.
diff --git a/nextest-runner/src/test_command.rs b/nextest-runner/src/test_command.rs
@@ -14,7 +14,10 @@ use once_cell::sync::Lazy;
 use std::{
     collections::{BTreeSet, HashMap},
     ffi::{OsStr, OsString},
+    fs::File,
+    io::{BufRead, BufReader},
 };
+use tracing::warn;
 
 mod imp;
 pub use imp::{Child, Output};
@@ -74,7 +77,8 @@ impl TestCommand {
         {
             // Convert the output directory to an absolute path.
             let out_dir = lctx.rust_build_meta.target_directory.join(out_dir);
-            cmd.env("OUT_DIR", out_dir);
+            cmd.env("OUT_DIR", &out_dir);
+            apply_build_script_env(&mut cmd, out_dir);
         }
 
         // Expose paths to non-test binaries at runtime so that relocated paths work.
@@ -191,6 +195,55 @@ fn apply_package_env(cmd: &mut std::process::Command, package: &PackageMetadata<
     );
 }
 
+/// Applies environment variables spcified by the build script via `cargo::rustc-env`
+fn apply_build_script_env(cmd: &mut std::process::Command, out_dir: Utf8PathBuf) {
+    let Some(out_dir_parent) = out_dir.parent() else {
+        warn!("could not determine parent directory of output directory {out_dir}");
+        return;
+    };
+    let Ok(out_file) = File::open(out_dir_parent.join("output")) else {
+        warn!("could not find build script output file at {out_dir_parent}/output");
+        return;
+    };
+    parse_build_script_output(
+        BufReader::new(out_file),
+        &out_dir_parent.join("output"),
+        |key, val| {
+            cmd.env(key, val);
+        },
+    );
+}
+
+/// Parses the build script output and calls the callback for each key value pair of `VAR=val`
+///
+/// This is mainly split out into a separate function from [`apply_build_script_env`] for easier
+/// unit testing
+fn parse_build_script_output<R, Cb>(out_file: R, out_file_path: &Utf8Path, mut callback: Cb)
+where
+    R: BufRead,
+    Cb: FnMut(&str, &str),
+{
+    for line in out_file.lines() {
+        let Ok(line) = line else {
+            warn!("found line with invalid UTF8 in build script output file at {out_file_path}");
+            continue;
+        };
+        // `cargo::rustc-env` is the official syntax since `cargo` 1.77, `cargo:rustc-env` is
+        // supported for backwards compatibility
+        let Some(key_val) = line
+            .strip_prefix("cargo::rustc-env=")
+            .or_else(|| line.strip_prefix("cargo:rustc-env="))
+        else {
+            continue;
+        };
+        let Some(split) = key_val.find('=') else {
+            warn!("rustc-env variable '{key_val}' has no value in {out_file_path}, skipping");
+            continue;
+        };
+        callback(&key_val[..split], &key_val[split + 1..]);
+    }
+}
+
 /// This is a workaround for a macOS SIP issue:
 /// https://github.com/nextest-rs/nextest/pull/84
 ///
@@ -236,3 +289,24 @@ pub(crate) fn apply_ld_dyld_env(cmd: &mut std::process::Command, dylib_path: &Os
         cmd.env("NEXTEST_".to_owned() + dylib_path_envvar(), dylib_path);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn parse_build_script_output() {
+        let out_file = "some_other_line\ncargo::rustc-env=NEW_VAR=new_val\ncargo:rustc-env=OLD_VAR=old_val\ncargo::rustc-env=NEW_MISSING_VALUE\ncargo:rustc-env=OLD_MISSING_VALUE";
+
+        let mut key_vals = Vec::new();
+        super::parse_build_script_output(
+            std::io::BufReader::new(std::io::Cursor::new(out_file)),
+            camino::Utf8Path::new("<test input>"),
+            |key, val| key_vals.push((key.to_owned(), val.to_owned())),
+        );
+
+        assert_eq!(key_vals.len(), 2);
+        assert_eq!(key_vals[0].0, "NEW_VAR");
+        assert_eq!(key_vals[0].1, "new_val");
+        assert_eq!(key_vals[1].0, "OLD_VAR");
+        assert_eq!(key_vals[1].1, "old_val");
+    }
+}
