DEV Community

Cover image for Testing Environment-Dependent Code in Rust
Federico Joaquín Barberón
Federico Joaquín Barberón

Posted on • Edited on • Originally published at federicobarberon.netlify.app

Testing Environment-Dependent Code in Rust

The problem

When we have a function that interacts with external resources, testing that code might be difficult because side-effects could impact on other tests. Environment variables are key-value pairs that are part of the process environment. One change in an env variable changes the process state, so this mutation persists throughout the process lifetime.

In Rust, unit tests are compiled into one binary, which means that all tests run in the same process. That is why we need to be careful when testing code that modifies env vars, because those changes may affect other tests and could produce an undesirable behaviour.

On top of that, Rust tests run in parallel by default, so we need to ensure that any thread modifying environment variables has exclusive access to the process environment.

Ensuring exclusive access to the Process Environment

One way to guarantee that there is only one thread with exclusive access to the process environment is to simply use one thread! We can achieve that by running cargo test -- --test-threads=1. This is the simplest solution to the problem, however, if the number of unit tests in our project grows a lot, the execution time of the tests could be annoying because all tests, including the ones that do not need to be single-threaded, run on a single thread anyway.

Another solution more appropriate to this case is to use the macro #[serial] from the serial_test crate. This macro allows us to select the tests that we want to run in serial, while leaving the others unchanged.

#[cfg(test)]
mod tests {
    use super::*;
    use serial_test::serial;

    #[test]
    fn normal_test1() {
        ...
    }

    #[test]
    fn normal_test2() {
        ...
    }

    #[test]
    #[serial]
    fn serial_test1() {
        ...
    }

    #[test]
    #[serial]
    fn serial_test2() {
        ...
    }
}
Enter fullscreen mode Exit fullscreen mode

In this example, we guarantee that serial_test1 and serial_test2 run in serial, while (maybe) at the same time normal_test1 and normal_test2 run in parallel, so we get the best of both worlds.

Why serial_test alone isn't enough

With serial_test we solve the problem of concurrent access to the process environment, but we still have the problem that changes to env vars persist across the tests. We need to find a way to restore the previous state of the env vars that were modified, so that other tests are not contaminated.

We may be tempted to do it with something like this:

#[test]
#[serial]
fn test() {
    let var = "PATH";
    let previous_state = std::env::var_os(var).unwrap();

    ... // test the function that modifies <var> env variable

    // SAFETY: There are no other threads accessing the process environment because all tests that do it have #[serial] macro.
    unsafe {
        std::env::set_var(var, previous_state);
    }
}
Enter fullscreen mode Exit fullscreen mode

However, this has a problem. The last statement may NEVER run if the test panics, so this is not a solution at all.

The RAII guard

Rust implements RAII (Resource Acquisition Is Initialization), which means that once the owner of some data goes out of scope, the drop() method of that data is called and the data is freed.

We can take advantage of this by implementing a common design pattern in Rust: the RAII guard. We create an object that holds the original values of the environment variables that we are going to change, and implement the Drop trait so that it restores the process environment state when the object goes out of scope (even if the test panics).

use std::{ffi::OsString, env};

struct EnvGuard {
    key: OsString,
    value: Option<OsString>
}

impl EnvGuard {
    // SAFETY:
    // The caller must guarantee exclusive access to the process environment
    // for the lifetime of this guard.
    unsafe fn capture(key: OsString) -> Self {
        Self {
            value: env::var_os(&key),
            key,
        }
    }
}

impl Drop for EnvGuard {
    fn drop(&mut self) {
        // SAFETY:
        // The caller of `capture` guaranteed exclusive access to the process
        // environment for the lifetime of this guard.
        unsafe {
            match &self.value {
                Some(val) => env::set_var(&self.key, val),
                None => env::remove_var(&self.key)
            }
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

Note that capture itself performs no unsafe operation — it only calls env::var_os, which is safe. It's marked unsafe to push the safety contract to the call site, since that contract ("exclusive access to the environment") covers the guard's entire lifetime, not just this one function.

Now we can use it in our tests

#[test]
#[serial]
fn test() {
    let var = "PATH";

    // SAFETY: There are no other threads accessing the process environment because all tests that do it have #[serial] macro.
    let _guard = unsafe { EnvGuard::capture(var.into()) };

    ... // test the function that modifies <var> env variable

    // When `_guard` goes out of scope, even if the test panics, the env var is restored to the original value.
}
Enter fullscreen mode Exit fullscreen mode

Result

We end up with a simple and idiomatic solution to a common problem of testing environment variables. It is worth mentioning that this solution is easy to extend to multiple variables, and also it is easy to adapt to other external state besides the environment variables.


If you spot something incorrect in this post, let me know — I'm still learning this too :)

Top comments (4)

Collapse
 
xulingfeng profile image
xulingfeng

The EnvGuard is nice. One extension worth considering: snapshot the full env at test start and diff at teardown. Catches leaks from dependencies that mutate env vars your code never touched. Clean first post.

Collapse
 
federicobarberon profile image
Federico Joaquín Barberón

Good one! My EnvGuard solution came from a shell in Rust that I'm working on as an exercise, and I didn't even consider the possibility of external dependencies modifying the env vars. I think it's worth paying the cost of maintaining a snapshot of the full env in tests to catch some potential bugs that would otherwise be difficult to find. Thanks for your contribution! :)

Collapse
 
xulingfeng profile image
xulingfeng

Thanks man, glad it helped 🙌 The RAII + full snapshot combo would make a solid setup. Env vars are one of those things nobody looks at — until they blow up a CI run at 2 AM 😄 Good luck with the shell project!

Thread Thread
 
federicobarberon profile image
Federico Joaquín Barberón

😂 Thank you!