Set status with arg rather than env

4 files changed, 47 insertions, 25 deletions

diff --git a/src/main.rs b/src/main.rs
index ab7e04e56..7cbe7d225 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,7 +6,7 @@ mod modules;
 mod print;
 
 use ansi_term::Style;
-use clap::App;
+use clap::{App, Arg};
 
 pub struct Segment {
     style: Style,
@@ -16,14 +16,19 @@ pub struct Segment {
 }
 
 fn main() {
-    App::new("Starship")
+    let args = App::new("Starship")
         .about("The cross-platform prompt for astronauts. ✨🚀")
         // pull the version number from Cargo.toml
         .version(crate_version!())
         // pull the authors from Cargo.toml
         .author(crate_authors!())
         .after_help("https://github.com/matchai/starship")
+        .arg(
+            Arg::with_name("status_code")
+                .help("The status code of the previously run command")
+                .required(true),
+        )
         .get_matches();
 
-    print::prompt();
+    print::prompt(args);
 }
diff --git a/src/modules/char.rs b/src/modules/char.rs
index 96387c740..b5d6183af 100644
--- a/src/modules/char.rs
+++ b/src/modules/char.rs
@@ -1,34 +1,49 @@
 use crate::Segment;
 use ansi_term::{Color, Style};
-use std::env;
+use clap::ArgMatches;
 
-pub fn segment() -> Segment {
+pub fn segment(args: &ArgMatches) -> Segment {
     const PROMPT_CHAR: &str = "➜ ";
     const COLOR_SUCCESS: Color = Color::Green;
     const COLOR_FAILURE: Color = Color::Red;
 
-    let default_prefix = Segment {
-        value: String::from("testPrefix"),
-        style: Style::default(),
-        prefix: None,
-        suffix: None,
-    };
-
     let color;
-    if let Ok(status) = env::var("status") {
-        if status == "0" {
-            color = COLOR_SUCCESS;
-        } else {
-            color = COLOR_FAILURE;
-        }
+    if args.value_of("status_code").unwrap() == "0" {
+        color = COLOR_SUCCESS;
     } else {
-        panic!("No status environment variable provided");
+        color = COLOR_FAILURE;
     }
 
     Segment {
-        prefix: Some(Box::new(default_prefix)),
+        prefix: None,
         value: String::from(PROMPT_CHAR),
-        style: Style::new().fg(color),
+        style: Style::from(color),
         suffix: None,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::{App, Arg};
+
+    #[test]
+    fn char_section_success_status() {
+        let args = App::new("starship")
+            .arg(Arg::with_name("status_code"))
+            .get_matches_from(vec!["starship", "0"]);
+
+        let segment = segment(&args);
+        assert_eq!(segment.style, Style::from(Color::Green));
+    }
+
+    #[test]
+    fn char_section_failure_status() {
+        let args = App::new("starship")
+            .arg(Arg::with_name("status_code"))
+            .get_matches_from(vec!["starship", "1"]);
+
+        let segment = segment(&args);
+        assert_eq!(segment.style, Style::from(Color::Red));
+    }
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 8fe6d63a4..52474cbcf 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -1,10 +1,11 @@
 mod char;
 
 use crate::Segment;
+use clap::ArgMatches;
 
-pub fn handle(module: &str) -> Segment {
+pub fn handle(module: &str, args: &ArgMatches) -> Segment {
     match module {
-        "char" => char::segment(),
+        "char" => char::segment(&args),
 
         _ => panic!("Unknown module: {}", module),
     }
diff --git a/src/print.rs b/src/print.rs
index 2f7ca0999..1fc6e535e 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -1,11 +1,12 @@
 use crate::modules;
 use crate::Segment;
+use clap::ArgMatches;
 
-pub fn prompt() {
+pub fn prompt(args: ArgMatches) {
     let default_prompt = vec!["char"];
 
     for module in default_prompt {
-        let segment = modules::handle(module);
+        let segment = modules::handle(module, &args);
         print_segment(segment);
     }
 }