diff --git a/src/cortex-cli/src/agent_cmd/tests.rs b/src/cortex-cli/src/agent_cmd/tests.rs index e2ff07f9f..18f7ba753 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/dag_cmd/helpers.rs b/src/cortex-cli/src/dag_cmd/helpers.rs index 7d0635486..65a240d2c 100644 --- a/src/cortex-cli/src/dag_cmd/helpers.rs +++ b/src/cortex-cli/src/dag_cmd/helpers.rs @@ -2,6 +2,7 @@ use anyhow::{Context, Result}; use cortex_agents::task::{TaskDag, TaskId, TaskSpec}; +use cortex_engine::config::find_cortex_home; use std::collections::HashMap; use std::path::PathBuf; @@ -9,10 +10,8 @@ use super::types::{DagExecutionStats, DagOutputFormat, DagSpecInput}; /// Get the DAG store path. pub fn get_dag_store_path() -> Result { - let data_dir = dirs::data_dir() - .or_else(|| dirs::home_dir().map(|h| h.join(".local/share"))) - .context("Could not determine data directory")?; - Ok(data_dir.join("cortex").join("dags")) + let cortex_home = find_cortex_home().context("Could not determine Cortex home directory")?; + Ok(cortex_home.join("dags")) } /// Load DAG specification from file. diff --git a/src/cortex-cli/src/dag_cmd/tests.rs b/src/cortex-cli/src/dag_cmd/tests.rs index a2bc60cb4..bf3ced5b3 100644 --- a/src/cortex-cli/src/dag_cmd/tests.rs +++ b/src/cortex-cli/src/dag_cmd/tests.rs @@ -1,12 +1,53 @@ //! Tests for DAG command functionality. -use super::helpers::convert_specs; +use super::helpers::{convert_specs, get_dag_store_path}; use super::types::{DagSpecInput, TaskSpecInput}; use cortex_agents::task::{DagHydrator, Task, TaskId, TaskSpec}; +use serial_test::serial; use std::collections::HashMap; +use std::env; +use std::ffi::{OsStr, OsString}; +use tempfile::TempDir; use super::executor::TaskExecutor; +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] fn test_load_yaml_spec() { let yaml = r#" @@ -31,6 +72,33 @@ tasks: assert_eq!(spec.tasks[1].depends_on, vec!["setup"]); } +#[test] +#[serial] +fn dag_store_path_uses_cortex_home() { + 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); + + let path = get_dag_store_path().unwrap(); + + assert_eq!(path, cortex_home.join("dags")); +} + +#[test] +#[serial] +fn dag_store_path_prefers_cortex_config_dir() { + let temp_dir = TempDir::new().unwrap(); + let cortex_home = temp_dir.path().join("cortex-home"); + let cortex_config_dir = temp_dir.path().join("cortex-config-dir"); + let _cortex_home = EnvVarGuard::set("CORTEX_HOME", &cortex_home); + let _config_dir = EnvVarGuard::set("CORTEX_CONFIG_DIR", &cortex_config_dir); + + let path = get_dag_store_path().unwrap(); + + assert_eq!(path, cortex_config_dir.join("dags")); +} + #[test] fn test_convert_specs() { let input = DagSpecInput {