feat(nats): Add NATS Context (#5900)

---------

Co-authored-by: David Knaack <davidkna@users.noreply.github.com>

7 files changed, 127 insertions, 0 deletions

diff --git a/src/configs/mod.rs b/src/configs/mod.rs
index 7e7aa2632..6f15638e0 100644
--- a/src/configs/mod.rs
+++ b/src/configs/mod.rs
@@ -55,6 +55,7 @@ pub mod localip;
 pub mod lua;
 pub mod memory_usage;
 pub mod meson;
+pub mod nats;
 pub mod nim;
 pub mod nix_shell;
 pub mod nodejs;
@@ -218,6 +219,8 @@ pub struct FullConfig<'a> {
     #[serde(borrow)]
     meson: meson::MesonConfig<'a>,
     #[serde(borrow)]
+    nats: nats::NatsConfig<'a>,
+    #[serde(borrow)]
     nim: nim::NimConfig<'a>,
     #[serde(borrow)]
     nix_shell: nix_shell::NixShellConfig<'a>,
diff --git a/src/configs/nats.rs b/src/configs/nats.rs
new file mode 100644
index 000000000..6404357aa
--- /dev/null
+++ b/src/configs/nats.rs
@@ -0,0 +1,26 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Deserialize, Serialize)]
+#[cfg_attr(
+    feature = "config-schema",
+    derive(schemars::JsonSchema),
+    schemars(deny_unknown_fields)
+)]
+#[serde(default)]
+pub struct NatsConfig<'a> {
+    pub format: &'a str,
+    pub symbol: &'a str,
+    pub style: &'a str,
+    pub disabled: bool,
+}
+
+impl<'a> Default for NatsConfig<'a> {
+    fn default() -> Self {
+        NatsConfig {
+            format: "[$symbol($name )]($style)",
+            symbol: "✉️ ",
+            style: "bold purple",
+            disabled: true,
+        }
+    }
+}
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index f1cc181c8..8acbbe114 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -37,6 +37,7 @@ pub const PROMPT_ORDER: &[&str] = &[
     "shlvl",
     "singularity",
     "kubernetes",
+    "nats",
     "directory",
     "vcsh",
     "fossil_branch",
diff --git a/src/module.rs b/src/module.rs
index c067aa3b2..ae8babca5 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -60,6 +60,7 @@ pub const ALL_MODULES: &[&str] = &[
     "lua",
     "memory_usage",
     "meson",
+    "nats",
     "nim",
     "nix_shell",
     "nodejs",
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index e842b4940..60eed9e55 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -52,6 +52,7 @@ mod localip;
 mod lua;
 mod memory_usage;
 mod meson;
+mod nats;
 mod nim;
 mod nix_shell;
 mod nodejs;
@@ -162,6 +163,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "lua" => lua::module(context),
             "memory_usage" => memory_usage::module(context),
             "meson" => meson::module(context),
+            "nats" => nats::module(context),
             "nim" => nim::module(context),
             "nix_shell" => nix_shell::module(context),
             "nodejs" => nodejs::module(context),
@@ -285,6 +287,7 @@ pub fn description(module: &str) -> &'static str {
         "meson" => {
             "The current Meson environment, if $MESON_DEVENV and $MESON_PROJECT_NAME are set"
         }
+        "nats" => "The current NATS context",
         "nim" => "The currently installed version of Nim",
         "nix_shell" => "The nix-shell environment",
         "nodejs" => "The currently installed version of NodeJS",
diff --git a/src/modules/nats.rs b/src/modules/nats.rs
new file mode 100644
index 000000000..4a6b90766
--- /dev/null
+++ b/src/modules/nats.rs
@@ -0,0 +1,89 @@
+use super::{Context, Module, ModuleConfig};
+use serde_json as json;
+
+use crate::configs::nats::NatsConfig;
+use crate::formatter::StringFormatter;
+
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("nats");
+    let config = NatsConfig::try_load(module.config);
+
+    if config.disabled {
+        return None;
+    };
+
+    let ctx_str = context
+        .exec_cmd("nats", &["context", "info", "--json"])?
+        .stdout;
+    let nats_context: json::Value = json::from_str(&ctx_str)
+        .map_err(|e| {
+            log::warn!("Error parsing nats context JSON: {}\n", e);
+            drop(e);
+        })
+        .ok()?;
+
+    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
+        formatter
+            .map_meta(|var, _| match var {
+                "symbol" => Some(config.symbol),
+                _ => None,
+            })
+            .map_style(|variable| match variable {
+                "style" => Some(Ok(config.style)),
+                _ => None,
+            })
+            .map(|variable| match variable {
+                "name" => Some(Ok(nats_context.get("name")?.as_str()?)),
+                _ => None,
+            })
+            .parse(None, Some(context))
+    });
+
+    module.set_segments(match parsed {
+        Ok(segments) => segments,
+        Err(error) => {
+            log::warn!("Error in module `nats`:\n{}", error);
+            return None;
+        }
+    });
+
+    Some(module)
+}
+
+#[cfg(test)]
+mod tests {
+    use nu_ansi_term::Color;
+    use std::io;
+
+    use crate::test::ModuleRenderer;
+
+    #[test]
+    fn show_context() -> io::Result<()> {
+        let actual = ModuleRenderer::new("nats")
+            .config(toml::toml! {
+                [nats]
+                format = "[$symbol$name](bold purple)"
+                symbol = ""
+                disabled = false
+            })
+            .collect();
+        let expected = Some(format!("{}", Color::Purple.bold().paint("localhost")));
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+
+    #[test]
+    fn test_with_symbol() -> io::Result<()> {
+        let actual = ModuleRenderer::new("nats")
+            .config(toml::toml! {
+                [nats]
+                format = "[$symbol$name](bold red)"
+                symbol = "✉️ "
+                disabled = false
+            })
+            .collect();
+        let expected = Some(format!("{}", Color::Red.bold().paint("✉️ localhost")));
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+}
diff --git a/src/utils.rs b/src/utils.rs
index d6b56e6c9..7bd1ee97e 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -315,6 +315,10 @@ Elixir 1.10 (compiled with Erlang/OTP 22)\n",
             stdout: String::from("LuaJIT 2.0.5 -- Copyright (C) 2005-2017 Mike Pall. http://luajit.org/\n"),
             stderr: String::default(),
         }),
+        "nats context info --json" => Some(CommandOutput{
+            stdout: String::from("{\"name\":\"localhost\",\"url\":\"nats://localhost:4222\"}"),
+            stderr: String::default(),
+        }),
         "nim --version" => Some(CommandOutput {
             stdout: String::from(
                 "\