From e5cdd9c1b3c17ffcf5416c39bf95dbb556b4f90f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?nils=20m=C3=A5s=C3=A9n?= <nils@piksel.se>
Date: Sat, 13 Mar 2021 09:35:50 +0100
Subject: feat(aws): add support for getting profile from awsu (#2451)

---
 docs/config/README.md |  3 +++
 src/modules/aws.rs    | 30 ++++++++++++++++++++++--------
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index 475981683..926e26dc0 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -249,6 +249,9 @@ The `aws` module shows the current AWS region and profile. This is based on
 When using [aws-vault](https://github.com/99designs/aws-vault) the profile
 is read from the `AWS_VAULT` env var.
 
+When using [awsu](https://github.com/kreuzwerker/awsu) the profile
+is read from the `AWSU_PROFILE` env var.
+
 ### Options
 
 | Option           | Default                                        | Description                                                     |
diff --git a/src/modules/aws.rs b/src/modules/aws.rs
index 516a2d2ba..72a57f72e 100644
--- a/src/modules/aws.rs
+++ b/src/modules/aws.rs
@@ -47,14 +47,14 @@ fn get_aws_region_from_config(context: &Context, aws_profile: Option<&str>) -> O
 }
 
 fn get_aws_profile_and_region(context: &Context) -> (Option<Profile>, Option<Region>) {
-    match (
-        context
-            .get_env("AWS_VAULT")
-            .or_else(|| context.get_env("AWS_PROFILE")),
-        context
-            .get_env("AWS_DEFAULT_REGION")
-            .or_else(|| context.get_env("AWS_REGION")),
-    ) {
+    let profile_env_vars = vec!["AWSU_PROFILE", "AWS_VAULT", "AWS_PROFILE"];
+    let profile = profile_env_vars
+        .iter()
+        .find_map(|env_var| context.get_env(env_var));
+    let region = context
+        .get_env("AWS_DEFAULT_REGION")
+        .or_else(|| context.get_env("AWS_REGION"));
+    match (profile, region) {
         (Some(p), Some(r)) => (Some(p), Some(r)),
         (None, Some(r)) => (None, Some(r)),
         (Some(ref p), None) => (
@@ -200,6 +200,20 @@ mod tests {
         assert_eq!(expected, actual);
     }
 
+    #[test]
+    fn profile_set_from_awsu() {
+        let actual = ModuleRenderer::new("aws")
+            .env("AWSU_PROFILE", "astronauts-awsu")
+            .env("AWS_PROFILE", "astronauts-profile")
+            .collect();
+        let expected = Some(format!(
+            "on {}",
+            Color::Yellow.bold().paint("☁️  astronauts-awsu ")
+        ));
+
+        assert_eq!(expected, actual);
+    }
+
     #[test]
     fn profile_and_region_set() {
         let actual = ModuleRenderer::new("aws")
-- 
cgit v1.2.3