diff --git a/antlir/antlir2/antlir2_vm/bzl/test.bzl b/antlir/antlir2/antlir2_vm/bzl/test.bzl index 920bdb967cd..d221a552c5f 100644 --- a/antlir/antlir2/antlir2_vm/bzl/test.bzl +++ b/antlir/antlir2/antlir2_vm/bzl/test.bzl @@ -22,10 +22,7 @@ def _impl(ctx: AnalysisContext) -> list[Provider]: cmd_args(ctx.attrs.vm_host[VMHostInfo].image[LayerInfo].subvol_symlink, format = "--image={}"), cmd_args(ctx.attrs.vm_host[VMHostInfo].machine_spec, format = "--machine-spec={}"), cmd_args(ctx.attrs.vm_host[VMHostInfo].runtime_spec, format = "--runtime-spec={}"), - cmd_args( - ["{}={}".format(k, v) for k, v in ctx.attrs.test[ExternalRunnerTestInfo].env.items()], - format = "--setenv={}", - ), + cmd_args([k for k in ctx.attrs.test[ExternalRunnerTestInfo].env], format = "--passenv={}"), ) test_args = cmd_args( ctx.attrs.test[ExternalRunnerTestInfo].test_type, diff --git a/antlir/antlir2/antlir2_vm/src/isolation.rs b/antlir/antlir2/antlir2_vm/src/isolation.rs index 3017dcb6fd2..97d51532aaf 100644 --- a/antlir/antlir2/antlir2_vm/src/isolation.rs +++ b/antlir/antlir2/antlir2_vm/src/isolation.rs @@ -8,7 +8,6 @@ use std::collections::BTreeMap; use std::collections::HashSet; use std::env; -use std::ffi::OsString; use std::path::PathBuf; use std::process::Command; @@ -87,21 +86,6 @@ impl Platform { } } -/// If these env exist, always pass them through. -const PASSTHROUGH_ENVS: &[&str] = &["RUST_LOG", "ANTLIR_BUCK"]; - -/// Generate default passthrough env vars -pub(crate) fn default_passthrough_envs() -> Vec { - PASSTHROUGH_ENVS - .iter() - .filter(|x| env::var(*x).is_ok()) - .map(|x| KvPair { - key: x.to_string(), - value: OsString::from(env::var(x).expect("must exist")), - }) - .collect() -} - /// Return IsolatedContext ready for executing a command inside isolation /// # Arguments /// * `image` - container image that would be used to run the VM @@ -136,21 +120,3 @@ pub(crate) fn is_isolated() -> Result { debug!("systemd-detect-virt returned: {}", virt); Ok(virt == "systemd-nspawn") } - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_default_envs() { - assert_eq!(default_passthrough_envs(), Vec::new()); - env::set_var("RUST_LOG", "hello"); - assert_eq!( - default_passthrough_envs(), - vec![KvPair { - key: "RUST_LOG".into(), - value: "hello".into(), - }], - ); - } -} diff --git a/antlir/antlir2/antlir2_vm/src/main.rs b/antlir/antlir2/antlir2_vm/src/main.rs index e6f4d6d5471..51a38818747 100644 --- a/antlir/antlir2/antlir2_vm/src/main.rs +++ b/antlir/antlir2/antlir2_vm/src/main.rs @@ -27,7 +27,6 @@ use anyhow::Context; use clap::Args; use clap::Parser; use clap::Subcommand; -use image_test_lib::KvPair; use image_test_lib::Test; use json_arg::JsonFile; use tempfile::tempdir; @@ -35,7 +34,6 @@ use tracing::debug; use tracing_subscriber::filter::LevelFilter; use tracing_subscriber::prelude::*; -use crate::isolation::default_passthrough_envs; use crate::isolation::is_isolated; use crate::isolation::isolated; use crate::isolation::Platform; @@ -44,6 +42,7 @@ use crate::types::MachineOpts; use crate::types::RuntimeOpts; use crate::types::VMArgs; use crate::utils::console_output_path_for_tpx; +use crate::utils::env_names_to_kvpairs; use crate::utils::log_command; use crate::vm::VM; @@ -84,10 +83,9 @@ struct IsolateCmdArgs { /// Path to container image. #[arg(long)] image: PathBuf, - /// Set these env vars in the container. If VM executes a command, these - /// env vars will also be prepended to the command. + /// Pass through these environment variables into the container and VM, if they exist. #[arg(long)] - setenv: Vec, + passenv: Vec, /// Args for run command #[clap(flatten)] run_cmd_args: RunCmdArgs, @@ -111,9 +109,8 @@ fn run(args: &RunCmdArgs) -> Result<()> { /// Enter isolated container and then respawn itself inside it with `run` /// command and its parameters. fn respawn(args: &IsolateCmdArgs) -> Result<()> { - let mut envs = default_passthrough_envs(); - envs.extend(args.setenv.clone()); let mut vm_args = args.run_cmd_args.vm_args.clone(); + let envs = env_names_to_kvpairs(args.passenv.clone()); vm_args.command_envs = envs.clone(); // Let's always capture console output unless it's console mode @@ -142,21 +139,6 @@ fn respawn(args: &IsolateCmdArgs) -> Result<()> { Ok(()) } -/// Merge all sources of our envs into final list of env vars we should use -/// everywhere for tests. Dedup is handled by functions that use the result. -fn get_test_envs(from_cli: &[KvPair]) -> Vec { - // This handles common envs like RUST_LOG - let mut envs = default_passthrough_envs(); - envs.extend_from_slice(from_cli); - // forward test runner env vars to the inner test - for (key, val) in std::env::vars() { - if key.starts_with("TEST_PILOT") { - envs.push((key, OsString::from(val)).into()); - } - } - envs -} - /// Validated `VMArgs` and other necessary metadata for tests. struct ValidatedVMArgs { /// VMArgs that will be passed into the VM with modified fields @@ -167,7 +149,7 @@ struct ValidatedVMArgs { /// Further validate `VMArgs` parsed by clap and generate a new `VMArgs` with /// content specific to test execution. -fn get_test_vm_args(orig_args: &VMArgs, cli_envs: &[KvPair]) -> Result { +fn get_test_vm_args(orig_args: &VMArgs, cli_envs: Vec) -> Result { if orig_args.timeout_secs.is_none() { return Err(anyhow!("Test command must specify --timeout-secs.")); } @@ -177,7 +159,15 @@ fn get_test_vm_args(orig_args: &VMArgs, cli_envs: &[KvPair]) -> Result R /// inputs instead of allowing caller to pass them in. Some inputs are parsed /// from the test command. fn test(args: &IsolateCmdArgs) -> Result<()> { - let validated_args = get_test_vm_args(&args.run_cmd_args.vm_args, &args.setenv)?; + let validated_args = get_test_vm_args(&args.run_cmd_args.vm_args, args.passenv.clone())?; let mut command = if validated_args.is_list { list_test_command(args, &validated_args) } else { @@ -292,24 +282,11 @@ fn main() -> Result<()> { #[cfg(test)] mod test { + use image_test_lib::KvPair; + use super::*; use crate::types::VMModeArgs; - #[test] - fn test_get_test_envs() { - env::set_var("RUST_LOG", "hello"); - env::set_var("TEST_PILOT_A", "A"); - let from_cli = vec![KvPair::from(("foo", "bar"))]; - assert_eq!( - get_test_envs(&from_cli), - vec![ - KvPair::from(("RUST_LOG", "hello")), - KvPair::from(("foo", "bar")), - KvPair::from(("TEST_PILOT_A", "A")), - ], - ) - } - #[test] fn test_get_test_vm_args() { let valid = VMArgs { @@ -320,25 +297,35 @@ mod test { }, ..Default::default() }; - let empty_env = Vec::::new(); let mut expected = valid.clone(); expected.mode.command = Some(vec![OsString::from("whatever")]); - let parsed = get_test_vm_args(&valid, &empty_env).expect("Parsing should succeed"); - assert_eq!(parsed.inner, expected); + let parsed = get_test_vm_args(&valid, vec![]).expect("Parsing should succeed"); + assert_eq!(parsed.inner.mode, expected.mode); assert!(!parsed.is_list); let mut timeout = valid.clone(); timeout.timeout_secs = None; - assert!(get_test_vm_args(&timeout, &empty_env).is_err()); + assert!(get_test_vm_args(&timeout, vec![]).is_err()); let mut output_dirs = valid.clone(); output_dirs.output_dirs = vec![PathBuf::from("/some")]; - assert!(get_test_vm_args(&output_dirs, &empty_env).is_err()); + assert!(get_test_vm_args(&output_dirs, vec![]).is_err()); - let mut command = valid; + let mut command = valid.clone(); command.mode.command = None; - assert!(get_test_vm_args(&command, &empty_env).is_err()); + assert!(get_test_vm_args(&command, vec![]).is_err()); command.mode.command = Some(vec![OsString::from("invalid")]); - assert!(get_test_vm_args(&command, &empty_env).is_err()); + assert!(get_test_vm_args(&command, vec![]).is_err()); + + let env_var_test = valid; + std::env::set_var("TEST_PILOT_A", "A"); + let parsed = get_test_vm_args(&env_var_test, vec![]).expect("Parsing should succeed"); + assert!( + parsed + .inner + .command_envs + .iter() + .any(|x| *x == KvPair::from(("TEST_PILOT_A", "A"))) + ); } } diff --git a/antlir/antlir2/antlir2_vm/src/utils.rs b/antlir/antlir2/antlir2_vm/src/utils.rs index 2647b12075b..44d8f940b1f 100644 --- a/antlir/antlir2/antlir2_vm/src/utils.rs +++ b/antlir/antlir2/antlir2_vm/src/utils.rs @@ -5,12 +5,14 @@ * LICENSE file in the root directory of this source tree. */ +use std::collections::HashSet; use std::fs; use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; use std::process::Command; +use image_test_lib::KvPair; use tracing::debug; use tracing::error; @@ -85,6 +87,24 @@ pub(crate) fn console_output_path_for_tpx() -> Result, std::io:: } } +/// Convert a list of env names into KvPair with its values +pub(crate) fn env_names_to_kvpairs(env_names: Vec) -> Vec { + let mut names: HashSet<_> = env_names.into_iter().collect(); + // If these env exist, always pass them through. + ["RUST_LOG", "RUST_BACKTRACE", "ANTLIR_BUCK"] + .iter() + .for_each(|name| { + names.insert(name.to_string()); + }); + names + .iter() + .filter_map(|name| match std::env::var(name) { + Ok(value) => Some(KvPair::from((name, value))), + _ => None, + }) + .collect() +} + #[cfg(test)] /// Helper function for converting qemu args to a single string for asserting in tests. /// This is usually only needed for string only functions like `contains`. @@ -97,6 +117,10 @@ pub(crate) fn qemu_args_to_string(args: &[std::ffi::OsString]) -> String { #[cfg(test)] mod test { + use std::collections::HashMap; + use std::env; + use std::ffi::OsString; + use super::*; #[test] @@ -110,4 +134,103 @@ mod test { format!("Program: `hello`. Args: `{:?}`", vec!["world"]), ); } + + struct EnvTest { + envs: Vec<(&'static str, &'static str)>, + passenv: Vec<&'static str>, + result: HashMap, + } + + const ALLOWED_ENV_NAMES: &[&str] = &[ + "RUST_LOG", + "RUST_BACKTRACE", + "ANTLIR_BUCK", + "TEST_PILOT_A", + "OTHER", + ]; + + impl EnvTest { + fn new( + envs: Vec<(&'static str, &'static str)>, + passenv: Vec<&'static str>, + result: Vec<(&'static str, &'static str)>, + ) -> Self { + // We have to clear all envs across tests, so we must know the ones to clear. + envs.iter().for_each(|(name, _)| { + assert!( + ALLOWED_ENV_NAMES.contains(name), + "{name} not allowed for testing" + ) + }); + Self { + envs, + passenv, + result: result + .into_iter() + .map(|(k, v)| (k.into(), v.into())) + .collect(), + } + } + + fn test(&self) { + ALLOWED_ENV_NAMES.iter().for_each(|name| { + env::remove_var(name); + }); + self.envs.iter().for_each(|(name, val)| { + env::set_var(name, val); + }); + let result: HashMap<_, _> = + env_names_to_kvpairs(self.passenv.iter().map(|s| s.to_string()).collect()) + .drain(..) + .map(|pair| (pair.key, pair.value)) + .collect(); + assert_eq!(result, self.result); + } + } + + #[test] + fn test_env_names_to_kvpairs() { + [ + EnvTest::new(vec![], vec![], vec![]), + // Reads nothing + EnvTest::new( + vec![], + vec!["RUST_LOG", "ANTLIR_BUCK", "TEST_PILOT_A", "OTHER"], + vec![], + ), + // Always pass through + EnvTest::new( + vec![ + ("RUST_LOG", "info"), + ("ANTLIR_BUCK", "1"), + ("TEST_PILOT_A", "A"), + ("OTHER", "other"), + ], + vec![], + vec![("RUST_LOG", "info"), ("ANTLIR_BUCK", "1")], + ), + // Selection + EnvTest::new( + vec![("TEST_PILOT_A", "A"), ("OTHER", "other")], + vec!["TEST_PILOT_A"], + vec![("TEST_PILOT_A", "A")], + ), + // Mixed + EnvTest::new( + vec![ + ("RUST_LOG", "info"), + ("TEST_PILOT_A", "A"), + ("OTHER", "other"), + ], + vec!["TEST_PILOT_A"], + vec![("TEST_PILOT_A", "A"), ("RUST_LOG", "info")], + ), + ] + .iter() + .enumerate() + .for_each(|(i, test)| { + println!("Running test #{i}"); + test.test(); + }); + } } diff --git a/antlir/antlir2/antlir2_vm/src/vm.rs b/antlir/antlir2/antlir2_vm/src/vm.rs index 6ab0ba625c0..d494c577591 100644 --- a/antlir/antlir2/antlir2_vm/src/vm.rs +++ b/antlir/antlir2/antlir2_vm/src/vm.rs @@ -5,7 +5,6 @@ * LICENSE file in the root directory of this source tree. */ -use std::collections::HashMap; use std::collections::HashSet; use std::ffi::OsString; use std::fs; @@ -26,7 +25,6 @@ use std::thread::JoinHandle; use std::time::Duration; use std::time::Instant; -use image_test_lib::KvPair; use thiserror::Error; use tracing::debug; use tracing::error; @@ -524,16 +522,8 @@ impl VM { // Just wait for the human that's trying to debug with console self.wait_for_timeout::<()>(f.into_inner(), start_ts, None)?; } else if !self.args.mode.container { - // Open ssh shell or run command over ssh - let dedup_envs: HashMap<_, _> = self - .args - .command_envs - .iter() - .cloned() - .map(|p| (p.key, p.value)) - .collect(); let mut ssh_cmd = GuestSSHCommand::new()?.ssh_cmd(); - dedup_envs.iter().map(KvPair::from).for_each(|kv| { + self.args.command_envs.iter().for_each(|kv| { ssh_cmd.arg(kv.to_os_string()); }); if let Some(command) = &self.args.mode.command { diff --git a/antlir/antlir2/antlir2_vm/tests/BUCK b/antlir/antlir2/antlir2_vm/tests/BUCK index 9aa763b8296..e4fd549a821 100644 --- a/antlir/antlir2/antlir2_vm/tests/BUCK +++ b/antlir/antlir2/antlir2_vm/tests/BUCK @@ -9,7 +9,9 @@ vm.cpp_test( srcs = ["test.cpp"], env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, + run_as_bundle = True, vm_host = get_vm(), ) @@ -18,7 +20,9 @@ vm.python_test( srcs = ["test.py"], env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, + run_as_bundle = True, vm_host = get_vm(), ) @@ -29,7 +33,9 @@ vm.rust_test( crate_root = "test.rs", env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, + run_as_bundle = True, vm_host = get_vm(), ) @@ -37,6 +43,7 @@ vm.sh_test( name = "sh-test", env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, test = "test.sh", vm_host = get_vm(), @@ -51,23 +58,13 @@ vm.sh_test( crate_root = "test.rs", env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, vm_host = target, ) for name, target in PRECONFIGURED_VM.items() ] -# Verify bundle -vm.cpp_test( - name = "cpp-test-bundle", - srcs = ["test.cpp"], - env = { - "ANTLIR2_TEST": "1", - }, - run_as_bundle = True, - vm_host = get_vm(), -) - # Test for sidecar rust_binary( name = "sidecar", diff --git a/antlir/antlir2/antlir2_vm/tests/test.cpp b/antlir/antlir2/antlir2_vm/tests/test.cpp index 02386e102c6..46db47bea86 100644 --- a/antlir/antlir2/antlir2_vm/tests/test.cpp +++ b/antlir/antlir2/antlir2_vm/tests/test.cpp @@ -7,6 +7,7 @@ #include #include +#include TEST(CppTest, TestIsRoot) { EXPECT_EQ(getuid(), 0); @@ -15,3 +16,9 @@ TEST(CppTest, TestIsRoot) { TEST(CppTest, TestEnvPropagated) { ASSERT_STREQ(std::getenv("ANTLIR2_TEST"), "1"); } + +TEST(CppTest, TestEnvArtifactExists) { + auto artifact = std::getenv("ENV_ARTIFACT"); + ASSERT_NE(nullptr, artifact); + ASSERT_TRUE(std::filesystem::exists(artifact)); +} diff --git a/antlir/antlir2/antlir2_vm/tests/test.py b/antlir/antlir2/antlir2_vm/tests/test.py index 9bf17b48e62..e3b3f2785ed 100644 --- a/antlir/antlir2/antlir2_vm/tests/test.py +++ b/antlir/antlir2/antlir2_vm/tests/test.py @@ -14,3 +14,8 @@ def test_is_root(self) -> None: def test_env_propagated(self) -> None: self.assertEqual("1", os.getenv("ANTLIR2_TEST")) + + def test_env_artifact_exists(self) -> None: + artifact = os.getenv("ENV_ARTIFACT") + self.assertNotEqual(None, artifact) + self.assertTrue(os.path.exists(artifact)) diff --git a/antlir/antlir2/antlir2_vm/tests/test.rs b/antlir/antlir2/antlir2_vm/tests/test.rs index fee064a1622..6677aed215d 100644 --- a/antlir/antlir2/antlir2_vm/tests/test.rs +++ b/antlir/antlir2/antlir2_vm/tests/test.rs @@ -5,6 +5,8 @@ * LICENSE file in the root directory of this source tree. */ +use std::path::Path; + #[link(name = "c")] extern "C" { fn geteuid() -> u32; @@ -19,3 +21,10 @@ fn is_root() { fn env_propagated() { assert_eq!("1", std::env::var("ANTLIR2_TEST").expect("env var missing")); } + +#[test] +fn env_artifact_exists() { + let artifact = std::env::var("ENV_ARTIFACT"); + assert!(artifact.is_ok()); + assert!(Path::new(&artifact.expect("Already checked")).exists()); +} diff --git a/antlir/antlir2/antlir2_vm/tests/test.sh b/antlir/antlir2/antlir2_vm/tests/test.sh index 2fdcff11490..c690b7fb13b 100755 --- a/antlir/antlir2/antlir2_vm/tests/test.sh +++ b/antlir/antlir2/antlir2_vm/tests/test.sh @@ -15,3 +15,11 @@ if [ "${ANTLIR2_TEST}" != "1" ]; then echo "Env var ANTLIR2_TEST is wrong or missing!" exit 2 fi + +if [ -z "${ENV_ARTIFACT}" ]; then + echo "Env var ENV_ARTIFACT is missing!" + exit 3 +elif ! [ -f "${ENV_ARTIFACT}" ]; then + echo "Env var ENV_ARTIFACT does not point to a valid target!" + exit 3 +fi