Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions src/cortex-cli/src/agent_cmd/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
7 changes: 3 additions & 4 deletions src/cortex-cli/src/dag_cmd/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,16 @@

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;

use super::types::{DagExecutionStats, DagOutputFormat, DagSpecInput};

/// Get the DAG store path.
pub fn get_dag_store_path() -> Result<PathBuf> {
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.
Expand Down
70 changes: 69 additions & 1 deletion src/cortex-cli/src/dag_cmd/tests.rs
Original file line number Diff line number Diff line change
@@ -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<OsString>,
}

impl EnvVarGuard {
fn set(key: &'static str, value: impl AsRef<OsStr>) -> 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#"
Expand All @@ -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 {
Expand Down