refactor(custom): various improvements (#3829)

3 files changed, 508 insertions, 149 deletions

diff --git a/src/configs/custom.rs b/src/configs/custom.rs
index 979fe1bd9..ac06dc1df 100644
--- a/src/configs/custom.rs
+++ b/src/configs/custom.rs
@@ -1,4 +1,4 @@
-use crate::config::VecOr;
+use crate::config::{Either, VecOr};
 
 use serde::{self, Deserialize, Serialize};
 
@@ -9,17 +9,22 @@ pub struct CustomConfig<'a> {
     pub format: &'a str,
     pub symbol: &'a str,
     pub command: &'a str,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub when: Option<&'a str>,
+    pub when: Either<bool, &'a str>,
     pub shell: VecOr<&'a str>,
     pub description: &'a str,
     pub style: &'a str,
     pub disabled: bool,
-    pub files: Vec<&'a str>,
-    pub extensions: Vec<&'a str>,
-    pub directories: Vec<&'a str>,
+    #[serde(alias = "files")]
+    pub detect_files: Vec<&'a str>,
+    #[serde(alias = "extensions")]
+    pub detect_extensions: Vec<&'a str>,
+    #[serde(alias = "directories")]
+    pub detect_folders: Vec<&'a str>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub os: Option<&'a str>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub use_stdin: Option<bool>,
+    pub ignore_timeout: bool,
 }
 
 impl<'a> Default for CustomConfig<'a> {
@@ -28,15 +33,17 @@ impl<'a> Default for CustomConfig<'a> {
             format: "[$symbol($output )]($style)",
             symbol: "",
             command: "",
-            when: None,
+            when: Either::First(false),
             shell: VecOr::default(),
             description: "<custom config>",
             style: "green bold",
             disabled: false,
-            files: Vec::default(),
-            extensions: Vec::default(),
-            directories: Vec::default(),
+            detect_files: Vec::default(),
+            detect_extensions: Vec::default(),
+            detect_folders: Vec::default(),
             os: None,
+            use_stdin: None,
+            ignore_timeout: false,
         }
     }
 }
diff --git a/src/modules/custom.rs b/src/modules/custom.rs
index 66880c3e7..27b7c7dab 100644
--- a/src/modules/custom.rs
+++ b/src/modules/custom.rs
@@ -1,11 +1,18 @@
 use std::env;
 use std::io::Write;
-use std::process::{Command, Output, Stdio};
+use std::path::Path;
+use std::process::{Command, Stdio};
+use std::time::Duration;
 use std::time::Instant;
 
-use super::{Context, Module, ModuleConfig, Shell};
+use process_control::{ChildExt, Control, Output};
 
-use crate::{configs::custom::CustomConfig, formatter::StringFormatter, utils::create_command};
+use super::{Context, Module, ModuleConfig};
+
+use crate::{
+    config::Either, configs::custom::CustomConfig, formatter::StringFormatter,
+    utils::create_command,
+};
 
 /// Creates a custom module with some configuration
 ///
@@ -29,15 +36,16 @@ pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
 
     let mut is_match = context
         .try_begin_scan()?
-        .set_files(&config.files)
-        .set_extensions(&config.extensions)
-        .set_folders(&config.directories)
+        .set_extensions(&config.detect_extensions)
+        .set_files(&config.detect_files)
+        .set_folders(&config.detect_folders)
         .is_match();
 
     if !is_match {
-        if let Some(when) = config.when {
-            is_match = exec_when(when, &config.shell.0);
-        }
+        is_match = match config.when {
+            Either::First(b) => b,
+            Either::Second(s) => exec_when(s, &config, context),
+        };
 
         if !is_match {
             return None;
@@ -58,12 +66,7 @@ pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
             })
             .map_no_escaping(|variable| match variable {
                 "output" => {
-                    if context.shell == Shell::Cmd && config.shell.0.is_empty() {
-                        log::error!("Executing custom commands with cmd shell is not currently supported. Please set a different shell with the \"shell\" option.");
-                        return None;
-                    }
-
-                    let output = exec_command(config.command, &config.shell.0)?;
+                    let output = exec_command(config.command, context, &config)?;
                     let trimmed = output.trim();
 
                     if trimmed.is_empty() {
@@ -89,112 +92,105 @@ pub fn module<'a>(name: &str, context: &'a Context) -> Option<Module<'a>> {
     Some(module)
 }
 
-/// Return the invoking shell, using `shell` and fallbacking in order to STARSHIP_SHELL and "sh"
-#[cfg(not(windows))]
-fn get_shell<'a, 'b>(shell_args: &'b [&'a str]) -> (std::borrow::Cow<'a, str>, &'b [&'a str]) {
+/// Return the invoking shell, using `shell` and fallbacking in order to STARSHIP_SHELL and "sh"/"cmd"
+fn get_shell<'a, 'b>(
+    shell_args: &'b [&'a str],
+    context: &Context,
+) -> (std::borrow::Cow<'a, str>, &'b [&'a str]) {
     if !shell_args.is_empty() {
         (shell_args[0].into(), &shell_args[1..])
-    } else if let Ok(env_shell) = std::env::var("STARSHIP_SHELL") {
+    } else if let Some(env_shell) = context.get_env("STARSHIP_SHELL") {
         (env_shell.into(), &[] as &[&str])
+    } else if cfg!(windows) {
+        // `/C` is added by `handle_shell`
+        ("cmd".into(), &[] as &[&str])
     } else {
         ("sh".into(), &[] as &[&str])
     }
 }
 
-/// Attempt to run the given command in a shell by passing it as `stdin` to `get_shell()`
-#[cfg(not(windows))]
-fn shell_command(cmd: &str, shell_args: &[&str]) -> Option<Output> {
-    let (shell, shell_args) = get_shell(shell_args);
-    let mut command = create_command(shell.as_ref()).ok()?;
+/// Attempt to run the given command in a shell by passing it as either `stdin` or an argument to `get_shell()`,
+/// depending on the configuration or by invoking a platform-specific falback shell if `shell` is empty.
+fn shell_command(cmd: &str, config: &CustomConfig, context: &Context) -> Option<Output> {
+    let (shell, shell_args) = get_shell(config.shell.0.as_ref(), context);
+    let mut use_stdin = config.use_stdin;
+
+    let mut command = match create_command(shell.as_ref()) {
+        Ok(command) => command,
+        // Don't attempt to use fallback shell if the user specified a shell
+        Err(error) if !shell_args.is_empty() => {
+            log::debug!(
+                "Error creating command with STARSHIP_SHELL, falling back to fallback shell: {}",
+                error
+            );
+
+            // Skip `handle_shell` and just set the shell and command
+            use_stdin = Some(!cfg!(windows));
+
+            if cfg!(windows) {
+                let mut c = create_command("cmd").ok()?;
+                c.arg("/C");
+                c
+            } else {
+                let mut c = create_command("/usr/bin/env").ok()?;
+                c.arg("sh");
+                c
+            }
+        }
+        _ => return None,
+    };
 
     command
+        .current_dir(&context.current_dir)
         .args(shell_args)
         .stdin(Stdio::piped())
         .stdout(Stdio::piped())
         .stderr(Stdio::piped());
 
-    handle_powershell(&mut command, &shell, shell_args);
+    let use_stdin = use_stdin.unwrap_or_else(|| handle_shell(&mut command, &shell, shell_args));
+
+    if !use_stdin {
+        command.arg(cmd);
+    }
 
     let mut child = match command.spawn() {
-        Ok(command) => command,
-        Err(err) => {
-            log::trace!("Error executing command: {:?}", err);
+        Ok(child) => child,
+        Err(error) => {
             log::debug!(
-                "Could not launch command with given shell or STARSHIP_SHELL env variable, retrying with /usr/bin/env sh"
+                "Failed to run command with given shell or STARSHIP_SHELL env variable:: {}",
+                error
             );
-
-            #[allow(clippy::disallowed_methods)]
-            Command::new("/usr/bin/env")
-                .arg("sh")
-                .stdin(Stdio::piped())
-                .stdout(Stdio::piped())
-                .stderr(Stdio::piped())
-                .spawn()
-                .ok()?
+            return None;
         }
     };
 
-    child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
-    child.wait_with_output().ok()
-}
-
-/// Attempt to run the given command in a shell by passing it as `stdin` to `get_shell()`,
-/// or by invoking cmd.exe /C.
-#[cfg(windows)]
-fn shell_command(cmd: &str, shell_args: &[&str]) -> Option<Output> {
-    let (shell, shell_args) = if !shell_args.is_empty() {
-        (
-            Some(std::borrow::Cow::Borrowed(shell_args[0])),
-            &shell_args[1..],
-        )
-    } else if let Some(env_shell) = std::env::var("STARSHIP_SHELL")
-        .ok()
-        .filter(|s| !cfg!(test) && !s.is_empty())
-    {
-        (Some(std::borrow::Cow::Owned(env_shell)), &[] as &[&str])
-    } else {
-        (None, &[] as &[&str])
-    };
-
-    if let Some(forced_shell) = shell {
-        let mut command = create_command(forced_shell.as_ref()).ok()?;
-
-        command
-            .args(shell_args)
-            .stdin(Stdio::piped())
-            .stdout(Stdio::piped())
-            .stderr(Stdio::piped());
+    if use_stdin {
+        child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
+    }
 
-        handle_powershell(&mut command, &forced_shell, shell_args);
+    let mut output = child.controlled_with_output();
 
-        if let Ok(mut child) = command.spawn() {
-            child.stdin.as_mut()?.write_all(cmd.as_bytes()).ok()?;
+    if !config.ignore_timeout {
+        output = output
+            .time_limit(Duration::from_millis(context.root_config.command_timeout))
+            .terminate_for_timeout()
+    }
 
-            return child.wait_with_output().ok();
+    match output.wait().ok()? {
+        None => {
+            log::warn!("Executing custom command {cmd:?} timed out.");
+            log::warn!("You can set command_timeout in your config to a higher value or set ignore_timeout to true for this module to allow longer-running commands to keep executing.");
+            None
         }
-
-        log::debug!(
-            "Could not launch command with given shell or STARSHIP_SHELL env variable, retrying with cmd.exe /C"
-        );
+        Some(status) => Some(status),
     }
-
-    let command = create_command("cmd")
-        .ok()?
-        .arg("/C")
-        .arg(cmd)
-        .stdin(Stdio::piped())
-        .stdout(Stdio::piped())
-        .stderr(Stdio::piped())
-        .spawn();
-
-    command.ok()?.wait_with_output().ok()
 }
 
 /// Execute the given command capturing all output, and return whether it return 0
-fn exec_when(cmd: &str, shell_args: &[&str]) -> bool {
+fn exec_when(cmd: &str, config: &CustomConfig, context: &Context) -> bool {
     log::trace!("Running '{}'", cmd);
 
-    if let Some(output) = shell_command(cmd, shell_args) {
+    if let Some(output) = shell_command(cmd, config, context) {
         if !output.status.success() {
             log::trace!("non-zero exit code '{:?}'", output.status.code());
             log::trace!(
@@ -216,10 +212,10 @@ fn exec_when(cmd: &str, shell_args: &[&str]) -> bool {
 }
 
 /// Execute the given command, returning its output on success
-fn exec_command(cmd: &str, shell_args: &[&str]) -> Option<String> {
-    log::trace!("Running '{}'", cmd);
+fn exec_command(cmd: &str, context: &Context, config: &CustomConfig) -> Option<String> {
+    log::trace!("Running '{cmd}'");
 
-    if let Some(output) = shell_command(cmd, shell_args) {
+    if let Some(output) = shell_command(cmd, config, context) {
         if !output.status.success() {
             log::trace!("Non-zero exit code '{:?}'", output.status.code());
             log::trace!(
@@ -241,14 +237,31 @@ fn exec_command(cmd: &str, shell_args: &[&str]) -> Option<String> {
 
 /// If the specified shell refers to PowerShell, adds the arguments "-Command -" to the
 /// given command.
-fn handle_powershell(command: &mut Command, shell: &str, shell_args: &[&str]) {
-    let is_powershell = shell.ends_with("pwsh.exe")
-        || shell.ends_with("powershell.exe")
-        || shell.ends_with("pwsh")
-        || shell.ends_with("powershell");
-
-    if is_powershell && shell_args.is_empty() {
-        command.arg("-NoProfile").arg("-Command").arg("-");
+/// Retruns `false` if the shell shell expects scripts as arguments, `true` if as `stdin`.
+fn handle_shell(command: &mut Command, shell: &str, shell_args: &[&str]) -> bool {
+    let shell_exe = Path::new(shell).file_stem();
+    let no_args = shell_args.is_empty();
+
+    match shell_exe.and_then(std::ffi::OsStr::to_str) {
+        Some("pwsh" | "powershell") => {
+            if no_args {
+                command.arg("-NoProfile").arg("-Command").arg("-");
+            }
+            true
+        }
+        Some("cmd") => {
+            if no_args {
+                command.arg("/C");
+            }
+            false
+        }
+        Some("nu") => {
+            if no_args {
+                command.arg("-c");
+            }
+            false
+        }
+        _ => true,
     }
 }
 
@@ -256,10 +269,15 @@ fn handle_powershell(command: &mut Command, shell: &str, shell_args: &[&str]) {
 mod tests {
     use super::*;
 
+    use crate::test::ModuleRenderer;
+    use ansi_term::Color;
+    use std::fs::File;
+    use std::io;
+
     #[cfg(not(windows))]
     const SHELL: &[&str] = &["/bin/sh"];
     #[cfg(windows)]
-    const SHELL: &[&str] = &[];
+    const SHELL: &[&str] = &["cmd"];
 
     #[cfg(not(windows))]
     const FAILING_COMMAND: &str = "false";
@@ -268,66 +286,394 @@ mod tests {
 
     const UNKNOWN_COMMAND: &str = "ydelsyiedsieudleylse dyesdesl";
 
+    fn render_cmd(cmd: &str) -> io::Result<Option<String>> {
+        let dir = tempfile::tempdir()?;
+        let cmd = cmd.to_owned();
+        let shell = SHELL.iter().map(|s| s.to_owned()).collect::<Vec<_>>();
+        let out = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "$output"
+                command = cmd
+                shell = shell
+                when = true
+                ignore_timeout = true
+            })
+            .collect();
+        dir.close()?;
+        Ok(out)
+    }
+
+    fn render_when(cmd: &str) -> io::Result<bool> {
+        let dir = tempfile::tempdir()?;
+        let cmd = cmd.to_owned();
+        let shell = SHELL.iter().map(|s| s.to_owned()).collect::<Vec<_>>();
+        let out = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                when = cmd
+                shell = shell
+                ignore_timeout = true
+            })
+            .collect()
+            .is_some();
+        dir.close()?;
+        Ok(out)
+    }
+
     #[test]
-    fn when_returns_right_value() {
-        assert!(exec_when("echo hello", SHELL));
-        assert!(!exec_when(FAILING_COMMAND, SHELL));
+    fn when_returns_right_value() -> io::Result<()> {
+        assert!(render_cmd("echo hello")?.is_some());
+        assert!(render_cmd(FAILING_COMMAND)?.is_none());
+        Ok(())
     }
 
     #[test]
-    fn when_returns_false_if_invalid_command() {
-        assert!(!exec_when(UNKNOWN_COMMAND, SHELL));
+    fn when_returns_false_if_invalid_command() -> io::Result<()> {
+        assert!(!render_when(UNKNOWN_COMMAND)?);
+        Ok(())
     }
 
     #[test]
     #[cfg(not(windows))]
-    fn command_returns_right_string() {
-        assert_eq!(exec_command("echo hello", SHELL), Some("hello\n".into()));
-        assert_eq!(
-            exec_command("echo 강남스타일", SHELL),
-            Some("강남스타일\n".into())
-        );
+    fn command_returns_right_string() -> io::Result<()> {
+        assert_eq!(render_cmd("echo hello")?, Some("hello".into()));
+        assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into()));
+        Ok(())
     }
 
     #[test]
     #[cfg(windows)]
-    fn command_returns_right_string() {
-        assert_eq!(exec_command("echo hello", SHELL), Some("hello\r\n".into()));
-        assert_eq!(
-            exec_command("echo 강남스타일", SHELL),
-            Some("강남스타일\r\n".into())
-        );
+    fn command_returns_right_string() -> io::Result<()> {
+        assert_eq!(render_cmd("echo hello")?, Some("hello".into()));
+        assert_eq!(render_cmd("echo 강남스타일")?, Some("강남스타일".into()));
+        Ok(())
     }
 
     #[test]
     #[cfg(not(windows))]
-    fn command_ignores_stderr() {
-        assert_eq!(
-            exec_command("echo foo 1>&2; echo bar", SHELL),
-            Some("bar\n".into())
-        );
-        assert_eq!(
-            exec_command("echo foo; echo bar 1>&2", SHELL),
-            Some("foo\n".into())
-        );
+    fn command_ignores_stderr() -> io::Result<()> {
+        assert_eq!(render_cmd("echo foo 1>&2; echo bar")?, Some("bar".into()));
+        assert_eq!(render_cmd("echo foo; echo bar 1>&2")?, Some("foo".into()));
+        Ok(())
     }
 
     #[test]
     #[cfg(windows)]
-    fn command_ignores_stderr() {
-        assert_eq!(
-            exec_command("echo foo 1>&2 & echo bar", SHELL),
-            Some("bar\r\n".into())
-        );
-        assert_eq!(
-            exec_command("echo foo& echo bar 1>&2", SHELL),
-            Some("foo\r\n".into())
-        );
+    fn command_ignores_stderr() -> io::Result<()> {
+        assert_eq!(render_cmd("echo foo 1>&2 & echo bar")?, Some("bar".into()));
+        assert_eq!(render_cmd("echo foo& echo bar 1>&2")?, Some("foo".into()));
+        Ok(())
+    }
+
+    #[test]
+    fn command_can_fail() -> io::Result<()> {
+        assert_eq!(render_cmd(FAILING_COMMAND)?, None);
+        assert_eq!(render_cmd(UNKNOWN_COMMAND)?, None);
+        Ok(())
+    }
+
+    #[test]
+    fn cwd_command() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let mut f = File::create(dir.path().join("a.txt"))?;
+        write!(f, "hello")?;
+        f.sync_all()?;
+
+        let cat = if cfg!(windows) { "type" } else { "cat" };
+        let cmd = format!("{cat} a.txt");
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                command = cmd
+                when = true
+                ignore_timeout = true
+            })
+            .collect();
+        let expected = Some(format!("{}", Color::Green.bold().paint("hello ")));
+
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn cwd_when() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        File::create(dir.path().join("a.txt"))?.sync_all()?;
+
+        let cat = if cfg!(windows) { "type" } else { "cat" };
+        let cmd = format!("{cat} a.txt");
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                when = cmd
+                ignore_timeout = true
+            })
+            .collect();
+        let expected = Some("test".to_owned());
+
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn use_stdin_false() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let shell = if cfg!(windows) {
+            vec![
+                "powershell".to_owned(),
+                "-NoProfile".to_owned(),
+                "-Command".to_owned(),
+            ]
+        } else {
+            vec!["sh".to_owned(), "-c".to_owned()]
+        };
+
+        // `use_stdin = false` doesn't like Korean on Windows
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                command = "echo test"
+                when = true
+                use_stdin = false
+                shell = shell
+                ignore_timeout = true
+            })
+            .collect();
+        let expected = Some(format!("{}", Color::Green.bold().paint("test ")));
+
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn use_stdin_true() -> io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let shell = if cfg!(windows) {
+            vec![
+                "powershell".to_owned(),
+                "-NoProfile".to_owned(),
+                "-Command".to_owned(),
+                "-".to_owned(),
+            ]
+        } else {
+            vec!["sh".to_owned()]
+        };
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                command = "echo 강남스타일"
+                when = true
+                use_stdin = true
+                ignore_timeout = true
+                shell = shell
+            })
+            .collect();
+        let expected = Some(format!("{}", Color::Green.bold().paint("강남스타일 ")));
+
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    #[cfg(not(windows))]
+    fn when_true_with_string() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                shell = ["sh"]
+                when = "true"
+                ignore_timeout = true
+            })
+            .collect();
+        let expected = Some("test".to_string());
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    #[cfg(not(windows))]
+    fn when_false_with_string() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                shell = ["sh"]
+                when = "false"
+                ignore_timeout = true
+            })
+            .collect();
+        let expected = None;
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn when_true_with_bool() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                when = true
+            })
+            .collect();
+        let expected = Some("test".to_string());
+        assert_eq!(expected, actual);
+
+        dir.close()
     }
 
     #[test]
-    fn command_can_fail() {
-        assert_eq!(exec_command(FAILING_COMMAND, SHELL), None);
-        assert_eq!(exec_command(UNKNOWN_COMMAND, SHELL), None);
+    #[cfg(not(windows))]
+    fn when_false_with_bool() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                when = false
+            })
+            .collect();
+        let expected = None;
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn timeout_short_cmd() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let shell = if cfg!(windows) {
+            "powershell".to_owned()
+        } else {
+            "sh".to_owned()
+        };
+
+        let when = if cfg!(windows) {
+            "$true".to_owned()
+        } else {
+            "true".to_owned()
+        };
+
+        // Use a long timeout to ensure that the test doesn't fail
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                command_timeout = 10000
+                [custom.test]
+                format = "test"
+                when = when
+                shell = shell
+                ignore_timeout = false
+            })
+            .collect();
+        let expected = Some("test".to_owned());
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn timeout_cmd() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        let shell = if cfg!(windows) {
+            "powershell".to_owned()
+        } else {
+            "sh".to_owned()
+        };
+
+        // Use a long timeout to ensure that the test doesn't fail
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                when = "sleep 3"
+                shell = shell
+                ignore_timeout = false
+            })
+            .collect();
+        let expected = None;
+        assert_eq!(expected, actual);
+
+        dir.close()
+    }
+
+    #[test]
+    fn config_aliases_work() -> std::io::Result<()> {
+        let dir = tempfile::tempdir()?;
+
+        File::create(dir.path().join("a.txt"))?;
+        std::fs::create_dir(dir.path().join("dir"))?;
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                files = ["a.txt"]
+            })
+            .collect();
+        let expected = Some("test".to_string());
+        assert_eq!(expected, actual);
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                extensions = ["txt"]
+            })
+            .collect();
+        let expected = Some("test".to_string());
+        assert_eq!(expected, actual);
+
+        let actual = ModuleRenderer::new("custom.test")
+            .path(dir.path())
+            .config(toml::toml! {
+                [custom.test]
+                format = "test"
+                directories = ["dir"]
+            })
+            .collect();
+        let expected = Some("test".to_string());
+        assert_eq!(expected, actual);
+
+        dir.close()
     }
 }
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 7392272d6..d58b73dff 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -163,6 +163,12 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "vagrant" => vagrant::module(context),
             "vcsh" => vcsh::module(context),
             "zig" => zig::module(context),
+            // Added for tests, avoid potential side effects in production code.
+            #[cfg(test)]
+            custom if custom.starts_with("custom.") => {
+                // SAFETY: We just checked that the module starts with "custom."
+                custom::module(custom.strip_prefix("custom.").unwrap(), context)
+            }
             _ => {
                 eprintln!("Error: Unknown module {}. Use starship module --list to list out all supported modules.", module);
                 None