[ENH] Use writer injection in CLI

rust/cli/Cargo.toml

@@ -40,6 +40,7 @@ tui-input = { version = "0.12.1", default-features = false, features = ["crosste
 textwrap = "0.16.2"
 futures = {workspace = true}
 supports-color = "3.0.2"
+async-trait = { workspace = true }
 
 
 [[bin]]

rust/cli/src/commands/copy.rs

@@ -27,7 +27,7 @@ pub enum CopyError {
     CollectionAlreadyExists(String),
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub struct CopyArgs {
     #[clap(
         long = "all",

rust/cli/src/commands/db.rs

@@ -58,21 +58,21 @@ impl fmt::Display for Language {
     }
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct DbArgs {
     #[clap(long, hide = true, help = "Flag to use during development")]
     dev: bool,
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct CreateArgs {
     #[clap(flatten)]
     db_args: DbArgs,
     #[clap(index = 1, help = "The name of the DB to create")]
     name: Option<String>,
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct DeleteArgs {
     #[clap(flatten)]
     db_args: DbArgs,
@@ -87,7 +87,7 @@ pub struct DeleteArgs {
     force: bool,
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct ConnectArgs {
     #[clap(flatten)]
     db_args: DbArgs,
@@ -100,13 +100,13 @@ pub struct ConnectArgs {
     language: Option<Language>,
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct ListArgs {
     #[clap(flatten)]
     db_args: DbArgs,
 }
 
-#[derive(Subcommand, Debug)]
+#[derive(Subcommand, Debug, Clone)]
 pub enum DbCommand {
     #[command(about = "Generate a connection snippet to a DB")]
     Connect(ConnectArgs),

rust/cli/src/commands/install.rs

@@ -40,7 +40,7 @@ pub enum InstallError {
     EnvFileWriteFailed,
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub struct InstallArgs {
     #[clap(index = 1, help = "The name of the sample app to install")]
     name: Option<String>,

rust/cli/src/commands/login.rs

@@ -18,7 +18,7 @@ use std::time::Duration;
 use thiserror::Error;
 use tokio::time::sleep;
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub struct LoginArgs {
     #[clap(long, help = "Profile name to associate with auth credentials")]
     profile: Option<String>,

rust/cli/src/commands/profile.rs

@@ -18,7 +18,7 @@ pub enum ProfileError {
     ProfileAlreadyExists(String),
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct DeleteArgs {
     #[clap(index = 1, help = "The name of the profile to delete")]
     name: String,
@@ -30,21 +30,21 @@ pub struct DeleteArgs {
     force: bool,
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct RenameArgs {
     #[clap(index = 1, help = "The name of the profile to rename")]
     name: String,
     #[clap(index = 2, help = "The new name for the profile to rename")]
     new_name: String,
 }
 
-#[derive(Args, Debug)]
+#[derive(Args, Debug, Clone)]
 pub struct UseArgs {
     #[clap(help = "The name of the profile to use as the active profile")]
     name: String,
 }
 
-#[derive(Subcommand, Debug)]
+#[derive(Subcommand, Debug, Clone)]
 pub enum ProfileCommand {
     #[command(about = "Delete profiles")]
     Delete(DeleteArgs),

rust/cli/src/commands/run.rs

@@ -1,9 +1,12 @@
 use crate::ui_utils::LOGO;
 use crate::utils::CliError;
+use crate::utils::UtilsError;
+use crate::{cli_writeln, CommandHandler};
 use chroma_frontend::config::FrontendServerConfig;
 use chroma_frontend::frontend_service_entrypoint_with_config;
 use clap::Parser;
 use colored::Colorize;
+use std::io::{stdout, Stdout, Write};
 use std::net::TcpListener;
 use std::sync::Arc;
 use thiserror::Error;
@@ -18,7 +21,7 @@ pub enum RunError {
     ServerStartFailed,
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub struct RunArgs {
     #[clap(
         index = 1,
@@ -50,75 +53,140 @@ pub struct RunArgs {
     port: Option<u16>,
 }
 
-fn validate_host(address: &String, port: u16) -> bool {
-    let socket = format!("{}:{}", address, port);
-    TcpListener::bind(&socket).is_ok()
+pub struct RunCommand<W: Write> {
+    args: RunArgs,
+    config: FrontendServerConfig,
+    writer: W,
 }
 
-fn override_default_config_with_args(args: RunArgs) -> Result<FrontendServerConfig, CliError> {
-    let mut config = FrontendServerConfig::single_node_default();
-
-    if let Some(path) = args.path {
-        config.persist_path = path;
+impl<W: Write> RunCommand<W> {
+    pub fn new(args: RunArgs, writer: W) -> Self {
+        Self {
+            args,
+            writer,
+            config: FrontendServerConfig::single_node_default(),
+        }
     }
 
-    if let Some(port) = args.port {
-        config.port = port;
+    fn validate_host(address: &String, port: u16) -> bool {
+        let socket = format!("{}:{}", address, port);
+        TcpListener::bind(&socket).is_ok()
     }
 
-    if let Some(host) = args.host {
-        config.listen_address = host;
-    }
+    fn override_default_config_with_args(&mut self) -> Result<(), CliError> {
+        if let Some(path) = &self.args.path {
+            self.config.persist_path = path.to_owned();
+        }
 
-    if !validate_host(&config.listen_address, config.port) {
-        return Err(RunError::AddressUnavailable(config.listen_address, config.port).into());
-    }
+        if let Some(port) = &self.args.port {
+            self.config.port = port.to_owned();
+        }
 
-    Ok(config)
-}
+        if let Some(host) = &self.args.host {
+            self.config.listen_address = host.to_owned();
+        }
+
+        if !Self::validate_host(&self.config.listen_address, self.config.port) {
+            return Err(RunError::AddressUnavailable(
+                self.config.listen_address.to_owned(),
+                self.config.port,
+            )
+            .into());
+        }
+
+        Ok(())
+    }
 
-fn display_run_message(config: &FrontendServerConfig) {
-    println!("{}", LOGO);
-    println!("Saving data to: {}", config.persist_path.bold());
-    println!(
-        "Connect to Chroma at: {}",
-        format!("http://localhost:{}", config.port)
+    fn run_message(&self) -> String {
+        let host = format!("http://localhost:{}", self.config.port)
             .underline()
-            .blue()
-    );
-    println!(
-        "Getting started guide: {}",
-        "https://docs.trychroma.com/docs/overview/getting-started\n"
+            .blue();
+
+        let docs = "https://docs.trychroma.com/docs/overview/getting-started\n"
             .underline()
-            .blue()
-    );
+            .blue();
+
+        format!(
+            "{}\nSaving data to: {}\nConnect to Chroma at: {}\nGetting started guide: {}",
+            LOGO,
+            self.config.persist_path.bold(),
+            host,
+            docs
+        )
+    }
+}
+
+impl RunCommand<Stdout> {
+    pub fn default(run_args: RunArgs) -> Self {
+        let stdout = stdout();
+        RunCommand::new(run_args, stdout)
+    }
 }
 
-pub fn run(args: RunArgs) -> Result<(), CliError> {
-    let config = match &args.config_path {
-        Some(config_path) => {
-            if !std::path::Path::new(config_path).exists() {
-                eprintln!(
-                    "Could not find {config_path:?} in {:?}",
-                    std::env::current_dir()
-                        .map(
-                            |p| String::from_utf8_lossy(p.as_os_str().as_encoded_bytes())
-                                .to_string()
-                        )
-                        .unwrap_or("<unknown>".to_string())
-                );
-                return Err(RunError::ConfigFileNotFound(config_path.to_string()).into());
+#[async_trait::async_trait]
+impl<W: Write + Send> CommandHandler for RunCommand<W> {
+    async fn run(&mut self) -> Result<(), CliError> {
+        match &self.args.config_path {
+            Some(config_path) => {
+                if !std::path::Path::new(config_path).exists() {
+                    return Err(RunError::ConfigFileNotFound(config_path.to_string()).into());
+                }
+                self.config = FrontendServerConfig::load_from_path(config_path)
             }
-            FrontendServerConfig::load_from_path(config_path)
-        }
-        None => override_default_config_with_args(args)?,
-    };
+            None => self.override_default_config_with_args()?,
+        };
 
-    display_run_message(&config);
+        cli_writeln!(self.writer, "{}", self.run_message())?;
 
-    let runtime = tokio::runtime::Runtime::new().map_err(|_| RunError::ServerStartFailed)?;
-    runtime.block_on(async {
-        frontend_service_entrypoint_with_config(Arc::new(()), Arc::new(()), &config).await;
-    });
-    Ok(())
+        frontend_service_entrypoint_with_config(Arc::new(()), Arc::new(()), &self.config).await;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::client::chroma_client::ChromaClient;
+    use crate::commands::run::{RunArgs, RunCommand};
+    use crate::CommandHandler;
+    use std::sync::Arc;
+    use tokio::sync::Mutex;
+
+    #[tokio::test]
+    async fn test_run() {
+        use tokio::time::{sleep, Duration};
+
+        let port = 8001;
+        let path = "test_data".to_string();
+        let run_args = RunArgs {
+            config_path: None,
+            path: Some(path.clone()),
+            port: Some(port),
+            host: None,
+        };
+
+        let writer = Vec::new();
+        let run_command_arc = Arc::new(Mutex::new(RunCommand::new(run_args, writer)));
+
+        let run_command_arc_clone = run_command_arc.clone();
+        let server_handle = tokio::spawn(async move {
+            let mut command = run_command_arc_clone.lock().await;
+            command.run().await.unwrap();
+        });
+
+        sleep(Duration::from_millis(500)).await;
+
+        let url = format!("http://localhost:{}", port);
+        let chroma_client = ChromaClient::local_default();
+        let response = chroma_client.healthcheck().await.unwrap();
+
+        server_handle.abort();
+
+        let run_command = run_command_arc.lock().await;
+        let message = run_command.run_message();
+        let output = String::from_utf8(run_command.writer.clone()).unwrap();
+        assert!(output.contains(message.as_str()));
+        assert!(output.contains(port.to_string().as_str()));
+        assert!(output.contains(path.as_str()));
+    }
 }

rust/cli/src/commands/vacuum.rs

@@ -39,7 +39,7 @@ pub enum VacuumError {
     VacuumFailed,
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub struct VacuumArgs {
     #[clap(long, help = "The path of your Chroma DB")]
     path: Option<String>,