Auto merge of #13204 - hi-rustin:rustin-patch-test-out-dir, r=ehuss

fix: set OUT_DIR for all units with build scripts

crates/cargo-test-support/src/lib.rs

@@ -1286,7 +1286,8 @@ pub trait TestEnv: Sized {
             .env_remove("RUSTFLAGS")
             .env_remove("SSH_AUTH_SOCK") // ensure an outer agent is never contacted
             .env_remove("USER") // not set on some rust-lang docker images
-            .env_remove("XDG_CONFIG_HOME"); // see #2345
+            .env_remove("XDG_CONFIG_HOME") // see #2345
+            .env_remove("OUT_DIR"); // see #13204
         if cfg!(target_os = "macos") {
             // Work-around a bug in macOS 10.15, see `link_or_copy` for details.
             self = self.env("__CARGO_COPY_DONT_LINK_DO_NOT_USE_THIS", "1");

src/cargo/core/compiler/context/mod.rs

@@ -12,6 +12,7 @@ use crate::util::errors::CargoResult;
 use crate::util::profile;
 use anyhow::{bail, Context as _};
 use filetime::FileTime;
+use itertools::Itertools;
 use jobserver::Client;
 
 use super::build_plan::BuildPlan;
@@ -185,6 +186,32 @@ impl<'a, 'cfg> Context<'a, 'cfg> {
             plan.output_plan(self.bcx.config);
         }
 
+        // Add `OUT_DIR` to env vars if unit has a build script.
+        let units_with_build_script = &self
+            .bcx
+            .roots
+            .iter()
+            .filter(|unit| self.build_scripts.contains_key(unit))
+            .dedup_by(|x, y| x.pkg.package_id() == y.pkg.package_id())
+            .collect::<Vec<_>>();
+        for unit in units_with_build_script {
+            for dep in &self.bcx.unit_graph[unit] {
+                if dep.unit.mode.is_run_custom_build() {
+                    let out_dir = self
+                        .files()
+                        .build_script_out_dir(&dep.unit)
+                        .display()
+                        .to_string();
+                    let script_meta = self.get_run_build_script_metadata(&dep.unit);
+                    self.compilation
+                        .extra_env
+                        .entry(script_meta)
+                        .or_insert_with(Vec::new)
+                        .push(("OUT_DIR".to_string(), out_dir));
+                }
+            }
+        }
+
         // Collect the result of the build into `self.compilation`.
         for unit in &self.bcx.roots {
             // Collect tests and executables.
@@ -213,26 +240,6 @@ impl<'a, 'cfg> Context<'a, 'cfg> {
                 }
             }
 
-            // If the unit has a build script, add `OUT_DIR` to the
-            // environment variables.
-            if unit.target.is_lib() {
-                for dep in &self.bcx.unit_graph[unit] {
-                    if dep.unit.mode.is_run_custom_build() {
-                        let out_dir = self
-                            .files()
-                            .build_script_out_dir(&dep.unit)
-                            .display()
-                            .to_string();
-                        let script_meta = self.get_run_build_script_metadata(&dep.unit);
-                        self.compilation
-                            .extra_env
-                            .entry(script_meta)
-                            .or_insert_with(Vec::new)
-                            .push(("OUT_DIR".to_string(), out_dir));
-                    }
-                }
-            }
-
             // Collect information for `rustdoc --test`.
             if unit.mode.is_doc_test() {
                 let mut unstable_opts = false;

tests/testsuite/test.rs

@@ -4897,3 +4897,49 @@ fn cargo_test_print_env_verbose() {
         )
         .run();
 }
+
+#[cargo_test]
+fn cargo_test_set_out_dir_env_var() {
+    let p = project()
+        .file(
+            "Cargo.toml",
+            r#"
+                [package]
+                name = "foo"
+                version = "0.0.1"
+                edition = "2021"
+            "#,
+        )
+        .file(
+            "src/lib.rs",
+            r#"
+                pub fn add(left: usize, right: usize) -> usize {
+                    left + right
+                }
+            "#,
+        )
+        .file(
+            "build.rs",
+            r#"
+                fn main() {}
+            "#,
+        )
+        .file(
+            "tests/case.rs",
+            r#"
+                #[cfg(test)]
+                pub mod tests {
+                    #[test]
+                    fn test_add() {
+                        assert!(std::env::var("OUT_DIR").is_ok());
+                        assert_eq!(foo::add(2, 5), 7);
+                    }
+                }
+            "#,
+        )
+        .build();
+
+    p.cargo("test").run();
+    p.cargo("test --package foo --test case -- tests::test_add --exact --nocapture")
+        .run();
+}