Allow setting environment variables in setup hook

README.md

@@ -44,6 +44,9 @@ use testscript_rs::testscript;
 fn test_my_cli() {
     testscript::run("testdata")
         .setup(|env| {
+            // Set environment variables for tests
+            env.set_env_var("MY_APP_CONFIG", "/path/to/config");
+
             // Compile your CLI tool
             std::process::Command::new("cargo")
                 .args(["build", "--bin", "my-cli"])

src/lib.rs

@@ -154,6 +154,10 @@ fn run(params: &mut RunParams, test_data_glob: &str) -> Result<()> {
 ///
 /// testscript::run("testdata")
 ///     .setup(|env| {
+///         // Set environment variables for your CLI
+///         env.set_env_var("MY_APP_MODE", "testing");
+///         env.set_env_var("LOG_LEVEL", "debug");
+///         
 ///         // Compile your CLI tool
 ///         std::process::Command::new("cargo")
 ///             .args(["build", "--bin", "my-cli"])
@@ -203,11 +207,30 @@ impl Builder {
 
     /// Add a setup function that runs before each test script
     ///
-    /// The setup function receives a reference to the test environment and can
-    /// perform actions like compiling binaries or setting up test data.
+    /// The setup function receives a mutable reference to the test environment and can
+    /// perform actions like compiling binaries, setting up test data, or setting environment variables.
+    ///
+    /// # Examples
+    /// ```no_run
+    /// use testscript_rs::testscript;
+    ///
+    /// testscript::run("testdata")
+    ///     .setup(|env| {
+    ///         // Set an environment variable
+    ///         env.set_env_var("FOO", "bar");
+    ///         
+    ///         // Compile a binary
+    ///         std::process::Command::new("cargo")
+    ///             .args(["build", "--bin", "my-cli"])
+    ///             .status()?;
+    ///         Ok(())
+    ///     })
+    ///     .execute()
+    ///     .unwrap();
+    /// ```
     pub fn setup<F>(mut self, func: F) -> Self
     where
-        F: Fn(&TestEnvironment) -> Result<()> + 'static,
+        F: Fn(&mut TestEnvironment) -> Result<()> + 'static,
     {
         self.params = self.params.setup(func);
         self

src/run/execution.rs

@@ -49,7 +49,7 @@ pub fn run_script_impl(script_path: &Path, params: &RunParams) -> Result<()> {
 
     // Run setup hook if provided
     if let Some(setup) = &params.setup {
-        setup(&env)?;
+        setup(&mut env)?;
     }
 
     // Track script updates if we're in update mode

src/run/params.rs

@@ -8,7 +8,7 @@ use std::collections::HashMap;
 pub type CommandFn = fn(&mut TestEnvironment, &[String]) -> Result<()>;
 
 /// Type alias for a setup function
-pub type SetupFn = Box<dyn Fn(&TestEnvironment) -> Result<()>>;
+pub type SetupFn = Box<dyn Fn(&mut TestEnvironment) -> Result<()>>;
 
 /// Configuration parameters for running tests
 pub struct RunParams {
@@ -73,7 +73,7 @@ impl RunParams {
     /// Set a setup function to run before each script
     pub fn setup<F>(mut self, func: F) -> Self
     where
-        F: Fn(&TestEnvironment) -> Result<()> + 'static,
+        F: Fn(&mut TestEnvironment) -> Result<()> + 'static,
     {
         self.setup = Some(Box::new(func));
         self

tests/setup_hook.rs

@@ -87,37 +87,66 @@ fn test_setup_hook_environment() {
     let temp_dir = TempDir::new().unwrap();
     let script_path = temp_dir.path().join("env_test.txt");
 
-    let script_content = r#"# Test setup hook can modify environment
-env TEST_FROM_SETUP=should_be_set
-exec echo "Value: $TEST_FROM_SETUP"
-stdout "Value: hello_from_setup""#;
+    // Use sh -c to allow shell variable expansion
+    let script_content = if cfg!(windows) {
+        r#"# Test setup hook can modify environment (Windows)
+exec cmd /c "echo Value: %TEST_FROM_SETUP%"
+stdout "Value: hello_from_setup""#
+    } else {
+        r#"# Test setup hook can modify environment (Unix)
+exec sh -c "echo Value: $TEST_FROM_SETUP"
+stdout "Value: hello_from_setup"#
+    };
 
     fs::write(&script_path, script_content).unwrap();
 
     // Create RunParams with setup hook that sets an environment variable
     let params = RunParams::new().setup(|env| {
-        // The setup hook has access to the TestEnvironment but can't modify it
-        // This demonstrates the current API limitation - we can access but not modify
-        println!("Setup running in: {}", env.work_dir.display());
-
-        // We could write files that contain environment setup
-        let env_script = env.work_dir.join("setup_env.sh");
-        std::fs::write(&env_script, "export TEST_FROM_SETUP=hello_from_setup")?;
+        // Now we can modify the environment directly in the setup hook
+        env.set_env_var("TEST_FROM_SETUP", "hello_from_setup");
         Ok(())
     });
 
-    let mut params = params;
-    // Manually set the environment variable for this test
-    params = params.condition("TEST_FROM_SETUP", true);
+    let result = run_script(&script_path, &params);
+    assert!(
+        result.is_ok(),
+        "Environment setup test should pass: {:?}",
+        result
+    );
+}
+
+#[test]
+fn test_setup_hook_multiple_env_vars() {
+    let temp_dir = TempDir::new().unwrap();
+    let script_path = temp_dir.path().join("multi_env_test.txt");
+
+    // Use sh -c to allow shell variable expansion
+    let script_content = if cfg!(windows) {
+        r#"# Test setup hook can set multiple environment variables (Windows)
+exec cmd /c "echo FOO=%FOO% BAR=%BAR% BAZ=%BAZ%"
+stdout "FOO=value1 BAR=value2 BAZ=value3""#
+    } else {
+        r#"# Test setup hook can set multiple environment variables (Unix)
+exec sh -c "echo FOO=$FOO BAR=$BAR BAZ=$BAZ"
+stdout "FOO=value1 BAR=value2 BAZ=value3"#
+    };
+
+    fs::write(&script_path, script_content).unwrap();
+
+    // Create RunParams with setup hook that sets multiple environment variables
+    let params = RunParams::new().setup(|env| {
+        env.set_env_var("FOO", "value1");
+        env.set_env_var("BAR", "value2");
+        env.set_env_var("BAZ", "value3");
+        Ok(())
+    });
 
     let result = run_script(&script_path, &params);
-    // This test demonstrates the current limitation - setup can't modify the running environment
-    if result.is_err() {
-        println!(
-            "Environment setup test failed (expected with current API): {:?}",
-            result
-        );
-    }
+    assert!(
+        result.is_ok(),
+        "Multiple environment variables test should pass: {:?}",
+        result
+    );
 }
 
 #[test]
@@ -144,3 +173,85 @@ exec echo "Should not execute""#;
         .to_string()
         .contains("Setup deliberately failed"));
 }
+
+#[test]
+fn test_setup_env_vars_presence_and_absence() {
+    let temp_dir = TempDir::new().unwrap();
+    let script_path = temp_dir.path().join("env_presence_test.txt");
+
+    // Test that verifies:
+    // 1. Variables set in setup ARE available in the test
+    // 2. Variables NOT set in setup are NOT available in the test
+    let script_content = if cfg!(windows) {
+        r#"# Test environment variable presence and absence (Windows)
+# Check that SET_VAR is available using printenv (no shell needed)
+exec cmd /c set SET_VAR
+stdout "SET_VAR=value_from_setup"
+
+# Also verify with shell expansion
+exec cmd /c "echo SET_VAR=%SET_VAR%"
+stdout "SET_VAR=value_from_setup"
+
+# Check that UNSET_VAR is not available
+! exec cmd /c set UNSET_VAR"#
+    } else {
+        r#"# Test environment variable presence and absence (Unix)
+# Check that SET_VAR is available using printenv (no shell needed)
+exec printenv SET_VAR
+stdout "value_from_setup"
+
+# Also verify with shell expansion
+exec sh -c "echo SET_VAR=$SET_VAR"
+stdout "SET_VAR=value_from_setup"
+
+# Check that UNSET_VAR is not available
+! exec printenv UNSET_VAR"#
+    };
+
+    fs::write(&script_path, script_content).unwrap();
+
+    // Create RunParams with setup hook that only sets SET_VAR, not UNSET_VAR
+    let params = RunParams::new().setup(|env| {
+        env.set_env_var("SET_VAR", "value_from_setup");
+        // Deliberately NOT setting UNSET_VAR to verify it's not available
+        Ok(())
+    });
+
+    let result = run_script(&script_path, &params);
+    assert!(
+        result.is_ok(),
+        "Environment presence/absence test should pass: {:?}",
+        result
+    );
+}
+
+#[test]
+fn test_env_command_prints_all_vars() {
+    let temp_dir = TempDir::new().unwrap();
+    let script_path = temp_dir.path().join("env_print_test.txt");
+
+    // Test that 'env' without args prints all environment variables
+    let script_content = r#"# Test env command prints all variables
+env VAR1=value1
+env VAR2=value2
+
+# The env command should be callable and set variables
+# (printing is done to stdout which we can't easily test here,
+# but we verify the vars are set by using them)
+exec sh -c "echo $VAR1"
+stdout "value1"
+
+exec sh -c "echo $VAR2"
+stdout "value2"
+"#;
+
+    fs::write(&script_path, script_content).unwrap();
+
+    let params = RunParams::new().setup(|env| {
+        env.set_env_var("SETUP_VAR", "from_setup");
+        Ok(())
+    });
+
+    let result = run_script(&script_path, &params);
+    assert!(result.is_ok(), "Env command test should pass: {:?}", result);
+}