azure_identity/
env.rs

1use azure_core::{
2    error::{ErrorKind, ResultExt},
3    Error,
4};
5use std::collections::HashMap;
6
7#[derive(Debug, Clone)]
8pub(crate) enum Env {
9    Process(ProcessEnv),
10    Mem(MemEnv),
11}
12
13impl Default for Env {
14    fn default() -> Self {
15        Self::Process(ProcessEnv)
16    }
17}
18
19impl Env {
20    pub fn var(&self, key: &str) -> azure_core::Result<String> {
21        match self {
22            Env::Process(env) => env.var(key),
23            Env::Mem(env) => env.var(key),
24        }
25    }
26}
27
28impl From<ProcessEnv> for Env {
29    fn from(env: ProcessEnv) -> Self {
30        Self::Process(env)
31    }
32}
33
34impl From<MemEnv> for Env {
35    fn from(env: MemEnv) -> Self {
36        Self::Mem(env)
37    }
38}
39
40/// The standard environment that gets variables from the process.
41#[derive(Debug, Clone, Default)]
42pub(crate) struct ProcessEnv;
43
44impl ProcessEnv {
45    fn var(&self, key: &str) -> azure_core::Result<String> {
46        std::env::var(key).with_context(ErrorKind::Io, || {
47            format!("environment variable {} not set", key)
48        })
49    }
50}
51
52/// An environment that stores and gets variables in memory.
53#[derive(Debug, Clone, Default)]
54pub(crate) struct MemEnv {
55    vars: HashMap<String, String>,
56}
57
58impl MemEnv {
59    fn var(&self, key: &str) -> azure_core::Result<String> {
60        self.vars.get(key).cloned().ok_or_else(|| {
61            Error::message(
62                ErrorKind::Io,
63                format!("environment variable {} not set", key),
64            )
65        })
66    }
67}
68
69impl From<&[(&str, &str)]> for Env {
70    fn from(pairs: &[(&str, &str)]) -> Self {
71        let mut vars = HashMap::new();
72        for (k, v) in pairs {
73            vars.insert(k.to_string(), v.to_string());
74        }
75        Self::Mem(MemEnv { vars })
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_env_var() {
85        let env = Env::from(&[("CHRISTMAS_GRINCH", "You're a mean one, Mr. Grinch")][..]);
86        assert_eq!(
87            env.var("CHRISTMAS_GRINCH").unwrap(),
88            "You're a mean one, Mr. Grinch"
89        );
90    }
91
92    // test ProcessEnv::var() returns an error when the environment variable is not set
93    #[test]
94    fn test_env_var_not_set() {
95        let env = ProcessEnv {};
96        assert!(env.var("CHRISTMAS_GRINCH").is_err());
97    }
98
99    // test MemEnv::var() returns an error when the environment variable is not set
100    #[test]
101    fn test_mem_env_var_not_set() {
102        let env = MemEnv::default();
103        assert!(env.var("CHRISTMAS_GRINCH").is_err());
104    }
105
106    // test MemEnv::var() returns valid entries when multiple environment variables are set
107    #[test]
108    fn test_mem_env_var_multiple() {
109        let env = Env::from(
110            &[
111                ("CHRISTMAS_GRINCH", "You're a mean one, Mr. Grinch"),
112                ("CHRISTMAS_TREE", "O Christmas Tree, O Christmas Tree"),
113                ("CHRISTMAS_SNOW", "Let it snow, let it snow, let it snow"),
114            ][..],
115        );
116        assert_eq!(
117            env.var("CHRISTMAS_GRINCH").unwrap(),
118            "You're a mean one, Mr. Grinch"
119        );
120        assert_eq!(
121            env.var("CHRISTMAS_TREE").unwrap(),
122            "O Christmas Tree, O Christmas Tree"
123        );
124        assert_eq!(
125            env.var("CHRISTMAS_SNOW").unwrap(),
126            "Let it snow, let it snow, let it snow"
127        );
128    }
129}