From d1ce35252830c6b3a5329374aa2364177eafa583 Mon Sep 17 00:00:00 2001
From: Aaron Kollasch <aaron@kollasch.dev>
Date: Mon, 8 Nov 2021 15:21:09 -0500
Subject: [PATCH] feat(ruby): Add environment variable checks to ruby module
 (#3206)

* feat(ruby): Add environment variable checks to ruby module

* docs(ruby): describe version detection method
---
 docs/config/README.md |  4 ++++
 src/configs/ruby.rs   |  2 ++
 src/modules/ruby.rs   | 35 ++++++++++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index 1d7b4a92..2ef2e40b 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -2642,6 +2642,9 @@ The module will be shown if any of the following conditions are met:
 - The current directory contains a `Gemfile` file
 - The current directory contains a `.ruby-version` file
 - The current directory contains a `.rb` file
+- The environment variables `RUBY_VERSION` or `RBENV_VERSION` are set
+
+Starship gets the current Ruby version by running `ruby -v`.
 
 ### Options
 
@@ -2653,6 +2656,7 @@ The module will be shown if any of the following conditions are met:
 | `detect_extensions` | `["rb"]`                             | Which extensions should trigger this module.                              |
 | `detect_files`      | `["Gemfile", ".ruby-version"]`       | Which filenames should trigger this module.                               |
 | `detect_folders`    | `[]`                                 | Which folders should trigger this module.                                 |
+| `detect_variables`  | `["RUBY_VERSION", "RBENV_VERSION"]`  | Which environment variables should trigger this module.                   |
 | `style`             | `"bold red"`                         | The style for the module.                                                 |
 | `disabled`          | `false`                              | Disables the `ruby` module.                                               |
 
diff --git a/src/configs/ruby.rs b/src/configs/ruby.rs
index 4ba0f057..ed6520ed 100644
--- a/src/configs/ruby.rs
+++ b/src/configs/ruby.rs
@@ -13,6 +13,7 @@ pub struct RubyConfig<'a> {
     pub detect_extensions: Vec<&'a str>,
     pub detect_files: Vec<&'a str>,
     pub detect_folders: Vec<&'a str>,
+    pub detect_variables: Vec<&'a str>,
 }
 
 impl<'a> Default for RubyConfig<'a> {
@@ -26,6 +27,7 @@ impl<'a> Default for RubyConfig<'a> {
             detect_extensions: vec!["rb"],
             detect_files: vec!["Gemfile", ".ruby-version"],
             detect_folders: vec![],
+            detect_variables: vec!["RUBY_VERSION", "RBENV_VERSION"],
         }
     }
 }
diff --git a/src/modules/ruby.rs b/src/modules/ruby.rs
index 1875193f..e045a8fd 100644
--- a/src/modules/ruby.rs
+++ b/src/modules/ruby.rs
@@ -8,6 +8,7 @@ use crate::formatter::{StringFormatter, VersionFormatter};
 /// Will display the Ruby version if any of the following criteria are met:
 ///     - Current directory contains a `.rb` file
 ///     - Current directory contains a `Gemfile` or `.ruby-version` file
+///     - The environment variables `RUBY_VERSION` or `RBENV_VERSION` are set
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let mut module = context.new_module("ruby");
     let config = RubyConfig::try_load(module.config);
@@ -19,7 +20,12 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         .set_folders(&config.detect_folders)
         .is_match();
 
-    if !is_rb_project {
+    let is_rb_env = &config
+        .detect_variables
+        .iter()
+        .any(|variable| context.get_env(variable).is_some());
+
+    if !is_rb_project && !is_rb_env {
         return None;
     }
 
@@ -130,6 +136,33 @@ mod tests {
         dir.close()
     }
 
+    #[test]
+    fn with_ruby_version_env() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let actual = ModuleRenderer::new("ruby")
+            .path(dir.path())
+            .env("RUBY_VERSION", "2.5.1")
+            .collect();
+
+        let expected = Some(format!("via {}", Color::Red.bold().paint("💎 v2.5.1 ")));
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
+    #[test]
+    fn with_rbenv_version_env() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+        let actual = ModuleRenderer::new("ruby")
+            .path(dir.path())
+            .env("RBENV_VERSION", "2.6.8")
+            .collect();
+
+        // rbenv variable is only detected; its value is not used
+        let expected = Some(format!("via {}", Color::Red.bold().paint("💎 v2.5.1 ")));
+        assert_eq!(expected, actual);
+        dir.close()
+    }
+
     #[test]
     fn test_format_ruby_version() {
         let config = RubyConfig::default();