Skip to content

Commit 9b117d4

Browse files
committed
fix: honor cortex home for stats
1 parent 7954d02 commit 9b117d4

2 files changed

Lines changed: 56 additions & 6 deletions

File tree

src/cortex-cli/src/agent_cmd/tests.rs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,9 @@
33
#[cfg(test)]
44
mod tests {
55
use crate::agent_cmd::cli::{CopyArgs, ExportArgs};
6-
use crate::agent_cmd::loader::{
7-
load_builtin_agents, parse_frontmatter, read_file_with_encoding,
8-
};
6+
use crate::agent_cmd::loader::{load_builtin_agents, parse_frontmatter};
97
use crate::agent_cmd::types::AgentMode;
8+
use crate::utils::file::read_file_with_encoding;
109

1110
#[test]
1211
fn test_read_file_with_utf8() {

src/cortex-cli/src/stats_cmd.rs

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
99
use anyhow::{Context, Result};
1010
use clap::Parser;
11+
use cortex_engine::config::find_cortex_home;
1112
use std::collections::HashMap;
1213
use std::path::PathBuf;
1314

@@ -192,9 +193,7 @@ impl StatsCli {
192193

193194
/// Get the cortex home directory.
194195
fn get_cortex_home() -> PathBuf {
195-
dirs::home_dir()
196-
.map(|h| h.join(".cortex"))
197-
.unwrap_or_else(|| PathBuf::from(".cortex"))
196+
find_cortex_home().unwrap_or_else(|_| PathBuf::from(".cortex"))
198197
}
199198

200199
/// Get pricing for a model.
@@ -700,6 +699,58 @@ fn format_cost(cost: f64) -> String {
700699
#[cfg(test)]
701700
mod tests {
702701
use super::*;
702+
use serial_test::serial;
703+
use std::env;
704+
use std::ffi::{OsStr, OsString};
705+
use tempfile::TempDir;
706+
707+
struct EnvVarGuard {
708+
key: &'static str,
709+
original: Option<OsString>,
710+
}
711+
712+
impl EnvVarGuard {
713+
fn set(key: &'static str, value: impl AsRef<OsStr>) -> Self {
714+
let original = env::var_os(key);
715+
// SAFETY: These tests are serialized and restore the environment on drop.
716+
unsafe {
717+
env::set_var(key, value);
718+
}
719+
Self { key, original }
720+
}
721+
722+
fn remove(key: &'static str) -> Self {
723+
let original = env::var_os(key);
724+
// SAFETY: These tests are serialized and restore the environment on drop.
725+
unsafe {
726+
env::remove_var(key);
727+
}
728+
Self { key, original }
729+
}
730+
}
731+
732+
impl Drop for EnvVarGuard {
733+
fn drop(&mut self) {
734+
// SAFETY: These tests are serialized and restore the environment before returning.
735+
unsafe {
736+
match &self.original {
737+
Some(value) => env::set_var(self.key, value),
738+
None => env::remove_var(self.key),
739+
}
740+
}
741+
}
742+
}
743+
744+
#[test]
745+
#[serial]
746+
fn test_get_cortex_home_uses_cortex_home_env() {
747+
let temp_dir = TempDir::new().unwrap();
748+
let cortex_home = temp_dir.path().join("custom-cortex-home");
749+
let _config_dir = EnvVarGuard::remove("CORTEX_CONFIG_DIR");
750+
let _cortex_home = EnvVarGuard::set("CORTEX_HOME", &cortex_home);
751+
752+
assert_eq!(get_cortex_home(), cortex_home);
753+
}
703754

704755
#[test]
705756
fn test_format_number() {

0 commit comments

Comments
 (0)