feat(aws): Add profile aliases (#3699)

* Rename m.aws.alias_region to alias_name

* Add aws profile aliases

* Document aws.profile_aliases, with examples

* Add tests for new aws.profile_aliases feature

* Tidy alias_handling a bit

2 files changed, 52 insertions, 12 deletions

diff --git a/src/configs/aws.rs b/src/configs/aws.rs
index 8cb2b2dd5..2bd6287b0 100644
--- a/src/configs/aws.rs
+++ b/src/configs/aws.rs
@@ -10,6 +10,7 @@ pub struct AwsConfig<'a> {
     pub style: &'a str,
     pub disabled: bool,
     pub region_aliases: HashMap<String, &'a str>,
+    pub profile_aliases: HashMap<String, &'a str>,
     pub expiration_symbol: &'a str,
 }
 
@@ -21,6 +22,7 @@ impl<'a> Default for AwsConfig<'a> {
             style: "bold yellow",
             disabled: false,
             region_aliases: HashMap::new(),
+            profile_aliases: HashMap::new(),
             expiration_symbol: "X",
         }
     }
diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index f59210657..9ff03a9ce 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -111,11 +111,11 @@ fn get_credentials_duration(context: &Context, aws_profile: Option<&Profile>) ->
     Some(expiration_date.timestamp() - chrono::Local::now().timestamp())
 }
 
-fn alias_region(region: String, aliases: &HashMap<String, &str>) -> String {
-    match aliases.get(&region) {
-        None => region,
-        Some(alias) => (*alias).to_string(),
-    }
+fn alias_name(name: Option<String>, aliases: &HashMap<String, &str>) -> Option<String> {
+    name.as_ref()
+        .and_then(|n| aliases.get(n))
+        .map(|&a| a.to_string())
+        .or(name)
 }
 
 fn get_credential_process(context: &Context, aws_profile: Option<&Profile>) -> Option<String> {
@@ -187,12 +187,6 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         return None;
     }
 
-    let mapped_region = if let Some(aws_region) = aws_region {
-        Some(alias_region(aws_region, &config.region_aliases))
-    } else {
-        None
-    };
-
     let duration = {
         get_credentials_duration(context, aws_profile.as_ref()).map(|duration| {
             if duration > 0 {
@@ -203,6 +197,10 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         })
     };
 
+    let mapped_region = alias_name(aws_region, &config.region_aliases);
+
+    let mapped_profile = alias_name(aws_profile, &config.profile_aliases);
+
     let parsed = StringFormatter::new(config.format).and_then(|formatter| {
         formatter
             .map_meta(|variable, _| match variable {
@@ -214,7 +212,7 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
                 _ => None,
             })
             .map(|variable| match variable {
-                "profile" => aws_profile.as_ref().map(Ok),
+                "profile" => mapped_profile.as_ref().map(Ok),
                 "region" => mapped_region.as_ref().map(Ok),
                 "duration" => duration.as_ref().map(Ok),
                 _ => None,
@@ -370,6 +368,46 @@ mod tests {
     }
 
     #[test]
+    fn profile_set_with_alias() {
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_PROFILE", "CORPORATION-CORP_astronauts_ACCESS_GROUP")
+            .env("AWS_REGION", "ap-northeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
+            .config(toml::toml! {
+                [aws.profile_aliases]
+                CORPORATION-CORP_astronauts_ACCESS_GROUP = "astro"
+            })
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astro (ap-northeast-2) ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn region_and_profile_both_set_with_alias() {
+        let actual = ModuleRenderer::new("aws")
+            .env("AWS_PROFILE", "CORPORATION-CORP_astronauts_ACCESS_GROUP")
+            .env("AWS_REGION", "ap-southeast-2")
+            .env("AWS_ACCESS_KEY_ID", "dummy")
+            .config(toml::toml! {
+                [aws.profile_aliases]
+                CORPORATION-CORP_astronauts_ACCESS_GROUP = "astro"
+                [aws.region_aliases]
+                ap-southeast-2 = "au"
+            })
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astro (au) ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
     fn credentials_file_is_ignored_when_is_directory() -> io::Result<()> {
         let dir = tempfile::tempdir()?;
         let config_path = dir.path().join("credentials");