1 files changed, 33 insertions, 18 deletions

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));
+    }
+}