diff --git a/src/cortex-cli/src/agent_cmd/tests.rs b/src/cortex-cli/src/agent_cmd/tests.rs index e2ff07f9..18f7ba75 100644 --- a/src/cortex-cli/src/agent_cmd/tests.rs +++ b/src/cortex-cli/src/agent_cmd/tests.rs @@ -3,10 +3,9 @@ #[cfg(test)] mod tests { use crate::agent_cmd::cli::{CopyArgs, ExportArgs}; - use crate::agent_cmd::loader::{ - load_builtin_agents, parse_frontmatter, read_file_with_encoding, - }; + use crate::agent_cmd::loader::{load_builtin_agents, parse_frontmatter}; use crate::agent_cmd::types::AgentMode; + use crate::utils::file::read_file_with_encoding; #[test] fn test_read_file_with_utf8() { diff --git a/src/cortex-cli/src/import_cmd.rs b/src/cortex-cli/src/import_cmd.rs index 696d93ae..2b1bf720 100644 --- a/src/cortex-cli/src/import_cmd.rs +++ b/src/cortex-cli/src/import_cmd.rs @@ -8,6 +8,7 @@ use std::collections::HashSet; use std::path::{Path, PathBuf}; use crate::styled_output::{print_info, print_success, print_warning}; +use cortex_engine::config::find_cortex_home; use cortex_engine::rollout::recorder::{RolloutRecorder, SessionMeta}; use cortex_engine::rollout::{SESSIONS_SUBDIR, get_rollout_path}; use cortex_protocol::{ @@ -45,9 +46,7 @@ impl ImportCommand { bail!("Error: Source path cannot be empty\n\nUsage: cortex import "); } - let cortex_home = dirs::home_dir() - .map(|h| h.join(".cortex")) - .ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?; + let cortex_home = get_cortex_home()?; // Read the export data let (json_content, is_from_url) = if self.source == "-" { @@ -236,6 +235,11 @@ impl ImportCommand { } } +/// Get the cortex home directory. +fn get_cortex_home() -> Result { + find_cortex_home().context("Could not determine Cortex home directory") +} + /// Fetch content from a URL. async fn fetch_url(url: &str) -> Result { // Use curl for fetching @@ -494,6 +498,58 @@ fn message_to_event(message: &ExportMessage, turn_id: &mut u64, cwd: &Path) -> R #[cfg(test)] mod tests { use super::*; + use serial_test::serial; + use std::env; + use std::ffi::{OsStr, OsString}; + use tempfile::TempDir; + + struct EnvVarGuard { + key: &'static str, + original: Option, + } + + impl EnvVarGuard { + fn set(key: &'static str, value: impl AsRef) -> Self { + let original = env::var_os(key); + // SAFETY: These tests are serialized and restore the environment on drop. + unsafe { + env::set_var(key, value); + } + Self { key, original } + } + + fn remove(key: &'static str) -> Self { + let original = env::var_os(key); + // SAFETY: These tests are serialized and restore the environment on drop. + unsafe { + env::remove_var(key); + } + Self { key, original } + } + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + // SAFETY: These tests are serialized and restore the environment before returning. + unsafe { + match &self.original { + Some(value) => env::set_var(self.key, value), + None => env::remove_var(self.key), + } + } + } + } + + #[test] + #[serial] + fn test_get_cortex_home_uses_cortex_home_env() { + let temp_dir = TempDir::new().unwrap(); + let cortex_home = temp_dir.path().join("custom-cortex-home"); + let _config_dir = EnvVarGuard::remove("CORTEX_CONFIG_DIR"); + let _cortex_home = EnvVarGuard::set("CORTEX_HOME", &cortex_home); + + assert_eq!(get_cortex_home().unwrap(), cortex_home); + } #[test] fn test_parse_export_json() {